From 30fe61d440c4fcd38615241882f969fe95a08356 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9A=D0=93=D0=91=20Lieu?=
 <60683892+hanhvs@users.noreply.github.com>
Date: Sat, 19 Sep 2026 08:47:58 +0700
Subject: [PATCH] fix(sql-file): extend relational constraint bypass to
 Postgres and SQL Server

---
 .../sql-file/SqlFileExecutionDialog.vue       |  15 +-
 crates/dbx-core/src/data/sql_file_import.rs   | 290 +++++++++++++++++-
 crates/dbx-sql/src/sql.rs                     |  13 +-
 3 files changed, 310 insertions(+), 8 deletions(-)

diff --git a/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue b/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue
index bff651096f..2da79970df 100644
--- a/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue
+++ b/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue
@@ -134,15 +134,24 @@ function resetPerFileState() {
 }
 
 const sqlConnections = computed(() => store.connections.filter((c) => !["redis", "mongodb", "elasticsearch", "easysearch", "meilisearch", "qdrant", "milvus", "weaviate", "chromadb", "etcd", "zookeeper", "consul", "mq", "nacos"].includes(c.db_type)));
-// Mirrors the core executor gate (`supports_connection_level_database_bootstrap_target`): the
-// MySQL-family types it runs for, so the constraint toggle appears wherever the backend honors it.
+// Mirrors the core executor gate (`relational_constraint_bypass_kind` in
+// sql_file_import.rs): MySQL-family types use the session-scoped
+// FOREIGN_KEY_CHECKS toggle, PostgreSQL-family types use DISABLE/ENABLE
+// TRIGGER ALL, and SQL Server uses NOCHECK/CHECK CONSTRAINT ALL. The toggle
+// appears wherever the backend implements one of those mechanisms.
 const MYSQL_BOOTSTRAP_IMPORT_TYPES = new Set(["mysql", "doris", "starrocks", "goldendb"]);
 const MYSQL_BOOTSTRAP_IMPORT_PROFILES = new Set(["mariadb", "tidb", "oceanbase", "custom_mysql", "doris", "starrocks", "selectdb", "goldendb"]);
+const POSTGRES_CONSTRAINT_BYPASS_TYPES = new Set(["postgres", "gaussdb", "opengauss"]);
 const isMysqlCompatibleTarget = computed(() => {
   const config = store.getConfig(connectionId.value);
   if (!config) return false;
   return MYSQL_BOOTSTRAP_IMPORT_TYPES.has(config.db_type) || (!!config.driver_profile && MYSQL_BOOTSTRAP_IMPORT_PROFILES.has(config.driver_profile.toLowerCase()));
 });
+const supportsRelationalConstraintBypass = computed(() => {
+  const config = store.getConfig(connectionId.value);
+  if (!config) return false;
+  return isMysqlCompatibleTarget.value || POSTGRES_CONSTRAINT_BYPASS_TYPES.has(config.db_type) || config.db_type === "sqlserver";
+});
 
 const selectedConnection = computed(() => sqlConnections.value.find((c) => c.id === connectionId.value));
 
@@ -799,7 +808,7 @@ watch(
             <Square v-else class="w-3.5 h-3.5 text-muted-foreground/40 shrink-0" />
             {{ t("sqlFile.continueOnError") }}
           </button>
-          <button v-if="isMysqlCompatibleTarget" type="button" class="flex items-center gap-2 text-xs text-left" :disabled="running" @click="skipRelationalConstraints = !skipRelationalConstraints">
+          <button v-if="supportsRelationalConstraintBypass" type="button" class="flex items-center gap-2 text-xs text-left" :disabled="running" @click="skipRelationalConstraints = !skipRelationalConstraints">
             <CheckSquare v-if="skipRelationalConstraints" class="w-3.5 h-3.5 text-primary shrink-0" />
             <Square v-else class="w-3.5 h-3.5 text-muted-foreground/40 shrink-0" />
             {{ t("sqlFile.skipRelationalConstraints") }}
diff --git a/crates/dbx-core/src/data/sql_file_import.rs b/crates/dbx-core/src/data/sql_file_import.rs
index 3c07b5292f..34f059f8ba 100644
--- a/crates/dbx-core/src/data/sql_file_import.rs
+++ b/crates/dbx-core/src/data/sql_file_import.rs
@@ -437,6 +437,152 @@ impl MySqlSqlFileExecutor {
     }
 }
 
