使用SQLAlchemy实现PostgreSQL批量插入并处理冲突更新

需求回顾 #

你要处理的场景是：

• 每日向PostgreSQL表插入约30000条数据，表结构包含id（主键）、category、createddate、updatedon四列
• 核心逻辑：
  • 若id已存在：更新updatedon为当日日期，同时把category替换成新传入的类别
  • 若id不存在：插入新行，且createddate和updatedon都设为当日日期

代码实现步骤 #

首先假设你已经用SQLAlchemy定义了对应的表模型（如果还没定义，先看这一步）：

from sqlalchemy import Column, Integer, String, Date
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()

# 替换成你的实际表名和字段类型
class DataTable(Base):
    __tablename__ = 'your_target_table'
    id = Column(Integer, primary_key=True)
    category = Column(String(50))  # 根据实际需求调整长度
    createddate = Column(Date)
    updatedon = Column(Date)

接下来是核心的Upsert逻辑代码，专门针对你的需求定制：

from sqlalchemy.dialects.postgresql import insert
from datetime import date
from sqlalchemy.orm import sessionmaker
# 这里假设你已经创建了数据库引擎engine
Session = sessionmaker(bind=engine)
session = Session()

# 模拟你的3万条待插入数据，实际场景中替换成你的数据源
data_batch = [
    {"id": 1001, "category": "tech"},
    {"id": 1002, "category": "life"},
    # ... 更多数据
]

# 获取当日日期，统一用这个值处理createddate和updatedon
today = date.today()

# 构建批量插入语句
insert_statement = insert(DataTable).values(
    [
        {
            "id": item["id"],
            "category": item["category"],
            "createddate": today,
            "updatedon": today
        } for item in data_batch
    ]
)

# 定义冲突后的更新规则：当id主键冲突时，更新category和updatedon
upsert_statement = insert_statement.on_conflict_do_update(
    index_elements=['id'],  # 指定判断冲突的主键字段
    set_={
        # 用插入语句中携带的新category覆盖旧值
        'category': insert_statement.excluded.category,
        # 把updatedon更新为当日日期
        'updatedon': today
    }
)

# 执行Upsert操作
try:
    with session.begin():
        session.execute(upsert_statement)
finally:
    session.close()

性能优化建议 #

因为你要处理3万条数据，一次性提交可能会导致事务过大，建议分批次处理（比如每1000条一批），这样能降低数据库压力，提升稳定性：

batch_size = 1000
total_data = len(data_batch)

try:
    with session.begin():
        for start_idx in range(0, total_data, batch_size):
            end_idx = min(start_idx + batch_size, total_data)
            current_batch = data_batch[start_idx:end_idx]
            
            insert_stmt = insert(DataTable).values(
                [{"id": item["id"], "category": item["category"], "createddate": today, "updatedon": today} 
                 for item in current_batch]
            )
            upsert_stmt = insert_stmt.on_conflict_do_update(
                index_elements=['id'],
                set_={'category': insert_stmt.excluded.category, 'updatedon': today}
            )
            session.execute(upsert_stmt)
finally:
    session.close()

关键逻辑说明 #

• insert(...)：构建基础的批量插入语句，给新数据的createddate和updatedon都设为当日日期
• on_conflict_do_update：这是PostgreSQL特有的Upsert语法，只有当指定的id主键发生冲突时才会触发更新
• insert_statement.excluded：代表插入语句中携带的待插入数据，用它来获取新的category值，保证更新的是最新传入的类别

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