feat(webserver): support filter by users for userEvents api (#1956)

* feat(webserver): support filter by users for userEvents api * add indexing * add unit test
TabbyML · Apr 24, 2024 · 320c559 · 320c559
1 parent 6239795
commit 320c559
Show file tree

Hide file tree

Showing 9 changed files with 106 additions and 7 deletions.
diff --git a/ee/tabby-db/migrations/0025_user-events.down.sql b/ee/tabby-db/migrations/0025_user-events.down.sql
@@ -1 +1,3 @@
+DROP INDEX idx_user_events_user_id;
+DROP INDEX idx_user_events_created_at;
 DROP TABLE user_events;
diff --git a/ee/tabby-db/migrations/0025_user-events.up.sql b/ee/tabby-db/migrations/0025_user-events.up.sql
@@ -6,4 +6,5 @@ CREATE TABLE user_events (
     payload BLOB NOT NULL,
     FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
 );
+CREATE INDEX idx_user_events_user_id ON user_events(user_id);
 CREATE INDEX idx_user_events_created_at ON user_events(created_at);
diff --git a/ee/tabby-db/schema.sqlite b/ee/tabby-db/schema.sqlite
diff --git a/ee/tabby-db/schema/schema.sql b/ee/tabby-db/schema/schema.sql
@@ -134,6 +134,7 @@ CREATE TABLE user_events(
   payload BLOB NOT NULL,
   FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE
 );
+CREATE INDEX idx_user_events_user_id ON user_events(user_id);
 CREATE INDEX idx_user_events_created_at ON user_events(created_at);
 CREATE TABLE refresh_tokens(
   id INTEGER PRIMARY KEY AUTOINCREMENT,

diff --git a/ee/tabby-db/src/user_events.rs b/ee/tabby-db/src/user_events.rs
@@ -44,12 +44,18 @@ impl DbConn {
         limit: Option<usize>,
         skip_id: Option<i32>,
         backwards: bool,
+        users: Vec<i64>,
         start: DateTimeUtc,
         end: DateTimeUtc,
     ) -> Result<Vec<UserEventDAO>> {
+        let users = users
+            .iter()
+            .map(|u| u.to_string())
+            .collect::<Vec<_>>()
+            .join(",");
+        let no_selected_users = users.is_empty();
         let condition = Some(format!(
-            "created_at >= '{}' AND created_at < '{}'",
-            start, end,
+            "created_at >= '{start}' AND created_at < '{end}' AND ({no_selected_users} OR user_id IN ({users}))"
         ));
         let events = query_paged_as!(
             UserEventDAO,

diff --git a/ee/tabby-webserver/graphql/schema.graphql b/ee/tabby-webserver/graphql/schema.graphql
@@ -259,7 +259,7 @@ type Query {
   jobs: [String!]!
   dailyStatsInPastYear(users: [ID!]): [CompletionStats!]!
   dailyStats(start: DateTimeUtc!, end: DateTimeUtc!, users: [ID!], languages: [Language!]): [CompletionStats!]!
-  userEvents(after: String, before: String, first: Int, last: Int, start: DateTimeUtc!, end: DateTimeUtc!): UserEventConnection!
+  userEvents(after: String, before: String, first: Int, last: Int, users: [ID!], start: DateTimeUtc!, end: DateTimeUtc!): UserEventConnection!
 }
 
 input NetworkSettingInput {

diff --git a/ee/tabby-webserver/src/schema/mod.rs b/ee/tabby-webserver/src/schema/mod.rs
@@ -422,10 +422,15 @@ impl Query {
 
     async fn user_events(
         ctx: &Context,
+
+        // pagination arguments
         after: Option<String>,
         before: Option<String>,
         first: Option<i32>,
         last: Option<i32>,
+
+        // filter arguments
+        users: Option<Vec<ID>>,
         start: DateTime<Utc>,
         end: DateTime<Utc>,
     ) -> Result<Connection<UserEvent>> {
@@ -438,7 +443,15 @@ impl Query {
             |after, before, first, last| async move {
                 ctx.locator
                     .user_event()
-                    .list(after, before, first, last, start, end)
+                    .list(
+                        after,
+                        before,
+                        first,
+                        last,
+                        users.unwrap_or_default(),
+                        start,
+                        end,
+                    )
                     .await
             },
         )

diff --git a/ee/tabby-webserver/src/schema/user_event.rs b/ee/tabby-webserver/src/schema/user_event.rs
@@ -5,7 +5,7 @@ use juniper::{GraphQLEnum, GraphQLObject, ID};
 use super::Context;
 use crate::{juniper::relay::NodeType, schema::Result};
 
-#[derive(GraphQLEnum)]
+#[derive(GraphQLEnum, Debug)]
 pub enum EventKind {
     Completion,
     Select,
@@ -47,6 +47,7 @@ pub trait UserEventService: Send + Sync {
         before: Option<String>,
         first: Option<usize>,
         last: Option<usize>,
+        users: Vec<ID>,
         start: DateTime<Utc>,
         end: DateTime<Utc>,
     ) -> Result<Vec<UserEvent>>;

diff --git a/ee/tabby-webserver/src/service/user_event.rs b/ee/tabby-webserver/src/service/user_event.rs
@@ -1,8 +1,10 @@
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
+use juniper::ID;
 use tabby_db::DbConn;
+use tracing::warn;
 
-use super::graphql_pagination_to_filter;
+use super::{graphql_pagination_to_filter, AsRowid};
 use crate::schema::{
     user_event::{UserEvent, UserEventService},
     Result,
@@ -24,17 +26,90 @@ impl UserEventService for UserEventServiceImpl {
         before: Option<String>,
         first: Option<usize>,
         last: Option<usize>,
+        users: Vec<ID>,
         start: DateTime<Utc>,
         end: DateTime<Utc>,
     ) -> Result<Vec<UserEvent>> {
+        let users = convert_ids(users);
         let (limit, skip_id, backwards) = graphql_pagination_to_filter(after, before, first, last)?;
         let events = self
             .db
-            .list_user_events(limit, skip_id, backwards, start.into(), end.into())
+            .list_user_events(limit, skip_id, backwards, users, start.into(), end.into())
             .await?;
         Ok(events
             .into_iter()
             .map(UserEvent::try_from)
             .collect::<Result<_, _>>()?)
     }
 }
+
+fn convert_ids(ids: Vec<ID>) -> Vec<i64> {
+    ids.into_iter()
+        .filter_map(|id| match id.as_rowid() {
+            Ok(rowid) => Some(rowid),
+            Err(_) => {
+                warn!("Ignoring invalid ID: {}", id);
+                None
+            }
+        })
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use assert_matches::assert_matches;
+    use chrono::{Days, Duration};
+
+    use super::*;
+    use crate::{schema::user_event::EventKind, service::AsID};
+
+    fn timestamp() -> u128 {
+        use std::time::{SystemTime, UNIX_EPOCH};
+        let start = SystemTime::now();
+        start
+            .duration_since(UNIX_EPOCH)
+            .expect("Time went backwards")
+            .as_millis()
+    }
+
+    #[tokio::test]
+    async fn test_list_user_events() {
+        let db = DbConn::new_in_memory().await.unwrap();
+        let user1 = db
+            .create_user("[email protected]".into(), Some("pass".into()), true)
+            .await
+            .unwrap();
+
+        db.create_user_event(user1, "view".into(), timestamp(), "".into())
+            .await
+            .unwrap();
+
+        let user2 = db
+            .create_user("[email protected]".into(), Some("pass".into()), true)
+            .await
+            .unwrap();
+
+        db.create_user_event(user2, "select".into(), timestamp(), "".into())
+            .await
+            .unwrap();
+
+        let svc = create(db);
+        let end = Utc::now() + Duration::days(1);
+        let start = end.checked_sub_days(Days::new(100)).unwrap();
+
+        // List without users should return all events
+        let events = svc
+            .list(None, None, None, None, vec![], start, end)
+            .await
+            .unwrap();
+        assert_eq!(events.len(), 2);
+
+        // Filter with user should return only events for that user
+        let events = svc
+            .list(None, None, None, None, vec![user1.as_id()], start, end)
+            .await
+            .unwrap();
+        assert_eq!(events.len(), 1);
+        assert_matches!(events[0].kind, EventKind::View);
+    }
+}