+/// Which database-family mechanism `skip_relational_constraints` uses. Unlike
+/// MySQL's session-scoped `FOREIGN_KEY_CHECKS`, the PostgreSQL and SQL Server
+/// mechanisms below are catalog-level (`ALTER TABLE ... DISABLE TRIGGER ALL` /
+/// `NOCHECK CONSTRAINT ALL`): the effect is visible to every connection in the
+/// pool immediately, so neither needs [`MySqlSqlFileExecutor`]'s pinned-connection
+/// bookkeeping. Both are applied to every table in the database, not just the
+/// tables referenced by the imported file, matching "temporarily disable
+/// relational constraints" at the database scope the option describes.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum RelationalConstraintBypassKind {
+    Mysql,
+    Postgres,
+    SqlServer,
+}
+
+fn relational_constraint_bypass_kind(target: &SqlFileImportTarget) -> Option<RelationalConstraintBypassKind> {
+    if crate::sql::supports_connection_level_database_bootstrap_target(
+        &target.db_type,
+        target.driver_profile.as_deref(),
+    ) {
+        return Some(RelationalConstraintBypassKind::Mysql);
+    }
+    match target.db_type {
+        DatabaseType::Postgres | DatabaseType::Gaussdb | DatabaseType::OpenGauss => {
+            Some(RelationalConstraintBypassKind::Postgres)
+        }
+        DatabaseType::SqlServer => Some(RelationalConstraintBypassKind::SqlServer),
+        _ => None,
+    }
+}
+
+// `pg_tables` lists every table the connection can see, including system
+// catalog tables owned by `pg_catalog`/`information_schema` (unlike
+// `information_schema.tables`, it has no built-in system-schema filter). Those
+// must be excluded explicitly or the loop fails trying to alter tables the
+// user cannot own, aborting the whole bypass. Temp schemas are excluded for
+// the same reason `POSTGRES_SCHEMA_INFOS_HIDE_SYSTEM_SQL` (schema.rs) does.
+// Foreign tables and empty partitioned parents are intentionally out of
+// scope: `DISABLE TRIGGER ALL` only matters for plain tables that can carry
+// FK/RI triggers.
+// Altering a table's triggers requires table ownership (or superuser). A
+// multi-schema database can easily contain tables the importing role does not
+// own; wrapping each ALTER in its own sub-transaction (BEGIN/EXCEPTION) means
+// one inaccessible table is skipped instead of aborting the bypass — and the
+// import it gates — for every table the role *can* alter. If every table hit
+// `insufficient_privilege` (a common case: `DISABLE TRIGGER` also affects
+// internal FK/RI triggers, which PostgreSQL restricts to the table owner or a
+// superuser), the toggle would otherwise be a silent no-op the caller cannot
+// distinguish from "there were no tables"; raise instead so the import fails
+// loudly rather than running with constraints the user believed were disabled.
+const POSTGRES_DISABLE_ALL_RELATIONAL_TRIGGERS_SQL: &str = "\
+    DO $dbx_disable_constraints$ \
+    DECLARE dbx_rel record; \
+    DECLARE dbx_total integer := 0; \
+    DECLARE dbx_disabled integer := 0; \
+    BEGIN \
+        FOR dbx_rel IN SELECT schemaname, tablename FROM pg_tables \
+            WHERE schemaname NOT IN ('pg_catalog', 'information_schema', 'pg_toast') \
+              AND schemaname NOT LIKE 'pg\\_temp\\_%' AND schemaname NOT LIKE 'pg\\_toast\\_temp\\_%' \
+        LOOP \
+            dbx_total := dbx_total + 1; \
+            BEGIN \
+                EXECUTE format('ALTER TABLE %I.%I DISABLE TRIGGER ALL', dbx_rel.schemaname, dbx_rel.tablename); \
+                dbx_disabled := dbx_disabled + 1; \
+            EXCEPTION WHEN insufficient_privilege THEN NULL; \
+            END; \
+        END LOOP; \
+        IF dbx_total > 0 AND dbx_disabled = 0 THEN \
+            RAISE EXCEPTION 'Could not disable relational constraints on any of % table(s): the connection role lacks owner/superuser privileges required to disable internal foreign-key triggers', dbx_total; \
+        END IF; \
+    END $dbx_disable_constraints$;";
+
+const POSTGRES_ENABLE_ALL_RELATIONAL_TRIGGERS_SQL: &str = "\
+    DO $dbx_enable_constraints$ \
+    DECLARE dbx_rel record; \
+    BEGIN \
+        FOR dbx_rel IN SELECT schemaname, tablename FROM pg_tables \
+            WHERE schemaname NOT IN ('pg_catalog', 'information_schema', 'pg_toast') \
+              AND schemaname NOT LIKE 'pg\\_temp\\_%' AND schemaname NOT LIKE 'pg\\_toast\\_temp\\_%' \
+        LOOP \
+            BEGIN \
+                EXECUTE format('ALTER TABLE %I.%I ENABLE TRIGGER ALL', dbx_rel.schemaname, dbx_rel.tablename); \
+            EXCEPTION WHEN insufficient_privilege THEN NULL; \
+            END; \
+        END LOOP; \
+    END $dbx_enable_constraints$;";
+
+// `sys.tables`/`sys.schemas` (rather than the undocumented `sp_msforeachtable`)
+// keeps this portable across on-prem SQL Server and Azure SQL Database.
+// `is_ms_shipped = 0` excludes system tables that cannot carry FK constraints.
+// `HAS_PERMS_BY_NAME(..., 'ALTER')` filters out tables the importing login
+// cannot alter (e.g. a multi-schema database where it does not own every
+// table) before building the batch, so one inaccessible table cannot abort
+// a single `ALTER TABLE ... NOCHECK CONSTRAINT ALL` statement covering every
+// table the login *can* alter.
+const SQLSERVER_DISABLE_ALL_FOREIGN_KEYS_SQL: &str = "\
+    DECLARE @dbx_fk_sql nvarchar(max) = N''; \
+    SELECT @dbx_fk_sql = @dbx_fk_sql + N'ALTER TABLE ' + QUOTENAME(s.name) + N'.' + QUOTENAME(t.name) \
+        + N' NOCHECK CONSTRAINT ALL;' \
+    FROM sys.tables t JOIN sys.schemas s ON s.schema_id = t.schema_id \
+    WHERE t.is_ms_shipped = 0 \
+      AND HAS_PERMS_BY_NAME(QUOTENAME(s.name) + N'.' + QUOTENAME(t.name), N'OBJECT', N'ALTER') = 1; \
+    IF @dbx_fk_sql <> N'' EXEC sys.sp_executesql @dbx_fk_sql;";
+
+// Restores enforcement for future writes with `WITH NOCHECK` (skips re-validating
+// rows written while constraints were disabled) so re-enabling never fails or
+// stalls on data an in-progress, possibly partial (continue-on-error) import left
+// behind; that mirrors how `mysqldump`/`pg_restore`-style tools re-enable checks.
+const SQLSERVER_ENABLE_ALL_FOREIGN_KEYS_SQL: &str = "\
+    DECLARE @dbx_fk_sql nvarchar(max) = N''; \
+    SELECT @dbx_fk_sql = @dbx_fk_sql + N'ALTER TABLE ' + QUOTENAME(s.name) + N'.' + QUOTENAME(t.name) \
+        + N' WITH NOCHECK CHECK CONSTRAINT ALL;' \
+    FROM sys.tables t JOIN sys.schemas s ON s.schema_id = t.schema_id \
+    WHERE t.is_ms_shipped = 0 \
+      AND HAS_PERMS_BY_NAME(QUOTENAME(s.name) + N'.' + QUOTENAME(t.name), N'OBJECT', N'ALTER') = 1; \
+    IF @dbx_fk_sql <> N'' EXEC sys.sp_executesql @dbx_fk_sql;";
+
+async fn set_relational_constraints_enabled(
+    state: &AppState,
+    request: &SqlFileRequest,
+    kind: RelationalConstraintBypassKind,
+    token: &CancellationToken,
+    enabled: bool,
+) -> Result<(), String> {
+    let sql = match (kind, enabled) {
+        (RelationalConstraintBypassKind::Postgres, false) => POSTGRES_DISABLE_ALL_RELATIONAL_TRIGGERS_SQL,
+        (RelationalConstraintBypassKind::Postgres, true) => POSTGRES_ENABLE_ALL_RELATIONAL_TRIGGERS_SQL,
+        (RelationalConstraintBypassKind::SqlServer, false) => SQLSERVER_DISABLE_ALL_FOREIGN_KEYS_SQL,
+        (RelationalConstraintBypassKind::SqlServer, true) => SQLSERVER_ENABLE_ALL_FOREIGN_KEYS_SQL,
+        (RelationalConstraintBypassKind::Mysql, _) => {
+            unreachable!("MySQL constraint bypass is handled by MySqlSqlFileExecutor, not this helper")
+        }
+    };
+    execute_sql_statement_with_options(
+        state,
+        &request.connection_id,
+        &request.database,
+        sql,
+        None,
+        Some(token.clone()),
+        QueryExecutionOptions::default(),
+    )
+    .await
+    .map(|_| ())
+}
+
 pub async fn execute_sql_file_content(
     state: &AppState,
     request: &SqlFileRequest,
@@ -471,8 +617,19 @@ pub async fn execute_sql_file_content(
     // MySQL-family imports need one pinned connection so `USE` and session
     // state survive across the whole file.
     let mut mysql_executor = MySqlSqlFileExecutor::build(state, request, import_target.as_ref()).await?;
+    let bypass_kind = import_target.as_ref().and_then(relational_constraint_bypass_kind);
+    let non_mysql_bypass = request.skip_relational_constraints
+        && mysql_executor.is_none()
+        && matches!(
+            bypass_kind,
+            Some(RelationalConstraintBypassKind::Postgres | RelationalConstraintBypassKind::SqlServer)
+        );
+    if non_mysql_bypass {
+        let kind = bypass_kind.expect("non_mysql_bypass implies bypass_kind is set");
+        set_relational_constraints_enabled(state, request, kind, &token, false).await?;
+    }
     let mut progress = SqlFileExecutionProgress::new();
-    execute_planned_statements_with_progress(
+    let import_result = execute_planned_statements_with_progress(
         state,
         request,
         &token,
@@ -482,7 +639,24 @@ pub async fn execute_sql_file_content(
         &mut progress,
         &mut emit,
     )
-    .await?;
+    .await;
+    if non_mysql_bypass {
+        let kind = bypass_kind.expect("non_mysql_bypass implies bypass_kind is set");
+        // Use a fresh token, not the (possibly already-cancelled) import token: the
+        // restore is catalog-level and must run even after the user cancels, or FK
+        // enforcement stays disabled database-wide for every session. The already
+        // cancelled `token` would fail `execute_sql_statement_with_options`'s
+        // pre-dispatch cancellation check before the restore SQL is ever sent.
+        if let Err(error) =
+            set_relational_constraints_enabled(state, request, kind, &CancellationToken::new(), true).await
+        {
+            log::error!(
+                "[sql_file_import] failed to restore relational constraints for connection {}: {error}",
+                request.connection_id
+            );
+        }
+    }
+    import_result?;
     emit_sql_file_terminal_progress(request, &token, started_at, &progress, &mut emit);
     Ok(())
 }
@@ -571,6 +745,20 @@ pub async fn execute_sql_file_paths(
         }
         executor.constraints_disabled = true;
     }
+    let bypass_kind = import_target.as_ref().and_then(relational_constraint_bypass_kind);
+    let non_mysql_bypass = request.skip_relational_constraints
+        && mysql_executor.is_none()
+        && matches!(
+            bypass_kind,
+            Some(RelationalConstraintBypassKind::Postgres | RelationalConstraintBypassKind::SqlServer)
+        );
+    if non_mysql_bypass {
+        let kind = bypass_kind.expect("non_mysql_bypass implies bypass_kind is set");
+        if let Err(error) = set_relational_constraints_enabled(state, request, kind, &token, false).await {
+            emit(sql_file_execution_error_progress(&request.execution_id, started_at, &progress, error.clone()));
+            return Err(error);
+        }
+    }
     let file_count = file_paths.len();
     let mut prev_statement_index = 0usize;
     let mut prev_success_count = 0usize;
@@ -768,6 +956,20 @@ pub async fn execute_sql_file_paths(
             let _ = executor.set_foreign_key_checks(state, &token, true).await;
         }
     }
+    if non_mysql_bypass {
+        let kind = bypass_kind.expect("non_mysql_bypass implies bypass_kind is set");
+        // Fresh token: see the matching comment in `execute_sql_file_content`. `token`
+        // here may already be cancelled (user cancel or terminal error), which would
+        // make `execute_sql_statement_with_options` refuse to send the restore SQL.
+        if let Err(error) =
+            set_relational_constraints_enabled(state, request, kind, &CancellationToken::new(), true).await
+        {
+            log::error!(
+                "[sql_file_import] failed to restore relational constraints for connection {}: {error}",
+                request.connection_id
+            );
+        }
+    }
     import_result
 }
 
