feat(db): create merged integration tables (#2075)

* feat(db): create merged integration tables * [autofix.ci] apply automated fixes * Apply suggestions * [autofix.ci] apply automated fixes * Apply suggestions * [autofix.ci] apply automated fixes * Change valid to error * [autofix.ci] apply automated fixes * Revert "Change valid to error" This reverts commit e9ac463. * [autofix.ci] apply automated fixes * Make access_token in update_integration_access_token optional * Reapply "Change valid to error" This reverts commit f97dd9a. * [autofix.ci] apply automated fixes --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
TabbyML · May 9, 2024 · 09c2940 · 09c2940
1 parent 9e480e6
commit 09c2940
Show file tree

Hide file tree

Showing 8 changed files with 707 additions and 297 deletions.
diff --git a/ee/tabby-db/migrations/0029_merged-provider-tables.down.sql b/ee/tabby-db/migrations/0029_merged-provider-tables.down.sql
@@ -0,0 +1,2 @@
+DROP TABLE integration_access_tokens;
+DROP TABLE provided_repositories;
diff --git a/ee/tabby-db/migrations/0029_merged-provider-tables.up.sql b/ee/tabby-db/migrations/0029_merged-provider-tables.up.sql
@@ -0,0 +1,22 @@
+CREATE TABLE integration_access_tokens(
+    id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
+    kind TEXT NOT NULL,
+    display_name TEXT NOT NULL,
+    access_token TEXT,
+    error TEXT,
+    created_at TIMESTAMP NOT NULL DEFAULT (DATETIME('now')),
+    updated_at TIMESTAMP NOT NULL DEFAULT (DATETIME('now'))
+);
+
+CREATE TABLE provided_repositories(
+    id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
+    integration_access_token_id INTEGER NOT NULL,
+    vendor_id TEXT NOT NULL,
+    name TEXT NOT NULL,
+    git_url TEXT NOT NULL,
+    active BOOLEAN NOT NULL DEFAULT FALSE,
+    created_at TIMESTAMP NOT NULL DEFAULT (DATETIME('now')),
+    updated_at TIMESTAMP NOT NULL DEFAULT (DATETIME('now')),
+    FOREIGN KEY (integration_access_token_id) REFERENCES integration_access_tokens(id) ON DELETE CASCADE,
+    CONSTRAINT idx_unique_provider_id_vendor_id UNIQUE (integration_access_token_id, vendor_id)
+);
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
@@ -171,3 +171,24 @@ CREATE TABLE gitlab_provided_repositories(
 CREATE INDEX gitlab_provided_repositories_updated_at ON gitlab_provided_repositories(
   updated_at
 );
+CREATE TABLE integration_access_tokens(
+  id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
+  kind TEXT NOT NULL,
+  display_name TEXT NOT NULL,
+  access_token TEXT,
+  error TEXT,
+  created_at TIMESTAMP NOT NULL DEFAULT(DATETIME('now')),
+  updated_at TIMESTAMP NOT NULL DEFAULT(DATETIME('now'))
+);
+CREATE TABLE provided_repositories(
+  id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
+  integration_access_token_id INTEGER NOT NULL,
+  vendor_id TEXT NOT NULL,
+  name TEXT NOT NULL,
+  git_url TEXT NOT NULL,
+  active BOOLEAN NOT NULL DEFAULT FALSE,
+  created_at TIMESTAMP NOT NULL DEFAULT(DATETIME('now')),
+  updated_at TIMESTAMP NOT NULL DEFAULT(DATETIME('now')),
+  FOREIGN KEY(integration_access_token_id) REFERENCES integration_access_tokens(id) ON DELETE CASCADE,
+  CONSTRAINT idx_unique_provider_id_vendor_id UNIQUE(integration_access_token_id, vendor_id)
+);
diff --git a/ee/tabby-db/schema/schema.svg b/ee/tabby-db/schema/schema.svg
diff --git a/ee/tabby-db/src/integration_access_tokens.rs b/ee/tabby-db/src/integration_access_tokens.rs
@@ -0,0 +1,155 @@
+use anyhow::{anyhow, Result};
+use sqlx::{prelude::FromRow, query, query_as};
+use tabby_db_macros::query_paged_as;
+
+use crate::{DateTimeUtc, DbConn};
+
+#[derive(FromRow)]
+pub struct IntegrationAccessTokenDAO {
+    pub id: i64,
+    pub kind: String,
+    pub error: Option<String>,
+    pub display_name: String,
+    pub access_token: Option<String>,
+    pub created_at: DateTimeUtc,
+    pub updated_at: DateTimeUtc,
+}
+
+impl DbConn {
+    pub async fn create_integration_access_token(
+        &self,
+        kind: &str,
+        name: &str,
+        access_token: &str,
+    ) -> Result<i64> {
+        let res = query!(
+            "INSERT INTO integration_access_tokens(kind, display_name, access_token) VALUES (?, ?, ?);",
+            kind,
+            name,
+            access_token
+        )
+        .execute(&self.pool)
+        .await?;
+        Ok(res.last_insert_rowid())
+    }
+
+    pub async fn get_integration_access_token(&self, id: i64) -> Result<IntegrationAccessTokenDAO> {
+        let provider = query_as!(
+            IntegrationAccessTokenDAO,
+            r#"SELECT
+                id,
+                kind,
+                error,
+                display_name,
+                access_token,
+                created_at AS "created_at: DateTimeUtc",
+                updated_at AS "updated_at: DateTimeUtc"
+            FROM integration_access_tokens WHERE id = ?;"#,
+            id
+        )
+        .fetch_one(&self.pool)
+        .await?;
+        Ok(provider)
+    }
+
+    pub async fn delete_integration_access_token(&self, id: i64) -> Result<()> {
+        let res = query!("DELETE FROM integration_access_tokens WHERE id = ?;", id)
+            .execute(&self.pool)
+            .await?;
+        if res.rows_affected() != 1 {
+            return Err(anyhow!("No integration access token to delete"));
+        }
+        Ok(())
+    }
+
+    pub async fn update_integration_access_token(
+        &self,
+        id: i64,
+        display_name: String,
+        access_token: Option<String>,
+    ) -> Result<()> {
+        let access_token = match access_token {
+            Some(access_token) => Some(access_token),
+            None => self.get_integration_access_token(id).await?.access_token,
+        };
+
+        let res = query!(
+            "UPDATE integration_access_tokens SET display_name = ?, access_token=? WHERE id = ?;",
+            display_name,
+            access_token,
+            id
+        )
+        .execute(&self.pool)
+        .await?;
+
+        if res.rows_affected() != 1 {
+            return Err(anyhow!(
+                "The specified integration access token does not exist"
+            ));
+        }
+
+        Ok(())
+    }
+
+    pub async fn update_integration_access_token_error(
+        &self,
+        id: i64,
+        error: Option<String>,
+    ) -> Result<()> {
+        query!(
+            "UPDATE integration_access_tokens SET updated_at = DATETIME('now'), error = ? WHERE id = ?",
+            error,
+            id
+        )
+        .execute(&self.pool)
+        .await?;
+        Ok(())
+    }
+
+    pub async fn list_integration_access_tokens(
+        &self,
+        ids: Vec<i64>,
+        kind: Option<String>,
+        limit: Option<usize>,
+        skip_id: Option<i32>,
+        backwards: bool,
+    ) -> Result<Vec<IntegrationAccessTokenDAO>> {
+        let mut conditions = vec![];
+
+        let id_condition = (!ids.is_empty()).then(|| {
+            let ids = ids
+                .into_iter()
+                .map(|id| id.to_string())
+                .collect::<Vec<_>>()
+                .join(", ");
+            format!("id in ({ids})")
+        });
+        conditions.extend(id_condition);
+
+        let kind_condition = kind.map(|kind| format!("kind = {kind}"));
+        conditions.extend(kind_condition);
+
+        let condition = (!conditions.is_empty()).then(|| conditions.join(" AND "));
+
+        let providers = query_paged_as!(
+            IntegrationAccessTokenDAO,
+            "integration_access_tokens",
+            [
+                "id",
+                "kind",
+                "error",
+                "display_name",
+                "access_token",
+                "created_at" as "created_at: DateTimeUtc",
+                "updated_at" as "updated_at: DateTimeUtc"
+            ],
+            limit,
+            skip_id,
+            backwards,
+            condition
+        )
+        .fetch_all(&self.pool)
+        .await?;
+        Ok(providers)
+    }
+}
diff --git a/ee/tabby-db/src/lib.rs b/ee/tabby-db/src/lib.rs
@@ -30,10 +30,12 @@ pub mod cache;
 mod email_setting;
 mod github_repository_provider;
 mod gitlab_repository_provider;
+mod integration_access_tokens;
 mod invitations;
 mod job_runs;
 mod oauth_credential;
 mod password_reset;
+mod provided_repositories;
 mod refresh_tokens;
 mod repositories;
 mod server_setting;

diff --git a/ee/tabby-db/src/provided_repositories.rs b/ee/tabby-db/src/provided_repositories.rs
@@ -0,0 +1,130 @@
+use anyhow::{anyhow, Result};
+use sqlx::{prelude::FromRow, query, query_as};
+use tabby_db_macros::query_paged_as;
+
+use crate::{DateTimeUtc, DbConn};
+
+#[derive(FromRow)]
+pub struct ProvidedRepositoryDAO {
+    pub id: i64,
+    pub vendor_id: String,
+    pub integration_access_token_id: i64,
+    pub name: String,
+    pub git_url: String,
+    pub active: bool,
+    pub created_at: DateTimeUtc,
+    pub updated_at: DateTimeUtc,
+}
+
+impl DbConn {
+    pub async fn upsert_provided_repository(
+        &self,
+        integration_access_token_id: i64,
+        vendor_id: String,
+        name: String,
+        git_url: String,
+    ) -> Result<i64> {
+        let res = query!(
+            "INSERT INTO provided_repositories (integration_access_token_id, vendor_id, name, git_url) VALUES ($1, $2, $3, $4)
+                ON CONFLICT(integration_access_token_id, vendor_id) DO UPDATE SET name = $3, git_url = $4, updated_at = DATETIME('now')",
+            integration_access_token_id,
+            vendor_id,
+            name,
+            git_url
+        ).execute(&self.pool).await?;
+        Ok(res.last_insert_rowid())
+    }
+
+    pub async fn delete_outdated_provided_repositories(
+        &self,
+        integration_access_token_id: i64,
+        cutoff_timestamp: DateTimeUtc,
+    ) -> Result<usize> {
+        let res = query!(
+            "DELETE FROM provided_repositories WHERE integration_access_token_id = ? AND updated_at < ?;",
+            integration_access_token_id,
+            cutoff_timestamp
+        ).execute(&self.pool).await?;
+        Ok(res.rows_affected() as usize)
+    }
+
+    pub async fn get_provided_repository(&self, id: i64) -> Result<ProvidedRepositoryDAO> {
+        let repo = query_as!(
+            ProvidedRepositoryDAO,
+            "SELECT id, vendor_id, name, git_url, active, integration_access_token_id, created_at, updated_at FROM provided_repositories WHERE id = ?",
+            id
+        )
+        .fetch_one(&self.pool)
+        .await?;
+        Ok(repo)
+    }
+
+    pub async fn list_provided_repositories(
+        &self,
+        provider_ids: Vec<i64>,
+        kind: Option<String>,
+        active: Option<bool>,
+        limit: Option<usize>,
+        skip_id: Option<i32>,
+        backwards: bool,
+    ) -> Result<Vec<ProvidedRepositoryDAO>> {
+        let mut conditions = vec![];
+
+        let provider_ids = provider_ids
+            .into_iter()
+            .map(|id| id.to_string())
+            .collect::<Vec<_>>()
+            .join(", ");
+        if !provider_ids.is_empty() {
+            conditions.push(format!("access_token_provider_id IN ({provider_ids})"));
+        }
+
+        let active_filter = active.map(|active| format!("active = {active}"));
+        conditions.extend(active_filter);
+
+        let kind_filter = kind.map(|kind| format!("kind = {kind}"));
+        conditions.extend(kind_filter);
+
+        let condition = (!conditions.is_empty()).then(|| conditions.join(" AND "));
+
+        let repos = query_paged_as!(
+            ProvidedRepositoryDAO,
+            "provided_repositories",
+            [
+                "id",
+                "vendor_id",
+                "name",
+                "git_url",
+                "active",
+                "integration_access_token_id",
+                "created_at" as "created_at: DateTimeUtc",
+                "updated_at" as "updated_at: DateTimeUtc"
+            ],
+            limit,
+            skip_id,
+            backwards,
+            condition
+        )
+        .fetch_all(&self.pool)
+        .await?;
+        Ok(repos)
+    }
+
+    pub async fn update_provided_repository_active(&self, id: i64, active: bool) -> Result<()> {
+        let not_active = !active;
+        let res = query!(
+            "UPDATE provided_repositories SET active = ? WHERE id = ? AND active = ?",
+            active,
+            id,
+            not_active
+        )
+        .execute(&self.pool)
+        .await?;
+
+        if res.rows_affected() != 1 {
+            return Err(anyhow!("Repository active status was not changed"));
+        }
+
+        Ok(())
+    }
+}
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		DROP TABLE integration_access_tokens;
		DROP TABLE provided_repositories;