Spark写入带标识列的SQL Server表如何开启IDENTITY_INSERT

你的需求可以通过原生Spark JDBC连接器实现，不需要更换其他组件，核心是要让Spark写入JDBC的会话提前执行SET IDENTITY_INSERT语句，因为该配置是会话级别生效，你之前的尝试要么参数冲突、要么会话不共享，所以不生效。

方案1：使用connectionInitSql参数（推荐，无需修改写入逻辑） #

Spark JDBC提供了connectionInitSql参数，会在每个JDBC连接创建完成后自动执行指定的SQL语句，刚好可以用来开启目标表的IDENTITY_INSERT配置，修改后的写入代码如下：

df.write.format("jdbc") \
      .mode(mode) \
      .option("url", jdbc_url) \
      .option("dbtable", table_name) \
      .option("user", jdbc_username) \
      .option("password", jdbc_password) \
      .option("driver", "com.microsoft.sqlserver.jdbc.SQLServerDriver") \
      # 新增以下参数，每个连接初始化时开启目标表的IDENTITY_INSERT
      .option("connectionInitSql", f"SET IDENTITY_INSERT {table_name} ON;") \
      # 如果你用overwrite模式，必须加下面这个参数避免删表重建丢失identity属性
      .option("truncate", "true") \
      .save()

注意事项 #

• IDENTITY_INSERT 同一时间只能在一个会话中对一张表生效，参数内只能指定你当前写入的目标表
• 如果表名包含特殊字符、关键字，需要用[]包裹，比如SET IDENTITY_INSERT [dbo].[UserTable] ON;

方案2：自定义foreachPartition写入（兼容低版本Spark） #

如果你使用的Spark版本低于2.3，不支持connectionInitSql参数，可以手动控制每个分区的JDBC连接，写入前先执行开启语句：

def write_to_sqlserver(partition):
    import java.sql.DriverManager
    conn = DriverManager.getConnection(jdbc_url, jdbc_username, jdbc_password)
    # 先开启当前会话的IDENTITY_INSERT
    cursor = conn.createStatement()
    cursor.execute(f"SET IDENTITY_INSERT {table_name} ON;")
    conn.setAutoCommit(False)
    # 批量写入数据，根据你的表字段调整INSERT语句和参数赋值逻辑
    insert_sql = f"INSERT INTO {table_name} (id, col1, col2) VALUES (?, ?, ?)"
    ps = conn.prepareStatement(insert_sql)
    for row in partition:
        ps.setInt(1, row.id)
        ps.setString(2, row.col1)
        ps.setInt(3, row.col2)
        ps.addBatch()
    ps.executeBatch()
    conn.commit()
    ps.close()
    conn.close()

df.foreachPartition(write_to_sqlserver)

通用注意事项 #

• 写入前需要确保DataFrame中的Id列值和目标表现有Id没有重复，否则会触发主键冲突报错
• 不需要手动执行SET IDENTITY_INSERT OFF，写入完成后JDBC会话销毁，配置自动失效，不会影响其他业务的写入操作

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