@@ -3140,4 +3342,88 @@ mod tests {
         );
         assert_eq!(emitted.last().unwrap().status, SqlFileStatus::Done);
     }
+
+    // No live PostgreSQL/SQL Server instance is available in this unit-test
+    // module (see the `#[ignore]`d live_* integration tests for those), so
+    // these regression tests assert on the generated SQL/control-flow source
+    // directly for the two properties raised in review:
+    //   1. the restore (re-enable) call must not be gated on the same
+    //      `CancellationToken` used for the import, or a user cancellation
+    //      permanently disables constraints database-wide (see
+    //      `execute_sql_file_content` / `execute_sql_file_paths`).
+    //   2. the PostgreSQL disable statement must fail loudly instead of
+    //      silently doing nothing when every table hits `insufficient_privilege`
+    //      (non-superuser importing roles cannot disable internal FK/RI
+    //      triggers on tables they do not own).
+
+    #[test]
+    fn relational_constraint_restore_never_reuses_the_import_cancellation_token() {
+        let source = include_str!("sql_file_import.rs");
+        for call_site in
+            source.match_indices("set_relational_constraints_enabled(state, request, kind, ").map(|(i, _)| i)
+        {
+            let after_call = &source[call_site..];
+            let args_end = after_call.find(')').unwrap_or(after_call.len());
+            let call = &after_call[..args_end];
+            let is_restore_call = call.trim_end().ends_with("true");
+            if is_restore_call {
+                assert!(
+                    call.contains("CancellationToken::new()"),
+                    "restore (enable=true) call must use a fresh CancellationToken, not the import's \
+                     (possibly already-cancelled) token, or constraints stay disabled database-wide \
+                     after a cancel: {call}"
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn postgres_constraint_disable_raises_when_every_table_lacks_privilege() {
+        assert!(
+            POSTGRES_DISABLE_ALL_RELATIONAL_TRIGGERS_SQL.contains("dbx_total > 0 AND dbx_disabled = 0"),
+            "the disable statement must detect and raise when it altered zero of the tables it found, \
+             instead of silently continuing as if the bypass succeeded"
+        );
+        assert!(
+            POSTGRES_DISABLE_ALL_RELATIONAL_TRIGGERS_SQL.contains("RAISE EXCEPTION"),
+            "must actually raise, not just count, so the caller's `execute_sql_statement_with_options` call \
+             surfaces an error instead of reporting success"
+        );
+    }
+
+    #[test]
+    fn postgres_and_sqlserver_disable_statements_exclude_system_schemas() {
+        for schema in ["pg_catalog", "information_schema", "pg_toast"] {
+            assert!(
+                POSTGRES_DISABLE_ALL_RELATIONAL_TRIGGERS_SQL.contains(schema),
+                "must exclude system schema {schema} or the bypass tries to alter catalog tables"
+            );
+        }
+        assert!(SQLSERVER_DISABLE_ALL_FOREIGN_KEYS_SQL.contains("is_ms_shipped = 0"));
+    }
+
+    #[test]
+    fn relational_constraint_bypass_kind_covers_postgres_family_and_sqlserver() {
+        let target = |db_type| SqlFileImportTarget {
+            db_type,
+            driver_profile: None,
+            compatibility_mode: SqlCompatibilityMode::NotApplicable,
+        };
+        assert_eq!(
+            relational_constraint_bypass_kind(&target(DatabaseType::Mysql)),
+            Some(RelationalConstraintBypassKind::Mysql)
+        );
+        for db_type in [DatabaseType::Postgres, DatabaseType::Gaussdb, DatabaseType::OpenGauss] {
+            assert_eq!(
+                relational_constraint_bypass_kind(&target(db_type)),
+                Some(RelationalConstraintBypassKind::Postgres),
+                "{db_type:?} should use the PostgreSQL trigger-disable bypass"
+            );
+        }
+        assert_eq!(
+            relational_constraint_bypass_kind(&target(DatabaseType::SqlServer)),
+            Some(RelationalConstraintBypassKind::SqlServer)
+        );
+        assert_eq!(relational_constraint_bypass_kind(&target(DatabaseType::Oracle)), None);
+    }
 }
diff --git a/crates/dbx-sql/src/sql.rs b/crates/dbx-sql/src/sql.rs
index 1654714b73..7bee198aa0 100644
--- a/crates/dbx-sql/src/sql.rs
+++ b/crates/dbx-sql/src/sql.rs
@@ -63,9 +63,16 @@ pub struct SqlFileRequest {
     pub selected_tables: Option<Vec<dbx_types::sql_file::SqlFileTable>>,
     #[serde(default)]
     pub part_cooldown_ms: u64,
-    /// Temporarily disable MySQL `FOREIGN_KEY_CHECKS` for this import and
-    /// restore them on completion, error, or cancellation. Only applies to
-    /// MySQL-compatible connections that reuse one pinned session.
+    /// Temporarily disable relational constraint enforcement across the entire
+    /// target database for this import, and restore it on completion, error,
+    /// or cancellation. The mechanism is engine-specific: MySQL-compatible
+    /// connections toggle session-scoped `FOREIGN_KEY_CHECKS`; PostgreSQL-family
+    /// connections (Postgres, GaussDB, openGauss) run `ALTER TABLE ... DISABLE
+    /// TRIGGER ALL` / `ENABLE TRIGGER ALL` for every table; SQL Server runs
+    /// `ALTER TABLE ... NOCHECK CONSTRAINT ALL` / `WITH NOCHECK CHECK CONSTRAINT
+    /// ALL` for every table. Has no effect on other connection types. This does
+    /// not bypass unique/primary-key violations on any engine — only foreign-key
+    /// and (on PostgreSQL) other trigger-enforced constraints.
     #[serde(default)]
     pub skip_relational_constraints: bool,
 }