SQLAlchemy运行单元格两次触发复合键重复约束错误求解

问题原因 #

你的代码存在两个核心问题导致第二次运行触发唯一键冲突：

1. 表名不一致：to_sql使用table_name.lower()指定表名，但后续COPY语句用的是原传入的table_name。如果传入的表名包含大小写差异（比如final.Product），PostgreSQL会将其视为不同表，导致replace清空的是一个空表，而COPY写入的是另一个已有数据的表。
2. 约束丢失风险：即使表名一致，df.head(0).to_sql(..., if_exists='replace')会删除原表并重建新表，但新表不会保留原有的复合主键等约束——你第一次运行后表仍有主键，可能是手动创建过约束，但第二次replace后约束会丢失，后续若重新添加约束再导入数据，就会因重复主键报错。

可行解决方案 #

方案1：统一表名，全量替换表数据 #

确保to_sql和COPY操作同一个表，用replace重建表后导入数据：

from sqlalchemy import create_engine
import io

def df_to_postgres(dataframe, table_name):
    engine = create_engine('postgresql://postgres:pass@localhost:5432/postgres')
    df = dataframe
    
    # 统一使用传入的表名，避免大小写差异
    df.head(0).to_sql(table_name, engine, if_exists='replace', index=False)
    
    # 若需要重建复合主键约束，添加以下语句
    with engine.connect() as conn:
        conn.execute(f'ALTER TABLE {table_name} ADD CONSTRAINT product_pkey PRIMARY KEY (sku, unit)')
        conn.commit()
    
    conn = engine.raw_connection()
    cur = conn.cursor()
    output = io.StringIO()
    df.to_csv(output, sep='\t', header=False, index=False)
    output.seek(0)
    cur.copy_expert(f'COPY {table_name} FROM STDIN', output) 
    conn.commit()

方案2：保留原表约束，清空数据后导入 #

如果不想丢失表的原有约束（如主键、索引），可以先清空表数据再导入，避免重建表：

from sqlalchemy import create_engine, inspect
import io

def df_to_postgres(dataframe, table_name):
    engine = create_engine('postgresql://postgres:pass@localhost:5432/postgres')
    df = dataframe
    inspector = inspect(engine)
    
    # 检查表是否存在，不存在则创建并添加约束
    if not inspector.has_table(table_name):
        df.head(0).to_sql(table_name, engine, index=False)
        with engine.connect() as conn:
            conn.execute(f'ALTER TABLE {table_name} ADD CONSTRAINT product_pkey PRIMARY KEY (sku, unit)')
            conn.commit()
    else:
        # 表存在则清空数据
        with engine.connect() as conn:
            conn.execute(f'TRUNCATE TABLE {table_name}')
            conn.commit()
    
    # COPY导入数据
    conn = engine.raw_connection()
    cur = conn.cursor()
    output = io.StringIO()
    df.to_csv(output, sep='\t', header=False, index=False)
    output.seek(0)
    cur.copy_expert(f'COPY {table_name} FROM STDIN', output)
    conn.commit()

方案3：增量更新，处理重复主键 #

如果需要保留原有数据，仅更新重复主键的记录，可使用临时表+ON CONFLICT语法：

from sqlalchemy import create_engine
import io

def df_to_postgres(dataframe, table_name):
    engine = create_engine('postgresql://postgres:pass@localhost:5432/postgres')
    df = dataframe
    temp_table = f"{table_name}_temp"
    
    # 创建临时表并导入数据
    df.head(0).to_sql(temp_table, engine, if_exists='replace', index=False)
    conn = engine.raw_connection()
    cur = conn.cursor()
    output = io.StringIO()
    df.to_csv(output, sep='\t', header=False, index=False)
    output.seek(0)
    cur.copy_expert(f'COPY {temp_table} FROM STDIN', output)
    
    # 合并到目标表：重复主键则更新指定字段，否则插入
    cur.execute(f'''
        INSERT INTO {table_name}
        SELECT * FROM {temp_table}
        ON CONFLICT (sku, unit) DO UPDATE SET
            -- 替换为你需要更新的字段，示例：
            product_name = EXCLUDED.product_name,
            price = EXCLUDED.price
    ''')
    conn.commit()
    
    # 删除临时表
    cur.execute(f'DROP TABLE {temp_table}')
    conn.commit()
    conn.close()

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