refactor

TabbyML · Dec 1, 2023 · bcd2f38 · bcd2f38
1 parent e9c9460
commit bcd2f38
Showing 1 changed file with 26 additions and 24 deletions.
diff --git a/ee/tabby-webserver/src/db.rs b/ee/tabby-webserver/src/db.rs
@@ -2,7 +2,7 @@ use std::{path::PathBuf, sync::Arc};
 
 use anyhow::Result;
 use lazy_static::lazy_static;
-use rusqlite::{params, OptionalExtension};
+use rusqlite::{params, OptionalExtension, Row};
 use rusqlite_migration::{AsyncMigrations, M};
 use tabby_common::path::tabby_root;
 use tokio_rusqlite::Connection;
@@ -47,6 +47,25 @@ pub struct User {
     pub is_admin: bool,
 }
 
+impl User {
+    fn select(clause: &str) -> String {
+        r#"SELECT id, email, password_encrypted, is_admin, created_at, updated_at FROM users WHERE "#
+            .to_owned()
+            + clause
+    }
+
+    fn from_row(row: &Row<'_>) -> std::result::Result<User, rusqlite::Error> {
+        Ok(User {
+            id: row.get(0)?,
+            email: row.get(1)?,
+            password_encrypted: row.get(2)?,
+            is_admin: row.get(3)?,
+            created_at: row.get(4)?,
+            updated_at: row.get(5)?,
+        })
+    }
+}
+
 async fn db_path() -> Result<PathBuf> {
     let db_dir = tabby_root().join("ee");
     tokio::fs::create_dir_all(db_dir.clone()).await?;
@@ -156,19 +175,11 @@ impl DbConn {
             .conn
             .call(move |c| {
                 c.query_row(
-                    r#"SELECT id, email, password_encrypted, is_admin, created_at, updated_at FROM users WHERE email = ?"#,
+                    User::select("email = ?").as_str(),
                     params![email],
-                    |row| {
-                        Ok(User {
-                            id: row.get(0)?,
-                            email: row.get(1)?,
-                            password_encrypted: row.get(2)?,
-                            is_admin: row.get(3)?,
-                            created_at: row.get(4)?,
-                            updated_at: row.get(5)?,
-                        })
-                    },
-                ).optional()
+                    User::from_row,
+                )
+                .optional()
             })
             .await?;
 
@@ -179,17 +190,8 @@ impl DbConn {
         let users = self
             .conn
             .call(move |c| {
-                let mut stmt = c.prepare(r#"SELECT id, email, password_encrypted, is_admin, created_at, updated_at FROM users WHERE is_admin"#)?;
-                let user_iter = stmt.query_map([], |row| {
-                    Ok(User {
-                            id: row.get(0)?,
-                            email: row.get(1)?,
-                            password_encrypted: row.get(2)?,
-                            is_admin: row.get(3)?,
-                            created_at: row.get(4)?,
-                            updated_at: row.get(5)?,
-                        })
-                })?;
+                let mut stmt = c.prepare(&User::select("is_admin"))?;
+                let user_iter = stmt.query_map([], User::from_row)?;
                 Ok(user_iter.filter_map(|x| x.ok()).collect::<Vec<_>>())
             })
             .await?;