使用psycopg3批量插入含空字符串复合主键列遇NULL转换问题

问题描述 #

我有一张PostgreSQL表test_ts，包含5列复合主键，其中context列是主键且默认值为空字符串。使用psycopg3批量插入DataFrame数据时，DataFrame中显式设置为空字符串的context值被转换为NULL，触发非空约束错误，需要解决这个问题以实现空字符串的正确插入。

表结构 #

test_ts = Table('test_ts', meta,
    Column('metric_id', ForeignKey('metric.id'), primary_key = True), 
    Column('entity_id', Integer, primary_key = True), 
    Column('date', DateTime, primary_key = True),              
    Column('freq', String, primary_key = True),           
    Column('context', String, primary_key = True, server_default=""),
    Column('value', String),
    Column("update_time", DateTime, server_default=func.now(), onupdate=func.current_timestamp()),
    Column('update_by', String, server_default=func.current_user()))

数据样例 #

date context  value freq  entity_id  metric_id
1  1999-02-01T05:00:00.000Z    test    101    D       1105          4
2  1999-02-02T05:00:00.000Z    test    102    D       1105          4
8  1999-02-01T05:00:00.000Z            201    D       1105          4
9  1999-02-02T05:00:00.000Z            202    D       1105          4

错误信息 #

Traceback (most recent call last):
  File "/workspaces/service/data_svc.py", line 121, in bulk_insert
    cur.execute(sql_insert)
  File "/usr/local/lib/python3.11/site-packages/psycopg/cursor.py", line 732, in execute
    raise ex.with_traceback(None)
psycopg.errors.NotNullViolation: null value in column "context" of relation "test_ts" violates not-null constraint
DETAIL:  Failing row contains (4, 1105, 1999-02-01 05:00:00, D, null, 201, 2024-03-27 21:23:42.778517, db_user).        

使用的SQL语句 #

创建临时表 #

drop table if exists tmp_tbl; CREATE UNLOGGED TABLE tmp_tbl AS SELECT date,context,value,freq,entity_id,metric_id FROM test_ts LIMIT 0

COPY数据到临时表 #

COPY tmp_tbl (date,context,value,freq,entity_id,metric_id) FROM STDIN (FORMAT CSV,  DELIMITER "	")

插入数据到目标表 #

insert into test_ts(date,context,value,freq,entity_id,metric_id) 
select * from tmp_tbl on conflict(metric_id,entity_id,date,freq,context) 
do update set value = EXCLUDED.value;drop table if exists tmp_tbl;

代码片段 #

with psycopg.connect(self.connect_str, autocommit=True) as conn:
    io_buf = io.StringIO()
    df.to_csv(io_buf, sep='\t', header=False, index=False)
    io_buf.seek(0)
    
    with conn.cursor() as cur:
          cur.execute(sql_create)
          with cur.copy(sql_copy) as copy:
              while data:=io_buf.read(self.batch_size):
                  copy.write(data)

          cur.execute(sql_insert)
conn.commit()

解决方案 #

方法1：修改INSERT语句，将NULL转为空字符串 #

在从临时表向目标表插入数据时，使用COALESCE函数将context字段的NULL值替换为空字符串，确保符合非空约束：

insert into test_ts(date,context,value,freq,entity_id,metric_id) 
select date, COALESCE(context, ''), value, freq, entity_id, metric_id from tmp_tbl 
on conflict(metric_id,entity_id,date,freq,context) 
do update set value = EXCLUDED.value;drop table if exists tmp_tbl;

方法2：修改COPY命令，强制空字段不转为NULL #

在COPY语句中添加FORCE_NOT_NULL(context)参数，告诉PostgreSQL不要将context字段的空值解析为NULL，而是保留为空字符串：

COPY tmp_tbl (date,context,value,freq,entity_id,metric_id) FROM STDIN (FORMAT CSV, DELIMITER '	', FORCE_NOT_NULL(context))

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