如何用SQLAlchemy实现基于CTE的批量自增ID插入？

原始需求对应的PostgreSQL SQL #

我们需要实现通过CTE获取表的最大ID，然后批量插入N条数据，每条数据的idmilityattache字段为最大ID依次加1至N，原始SQL如下：

with cte as (select max(idmilityattache)
             from dopinfo.milityattaches)
insert
into dopinfo.milityattaches
values ((select * from cte) + 1, 1, 'a', '2023-03-02'::date, 'a'),
       ((select * from cte) + 2, 1, 'a', '2023-03-02'::date, 'a')

失败的尝试代码 #

最初尝试的代码存在问题：循环中每次单独查询max ID，不仅效率低下，还可能因并发导致ID冲突，且未复用CTE的原子性逻辑：

attaches_values = [
            {
                Milityattache.nameemployee: literal(self.validated['fio']),
                Milityattache.idmilitaryrank: literal(self.validated['mr']),
                Milityattache.datestart: literal(self.validated['datestart']),
                Milityattache.telefons: literal(self.validated['telefons']),
                Milityattache.idmilityattache: self.session.query(
                    case(
                        [(func.max(Milityattache.idmilityattache).isnot(None),
                          func.max(Milityattache.idmilityattache) + num)], else_=literal(value=num)
                        )
                    ).scalar()
            }
            for num in range(1, len(self.validated['country'])+1)
        ]
        res = self.session.execute(
            insert(
                Milityattache

            ).values(
                attaches_values
            ).returning(
                Milityattache.idmilityattache
            )
        )

可行的解决方案 #

正确的实现方式是先构建与原始SQL对应的CTE，复用该CTE生成每条插入数据的自增ID，确保所有插入操作基于同一个max ID值，避免并发问题和重复查询：

max_id_cte = self.session.query(
        case(
            [(func.max(Milityattache.idmilityattache).isnot(None),
              func.max(Milityattache.idmilityattache) + 1)], else_=literal(value=1)
        ).label("m")
    ).cte("cte")

    ins = insert(
        Milityattache,
        [
            {
                'nameemployee': self.validated['fio'],
                'idmilitaryrank': self.validated['mr'],
                'datestart': self.validated['datestart'],
                'telefons': self.validated['telefons'],
                'idmilityattache': select([max_id_cte.c.m + _id])
            }
            for _id in range(len(self.validated['country']))
        ]
    ).returning(Milityattache.idmilityattache)

方案说明 #

• 先通过CTE查询得到基础起始ID（表为空时从1开始）
• 批量插入时，通过max_id_cte.c.m + _id生成连续自增ID，确保所有ID基于同一个起始值计算
• 完全对应原始PostgreSQL SQL的逻辑，保证操作的原子性与一致性

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