chore(webserver): Add more languages to analytics (TabbyML#1790)

* chore(webserver): Add more languages to analytics * [autofix.ci] apply automated fixes * Fix languages, implement Other query, and fix analytics bug * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes (attempt 2/3) * [autofix.ci] apply automated fixes (attempt 3/3) * Fix query * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes (attempt 2/3) * [autofix.ci] apply automated fixes (attempt 3/3) * update * [autofix.ci] apply automated fixes --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Meng Zhang <[email protected]>
CatsPunch · Apr 10, 2024 · 7ff97a5 · 7ff97a5
1 parent f38192f
commit 7ff97a5
Show file tree

Hide file tree

Showing 8 changed files with 136 additions and 31 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -38,6 +38,7 @@ derive_builder = "0.12.0"
 futures = "0.3.28"
 async-stream = "0.3.5"
 regex = "1.10.0"
+strum = { version = "0.24", features = ["derive"] }
 thiserror = "1.0.49"
 utoipa = "3.3"
 axum = "0.6"

diff --git a/crates/tabby/Cargo.toml b/crates/tabby/Cargo.toml
@@ -32,7 +32,7 @@ serde_json = { workspace = true }
 tower-http = { version = "0.4.0", features = ["cors", "timeout"] }
 clap = { version = "4.3.0", features = ["derive"] }
 lazy_static = { workspace = true }
-strum = { version = "0.24", features = ["derive"] }
+strum = { workspace = true }
 strfmt = "0.2.4"
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }

diff --git a/ee/tabby-db/src/user_completions.rs b/ee/tabby-db/src/user_completions.rs
@@ -20,6 +20,7 @@ pub struct UserCompletionDAO {
     pub updated_at: DateTimeUtc,
 }
 
