Databricks写入DataFrame到Azure SQL时如何禁用启用主键外键约束

Azure SQL基于SQL Server引擎实现，可通过JDBC连接在写入前后执行原生T-SQL命令完成约束的临时禁用与恢复，全程不需要手动修改表结构定义，具体实现步骤如下：

前置配置 #

先提前定义好Azure SQL的JDBC连接参数，避免硬编码敏感信息：

# Azure SQL JDBC连接配置
jdbc_url = "jdbc:sqlserver://<你的Azure SQL服务名>.database.windows.net:1433;database=<目标库名>;encrypt=true;trustServerCertificate=false;hostNameInCertificate=*.database.windows.net;loginTimeout=30;"
connection_props = {
  "user": "<数据库访问账号>",
  "password": "<数据库访问密码>",
  "driver": "com.microsoft.sqlserver.jdbc.SQLServerDriver"
}
target_schema = "dbo" # 替换为目标表所属schema
target_table = "your_target_table" # 替换为目标表名

写入前禁用关联约束 #

通过系统视图自动查询目标表关联的所有外键（含当前表的外键、其他表引用当前表主键的外键），可按需选择是否禁用主键/唯一键约束，执行禁用后留存后续启用需要的命令列表：

def disable_related_constraints(spark, jdbc_url, conn_props, schema, table, disable_pk=False):
    driver_manager = spark._jvm.java.sql.DriverManager
    conn = driver_manager.getConnection(jdbc_url, conn_props["user"], conn_props["password"])
    statement = conn.createStatement()
    try:
        # 基础查询：所有关联外键
        base_query = """
        SELECT DISTINCT 
            'ALTER TABLE [' + s.name + '].[' + t.name + '] NOCHECK CONSTRAINT [' + fk.name + ']' as disable_cmd,
            'ALTER TABLE [' + s.name + '].[' + t.name + '] WITH CHECK CHECK CONSTRAINT [' + fk.name + ']' as enable_cmd
        FROM sys.foreign_keys fk
        JOIN sys.tables t ON fk.parent_object_id = t.object_id
        JOIN sys.schemas s ON t.schema_id = s.schema_id
        WHERE fk.parent_object_id = OBJECT_ID('[{schema}].[{table}]')
           OR fk.referenced_object_id = OBJECT_ID('[{schema}].[{table}]')
        """
        # 按需拼接主键/唯一键查询逻辑
        if disable_pk:
            base_query += """
            UNION ALL
            SELECT 
                'ALTER TABLE [{schema}].[{table}] NOCHECK CONSTRAINT [' + kc.name + ']' as disable_cmd,
                'ALTER TABLE [{schema}].[{table}] WITH CHECK CHECK CONSTRAINT [' + kc.name + ']' as enable_cmd
            FROM sys.key_constraints kc
            WHERE kc.parent_object_id = OBJECT_ID('[{schema}].[{table}]')
              AND kc.type IN ('PK', 'UQ')
            """
        # 格式化查询语句
        final_query = base_query.format(schema=schema, table=table)
        rs = statement.executeQuery(final_query)

        disable_cmds = []
        enable_cmds = []
        while rs.next():
            disable_cmds.append(rs.getString("disable_cmd"))
            enable_cmds.append(rs.getString("enable_cmd"))
        
        # 执行所有禁用命令
        for cmd in disable_cmds:
            statement.execute(cmd)
        conn.commit()
        return enable_cmds
    except Exception as e:
        conn.rollback()
        raise e
    finally:
        statement.close()
        conn.close()

# 执行约束禁用，默认不禁用主键/唯一键，需要的话把disable_pk设为True
enable_cmds = disable_related_constraints(spark, jdbc_url, connection_props, target_schema, target_table)

执行DataFrame写入 #

写入时注意如果使用overwrite模式必须开启truncate选项，否则Spark会默认删除表后重建，导致原有约束全部丢失：

# 替换为你自己的PySpark DataFrame变量名
source_df.write \
    .format("jdbc") \
    .option("url", jdbc_url) \
    .option("dbtable", f"[{target_schema}].[{target_table}]") \
    .option("user", connection_props["user"]) \
    .option("password", connection_props["password"]) \
    .option("driver", connection_props["driver"]) \
    .option("truncate", "true") \
    .mode("overwrite") \
    .save()

如果使用append写入模式，可以移除truncate配置

写入完成后重新启用约束 #

启用时必须添加WITH CHECK参数，否则约束会被标记为不可信，既不会校验存量数据，也不会被查询优化器使用：

def enable_all_constraints(spark, jdbc_url, conn_props, enable_cmd_list):
    driver_manager = spark._jvm.java.sql.DriverManager
    conn = driver_manager.getConnection(jdbc_url, conn_props["user"], conn_props["password"])
    statement = conn.createStatement()
    try:
        # 倒序执行启用命令，先启用主键/唯一键再启用外键，避免依赖报错
        for cmd in reversed(enable_cmd_list):
            statement.execute(cmd)
        conn.commit()
    except Exception as e:
        conn.rollback()
        raise RuntimeError("约束启用失败，写入数据存在违反主键/外键规则的脏数据，请排查数据一致性") from e
    finally:
        statement.close()
        conn.close()

# 执行约束启用
enable_all_constraints(spark, jdbc_url, connection_props, enable_cmds)

关键注意事项 #

• 多层外键依赖场景（比如A表引用B表，B表引用目标表），需要把约束查询逻辑改成递归CTE遍历所有层级依赖，上述脚本默认仅处理直接关联的一层约束
• 非必要不要禁用主键、唯一键约束，这类约束绑定索引，禁用后重新启用需要全表校验数据，大表场景耗时极高
• 禁用约束前确认没有其他业务作业写入关联表，否则会出现绕过约束的脏数据
• 约束启用阶段的耗时和写入数据量正相关，大表写入场景需要预留足够的作业超时时间

内容的提问来源于stack exchange，提问作者Arjun R