Perform status update logic in db layer

TabbyML · Apr 26, 2024 · ff7c5b8 · ff7c5b8
1 parent 11673e8
commit ff7c5b8
Show file tree

Hide file tree

Showing 6 changed files with 73 additions and 13 deletions.
diff --git a/ee/tabby-db/src/github_repository_provider.rs b/ee/tabby-db/src/github_repository_provider.rs
@@ -96,11 +96,12 @@ impl DbConn {
         Ok(())
     }
 
-    pub async fn update_github_provider_synced_at(&self, id: i64, success: bool) -> Result<()> {
+    pub async fn update_github_provider_sync_status(&self, id: i64, success: bool) -> Result<()> {
         let time = success.then_some(DateTimeUtc::now());
         query!(
-            "UPDATE github_repository_provider SET synced_at = ? WHERE id = ?",
+            "UPDATE github_repository_provider SET synced_at = ?, access_token = IIF(?, access_token, NULL) WHERE id = ?",
             time,
+            success,
             id
         )
         .execute(&self.pool)

diff --git a/ee/tabby-db/src/gitlab_repository_provider.rs b/ee/tabby-db/src/gitlab_repository_provider.rs
@@ -96,11 +96,12 @@ impl DbConn {
         Ok(())
     }
 
-    pub async fn update_gitlab_provider_synced_at(&self, id: i64, success: bool) -> Result<()> {
+    pub async fn update_gitlab_provider_sync_status(&self, id: i64, success: bool) -> Result<()> {
         let time = success.then_some(DateTimeUtc::now());
         query!(
-            "UPDATE gitlab_repository_provider SET synced_at = ? WHERE id = ?",
+            "UPDATE gitlab_repository_provider SET synced_at = ?, access_token = IIF(?, access_token, NULL) WHERE id = ?",
             time,
+            success,
             id
         )
         .execute(&self.pool)

diff --git a/ee/tabby-webserver/src/schema/types.rs b/ee/tabby-webserver/src/schema/types.rs
@@ -32,7 +32,7 @@ pub enum RepositoryProviderStatus {
 }
 
 impl RepositoryProviderStatus {
-    pub fn get_status(access_token_set: bool, synced_at_set: bool) -> Self {
+    pub fn new(access_token_set: bool, synced_at_set: bool) -> Self {
         match (access_token_set, synced_at_set) {
             (true, true) => RepositoryProviderStatus::Ready,
             (true, false) => RepositoryProviderStatus::Pending,

diff --git a/ee/tabby-webserver/src/service/dao.rs b/ee/tabby-webserver/src/service/dao.rs
@@ -131,7 +131,7 @@ impl From<GithubRepositoryProviderDAO> for GithubRepositoryProvider {
         Self {
             display_name: value.display_name,
             id: value.id.as_id(),
-            status: RepositoryProviderStatus::get_status(
+            status: RepositoryProviderStatus::new(
                 value.access_token.is_some(),
                 value.synced_at.is_some(),
             ),
@@ -158,7 +158,7 @@ impl From<GitlabRepositoryProviderDAO> for GitlabRepositoryProvider {
         Self {
             display_name: value.display_name,
             id: value.id.as_id(),
-            status: RepositoryProviderStatus::get_status(
+            status: RepositoryProviderStatus::new(
                 value.access_token.is_some(),
                 value.synced_at.is_some(),
             ),

diff --git a/ee/tabby-webserver/src/service/github_repository_provider.rs b/ee/tabby-webserver/src/service/github_repository_provider.rs
@@ -182,12 +182,8 @@ impl GithubRepositoryProviderService for GithubRepositoryProviderServiceImpl {
         id: ID,
         success: bool,
     ) -> Result<()> {
-        let id = id.as_rowid()?;
-        if !success {
-            self.db.reset_github_provider_access_token(id).await?;
-        }
         self.db
-            .update_github_provider_synced_at(id, success)
+            .update_github_provider_sync_status(id.as_rowid()?, success)
             .await?;
         Ok(())
     }
@@ -385,6 +381,37 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_sync_status() {
+        let db = DbConn::new_in_memory().await.unwrap();
+        let service = create(db.clone());
+
+        let provider_id = db
+            .create_github_provider("provider1".into(), "token".into())
+            .await
+            .unwrap();
+
+        service
+            .update_github_repository_provider_sync_status(provider_id.as_id(), true)
+            .await
+            .unwrap();
+
+        let provider = db.get_github_provider(provider_id).await.unwrap();
+
+        assert!(provider.access_token.is_some());
+        assert!(provider.synced_at.is_some());
+
+        service
+            .update_github_repository_provider_sync_status(provider_id.as_id(), false)
+            .await
+            .unwrap();
+
+        let provider = db.get_github_provider(provider_id).await.unwrap();
+
+        assert!(provider.access_token.is_none());
+        assert!(provider.synced_at.is_none());
+    }
+
     #[tokio::test]
     async fn test_delete_outdated_repos() {
         let db = DbConn::new_in_memory().await.unwrap();

diff --git a/ee/tabby-webserver/src/service/gitlab_repository_provider.rs b/ee/tabby-webserver/src/service/gitlab_repository_provider.rs
@@ -186,7 +186,7 @@ impl GitlabRepositoryProviderService for GitlabRepositoryProviderServiceImpl {
         success: bool,
     ) -> Result<()> {
         self.db
-            .update_gitlab_provider_synced_at(id.as_rowid()?, success)
+            .update_gitlab_provider_sync_status(id.as_rowid()?, success)
             .await?;
         Ok(())
     }
@@ -353,6 +353,37 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_sync_status() {
+        let db = DbConn::new_in_memory().await.unwrap();
+        let service = create(db.clone());
+
+        let provider_id = db
+            .create_gitlab_provider("provider1".into(), "token".into())
+            .await
+            .unwrap();
+
+        service
+            .update_gitlab_repository_provider_sync_status(provider_id.as_id(), true)
+            .await
+            .unwrap();
+
+        let provider = db.get_gitlab_provider(provider_id).await.unwrap();
+
+        assert!(provider.access_token.is_some());
+        assert!(provider.synced_at.is_some());
+
+        service
+            .update_gitlab_repository_provider_sync_status(provider_id.as_id(), false)
+            .await
+            .unwrap();
+
+        let provider = db.get_gitlab_provider(provider_id).await.unwrap();
+
+        assert!(provider.access_token.is_none());
+        assert!(provider.synced_at.is_none());
+    }
+
     #[tokio::test]
     async fn test_provided_git_urls() {
         let db = DbConn::new_in_memory().await.unwrap();