refactor(db): Update UserDAO to use query_as! and use i64 IDs in most pl…

ee/tabby-db/src/job_runs.rs

@@ -39,7 +39,7 @@ impl DbConn {
         Ok(rowid as i32)
     }
 
-    pub async fn update_job_stdout(&self, job_id: i32, stdout: String) -> Result<()> {
+    pub async fn update_job_stdout(&self, job_id: i64, stdout: String) -> Result<()> {
         query!(
             r#"UPDATE job_runs SET stdout = stdout || ?, updated_at = datetime('now') WHERE id = ?"#,
             stdout,
@@ -48,7 +48,7 @@ impl DbConn {
         Ok(())
     }
 
-    pub async fn update_job_stderr(&self, job_id: i32, stderr: String) -> Result<()> {
+    pub async fn update_job_stderr(&self, job_id: i64, stderr: String) -> Result<()> {
         query!(
             r#"UPDATE job_runs SET stderr = stderr || ?, updated_at = datetime('now') WHERE id = ?"#,
             stderr,
@@ -57,7 +57,7 @@ impl DbConn {
         Ok(())
     }
 
-    pub async fn update_job_status(&self, job_id: i32, exit_code: i32) -> Result<()> {
+    pub async fn update_job_status(&self, job_id: i64, exit_code: i32) -> Result<()> {
         query!(
             r#"UPDATE job_runs SET end_ts = datetime('now'), exit_code = ?, updated_at = datetime('now') WHERE id = ?"#,
             exit_code,

ee/tabby-db/src/lib.rs

@@ -82,7 +82,7 @@
     select.as_string()
 }
 
 fn make_pagination_query(
     table_name: &str,
     field_names: &[&str],
     limit: Option<usize>,
@@ -355,7 +355,7 @@
 pub mod testutils {
     use super::*;
 
-    pub async fn create_user(conn: &DbConn) -> i32 {
+    pub async fn create_user(conn: &DbConn) -> i64 {
         let email: &str = "[email protected]";
         let password: &str = "123456789";
         conn.create_user(email.to_string(), Some(password.to_string()), true)

ee/tabby-db/src/password_reset.rs

@@ -66,15 +66,15 @@ impl DbConn {
             .ok_or_else(|| anyhow!("Invalid code"))?;
 
         let user_res = self
-            .get_user(password_reset.user_id as i32)
+            .get_user(password_reset.user_id)
             .await?
             .filter(|user| user.active)
             .ok_or_else(|| anyhow!("Invalid code"))?;
 
         if Utc::now().signed_duration_since(*password_reset.created_at) > Duration::minutes(15) {
             Err(anyhow!("Invalid code"))
         } else {
-            Ok(user_res.id as i64)
+            Ok(user_res.id)
         }
     }
 

ee/tabby-db/src/refresh_tokens.rs

@@ -28,7 +28,7 @@ impl RefreshTokenDAO {
 
 /// db read/write operations for `refresh_tokens` table
 impl DbConn {
-    pub async fn create_refresh_token(&self, user_id: i32) -> Result<String> {
+    pub async fn create_refresh_token(&self, user_id: i64) -> Result<String> {
         let token = generate_refresh_token(0);
         let res = query!(
             r#"INSERT INTO refresh_tokens (user_id, token, expires_at) VALUES (?, ?, datetime('now', '+7 days'))"#,
@@ -43,8 +43,8 @@ impl DbConn {
         Ok(token)
     }
 
-    pub async fn renew_refresh_token(&self, id: i32, old: &str) -> Result<String> {
-        let new = generate_refresh_token(id as i64);
+    pub async fn renew_refresh_token(&self, id: i64, old: &str) -> Result<String> {
+        let new = generate_refresh_token(id);
         let res = query!(
             "UPDATE refresh_tokens SET token = $1, expires_at = datetime('now', '+7 days') WHERE token = $2 AND id = $3",
             new,
@@ -82,7 +82,7 @@ impl DbConn {
         Ok(token)
     }
 
-    pub async fn delete_tokens_by_user_id(&self, id: i32) -> Result<()> {
+    pub async fn delete_tokens_by_user_id(&self, id: i64) -> Result<()> {
         query!("DELETE FROM refresh_tokens WHERE user_id = ?", id)
             .execute(&self.pool)
             .await?;

ee/tabby-db/src/user_completions.rs

@@ -24,7 +24,7 @@ impl DbConn {
     pub async fn create_user_completion(
         &self,
         ts: u128,
-        user_id: i32,
+        user_id: i64,
         completion_id: String,
         language: String,
     ) -> Result<i32> {

ee/tabby-db/src/users.rs

@@ -1,18 +1,18 @@
 use anyhow::{anyhow, bail, Result};
-use chrono::{DateTime, Utc};
-use sqlx::{query, query_scalar, FromRow};
+use sqlx::{query, query_as, query_scalar, FromRow};
+use tabby_db_macros::query_paged_as;
 use uuid::Uuid;
 
 use super::DbConn;
-use crate::{make_pagination_query, SQLXResultExt};
+use crate::{DateTimeUtc, SQLXResultExt};
 
 #[allow(unused)]
 #[derive(FromRow)]
 pub struct UserDAO {
-    pub created_at: DateTime<Utc>,
-    pub updated_at: DateTime<Utc>,
+    pub created_at: DateTimeUtc,
+    pub updated_at: DateTimeUtc,
 
-    pub id: i32,
+    pub id: i64,
     pub email: String,
     pub password_encrypted: Option<String>,
     pub is_admin: bool,
@@ -22,15 +22,19 @@
     pub active: bool,
 }
 
-static OWNER_USER_ID: i32 = 1;
+static OWNER_USER_ID: i64 = 1;
 
-impl UserDAO {
-    fn select(clause: &str) -> String {
-        r#"SELECT id, email, password_encrypted, is_admin, created_at, updated_at, auth_token, active FROM users WHERE "#
-            .to_owned()
-            + clause
+macro_rules! select {
+    ($str:literal $(,)? $($val:expr),*) => {
+        query_as!(
+            UserDAO,
+            r#"SELECT id as "id!", email, password_encrypted, is_admin, created_at as "created_at!", updated_at as "updated_at!", auth_token, active FROM users WHERE "# + $str,
+            $($val),*
+        )
     }
+}
 
+impl UserDAO {
     pub fn is_owner(&self) -> bool {
         self.id == OWNER_USER_ID
     }
@@ -43,7 +47,7 @@
         email: String,
         password_encrypted: Option<String>,
         is_admin: bool,
-    ) -> Result<i32> {
+    ) -> Result<i64> {
         self.create_user_impl(email, password_encrypted, is_admin, None)
             .await
     }
@@ -53,8 +57,8 @@
         email: String,
         password_encrypted: Option<String>,
         is_admin: bool,
-        invitation_id: i32,
-    ) -> Result<i32> {
+        invitation_id: i64,
+    ) -> Result<i64> {
         self.create_user_impl(email, password_encrypted, is_admin, Some(invitation_id))
             .await
     }
@@ -64,8 +68,8 @@
         email: String,
         password_encrypted: Option<String>,
         is_admin: bool,
-        invitation_id: Option<i32>,
-    ) -> Result<i32> {
+        invitation_id: Option<i64>,
+    ) -> Result<i64> {
         let mut transaction = self.pool.begin().await?;
         if let Some(invitation_id) = invitation_id {
             query!("DELETE FROM invitations WHERE id = ?", invitation_id)
@@ -84,31 +88,25 @@
             self.cache.active_admin_count.invalidate().await;
         }
 
-        Ok(res.last_insert_rowid() as i32)
+        Ok(res.last_insert_rowid())
     }
 
-    pub async fn get_user(&self, id: i32) -> Result<Option<UserDAO>> {
-        let user = sqlx::query_as(&UserDAO::select("id = ?"))
-            .bind(id)
-            .fetch_optional(&self.pool)
-            .await?;
+    pub async fn get_user(&self, id: i64) -> Result<Option<UserDAO>> {
+        let user = select!("id = ?", id).fetch_optional(&self.pool).await?;
 
         Ok(user)
     }
 
     pub async fn get_user_by_email(&self, email: &str) -> Result<Option<UserDAO>> {
-        let user = sqlx::query_as(&UserDAO::select("email = ?"))
-            .bind(email)
+        let user = select!("email = ?", email)
             .fetch_optional(&self.pool)
             .await?;
 
         Ok(user)
     }
 
     pub async fn list_admin_users(&self) -> Result<Vec<UserDAO>> {
-        let users = sqlx::query_as(&UserDAO::select("is_admin"))
-            .fetch_all(&self.pool)
-            .await?;
+        let users = select!("is_admin").fetch_all(&self.pool).await?;
 
         Ok(users)
     }
@@ -119,24 +117,26 @@
         skip_id: Option<i32>,
         backwards: bool,
     ) -> Result<Vec<UserDAO>> {
-        let query = make_pagination_query(
+        let users = query_paged_as!(
+            UserDAO,
             "users",
-            &[
-                "id",
+            [
+                "id"!,
                 "email",
                 "password_encrypted",
                 "is_admin",
-                "created_at",
-                "updated_at",
+                "created_at"!,
+                "updated_at"!,
                 "auth_token",
-                "active",
+                "active"
             ],
             limit,
             skip_id,
-            backwards,
-        );
+            backwards
+        )
+        .fetch_all(&self.pool)
+        .await?;
 
-        let users = sqlx::query_as(&query).fetch_all(&self.pool).await?;
         Ok(users)
     }
 
@@ -157,7 +157,7 @@
         Ok(id)
     }
 
-    pub async fn reset_user_auth_token_by_id(&self, id: i32) -> Result<()> {
+    pub async fn reset_user_auth_token_by_id(&self, id: i64) -> Result<()> {
         let updated_at = chrono::Utc::now();
         let token = generate_auth_token();
         query!(
@@ -172,7 +172,7 @@
         Ok(())
     }
 
-    pub async fn update_user_active(&self, id: i32, active: bool) -> Result<()> {
+    pub async fn update_user_active(&self, id: i64, active: bool) -> Result<()> {
         let not_active = !active;
         let changed = query!(
             "UPDATE users SET active = ? WHERE id = ? AND active = ?",
@@ -191,7 +191,7 @@
         Ok(())
     }
 
-    pub async fn update_user_role(&self, id: i32, is_admin: bool) -> Result<()> {
+    pub async fn update_user_role(&self, id: i64, is_admin: bool) -> Result<()> {
         let not_admin = !is_admin;
         let changed = query!(
             "UPDATE users SET is_admin = ? WHERE id = ? AND is_admin = ?",
@@ -210,7 +210,7 @@
         }
     }
 
-    pub async fn update_user_password(&self, id: i32, password_encrypted: String) -> Result<()> {
+    pub async fn update_user_password(&self, id: i64, password_encrypted: String) -> Result<()> {
         query!(
             "UPDATE users SET password_encrypted = ? WHERE id = ?",
             password_encrypted,
@@ -221,14 +221,14 @@
         Ok(())
     }
 
-    pub async fn update_user_avatar(&self, id: i32, avatar: Option<Box<[u8]>>) -> Result<()> {
+    pub async fn update_user_avatar(&self, id: i64, avatar: Option<Box<[u8]>>) -> Result<()> {
         query!("UPDATE users SET avatar = ? WHERE id = ?;", avatar, id)
             .execute(&self.pool)
             .await?;
         Ok(())
     }
 
-    pub async fn get_user_avatar(&self, id: i32) -> Result<Option<Box<[u8]>>> {
+    pub async fn get_user_avatar(&self, id: i64) -> Result<Option<Box<[u8]>>> {
         let avatar = query_scalar!("SELECT avatar FROM users WHERE id = ?", id)
             .fetch_one(&self.pool)
             .await?;
@@ -341,7 +341,7 @@
     async fn test_list_users_with_filter() {
         let conn = DbConn::new_in_memory().await.unwrap();
 
-        let empty: Vec<i32> = vec![];
+        let empty: Vec<i64> = vec![];
         let to_ids = |users: Vec<UserDAO>| users.into_iter().map(|u| u.id).collect::<Vec<_>>();
 
         // empty