+#[derive(FromRow)]
 pub struct UserCompletionDailyStatsDAO {
     pub start: DateTime<Utc>,
     pub completions: i32,
@@ -100,31 +101,44 @@ impl DbConn {
         end: DateTime<Utc>,
         users: Vec<i64>,
         languages: Vec<String>,
+        not_languages: Vec<String>,
     ) -> Result<Vec<UserCompletionDailyStatsDAO>> {
         let users = users
             .iter()
             .map(|u| u.to_string())
             .collect::<Vec<_>>()
             .join(",");
-        let languages = languages.join(",");
-        let res = sqlx::query_as!(
-            UserCompletionDailyStatsDAO,
-            r#"
-        SELECT CAST(STRFTIME('%s', DATE(created_at)) AS TIMESTAMP) as "start!: DateTime<Utc>",
-               SUM(1) as "completions!: i32",
-               SUM(selects) as "selects!: i32"
-        FROM user_completions
-        WHERE created_at >= ?1 AND created_at < ?2
-            AND (?3 = '' OR user_id IN (?3))
-            AND (?4 = '' OR language IN (?4))
-        GROUP BY 1
-        ORDER BY 1 ASC
-        "#,
-            start,
-            end,
-            users,
-            languages
-        )
+        let languages = languages
+            .into_iter()
+            .map(|l| format!("'{}'", l))
+            .collect::<Vec<_>>()
+            .join(",");
+        let not_languages = not_languages
+            .into_iter()
+            .map(|l| format!("'{}'", l))
+            .collect::<Vec<_>>()
+            .join(",");
+
+        let res = sqlx::query_as(&format!(
+                r#"
+            SELECT CAST(STRFTIME('%s', DATE(created_at)) AS TIMESTAMP) as start,
+                   SUM(1) as completions,
+                   SUM(selects) as selects
+            FROM user_completions
+            WHERE created_at >= ?1 AND created_at < ?2
+                AND ({no_selected_users} OR user_id IN ({users}))
+                AND (({no_selected_languages} OR language IN ({languages})) AND (language NOT IN ({not_languages})))
+            GROUP BY 1
+            ORDER BY 1 ASC
+            "#,
+            no_selected_users = users.is_empty(),
+            no_selected_languages = languages.is_empty(),
+        ))
+        .bind(start)
+        .bind(end)
+        .bind(users)
+        .bind(languages)
+        .bind(not_languages)
         .fetch_all(&self.pool)
         .await?;
         Ok(res)

diff --git a/ee/tabby-webserver/Cargo.toml b/ee/tabby-webserver/Cargo.toml
@@ -31,6 +31,7 @@ reqwest = { workspace = true, features = ["json"] }
 rust-embed = "8.0.0"
 serde.workspace = true
 serde_json.workspace = true
+strum.workspace = true
 tabby-common = { path = "../../crates/tabby-common" }
 tabby-db = { path = "../../ee/tabby-db" }
 tarpc = { version = "0.33.0", features = ["serde-transport"] }

diff --git a/ee/tabby-webserver/graphql/schema.graphql b/ee/tabby-webserver/graphql/schema.graphql
@@ -1,6 +1,17 @@
 enum Language {
   RUST
   PYTHON
+  JAVA
+  KOTLIN
+  JAVASCRIPT
+  TYPESCRIPT
+  GO
+  RUBY
+  CSHARP
+  C
+  CPP
+  SOLIDITY
+  OTHER
 }
 
 type GithubRepositoryProviderEdge {

diff --git a/ee/tabby-webserver/src/schema/analytic.rs b/ee/tabby-webserver/src/schema/analytic.rs
@@ -1,8 +1,7 @@
-use std::fmt::Display;
-
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
 use juniper::{GraphQLEnum, GraphQLObject, ID};
+use strum::{EnumIter, IntoEnumIterator};
 
 use crate::schema::Result;
 
@@ -15,19 +14,46 @@ pub struct CompletionStats {
     pub selects: i32,
 }
 
-// FIXME(boxbeam): Adding more languages.
-#[derive(GraphQLEnum, Clone, Debug)]
+#[derive(GraphQLEnum, Clone, Debug, Eq, PartialEq, EnumIter)]
 pub enum Language {
     Rust,
     Python,
+    Java,
+    Kotlin,
+    Javascript,
+    Typescript,
+    Go,
+    Ruby,
+    CSharp,
+    C,
+    CPP,
+    Solidity,
+    Other,
 }
 
-impl Display for Language {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+impl Language {
+    pub fn all_known() -> impl Iterator<Item = Language> {
+        Language::iter().filter(|l| l != &Language::Other)
+    }
+
+    pub fn to_strings(&self) -> impl IntoIterator<Item = String> {
         match self {
-            Language::Rust => write!(f, "rust"),
-            Language::Python => write!(f, "python"),
+            Language::Rust => vec!["rust"],
+            Language::Python => vec!["python"],
+            Language::Java => vec!["java"],
+            Language::Kotlin => vec!["kotlin"],
+            Language::Javascript => vec!["javascript", "javascriptreact"],
+            Language::Typescript => vec!["typescript", "typescriptreact"],
+            Language::Go => vec!["go"],
+            Language::Ruby => vec!["ruby"],
+            Language::CSharp => vec!["csharp"],
+            Language::C => vec!["c"],
+            Language::CPP => vec!["cpp"],
+            Language::Solidity => vec!["solidity"],
+            Language::Other => vec!["other"],
         }
+        .into_iter()
+        .map(|s| s.to_string())
     }
 }
 

diff --git a/ee/tabby-webserver/src/service/analytic.rs b/ee/tabby-webserver/src/service/analytic.rs
@@ -39,13 +39,29 @@ impl AnalyticService for AnalyticServiceImpl {
         start: DateTime<Utc>,
         end: DateTime<Utc>,
         users: Vec<ID>,
-        languages: Vec<Language>,
+        mut languages: Vec<Language>,
     ) -> Result<Vec<CompletionStats>> {
         let users = convert_ids(users);
-        let languages = languages.into_iter().map(|l| l.to_string()).collect();
+
+        let include_other_languages = languages.iter().any(|l| l == &Language::Other);
+        let not_languages = if include_other_languages {
+            Some(Language::all_known().flat_map(|l| l.to_strings()).collect())
+        } else {
+            None
+        };
+
+        languages.retain(|l| l != &Language::Other);
+
+        let languages = languages.into_iter().flat_map(|l| l.to_strings()).collect();
         let stats = self
             .db
-            .compute_daily_stats(start, end, users, languages)
+            .compute_daily_stats(
+                start,
+                end,
+                users,
+                languages,
+                not_languages.unwrap_or_default(),
+            )
             .await?;
         let stats = stats
             .into_iter()
@@ -154,4 +170,39 @@ mod tests {
         assert_eq!(1, stats[0].completions);
         assert_eq!(1, stats[0].selects);
     }
+
+    #[tokio::test]
+    async fn test_other_langs() {
+        let db = DbConn::new_in_memory().await.unwrap();
+
+        let user_id = db
+            .create_user("[email protected]".into(), Some("pass".into()), true)
+            .await
+            .unwrap();
+
+        db.create_user_completion(
+            timestamp(),
+            user_id,
+            "completion_id".into(),
+            "testlang".into(),
+        )
+        .await
+        .unwrap();
+
+        db.create_user_completion(timestamp(), user_id, "completion_id2".into(), "rust".into())
+            .await
+            .unwrap();
+
+        let service = new_analytic_service(db);
+        let end = Utc::now();
+        let start = end.checked_sub_days(Days::new(100)).unwrap();
+
+        let stats = service
+            .daily_stats(start, end, vec![], vec![Language::Other])
+            .await
+            .unwrap();
+
+        assert_eq!(1, stats.len());
+        assert_eq!(1, stats[0].completions);
+    }
 }