DML RETURNING omits other mapped cols due to bulk insert assumptions

[zzzeek]

Discussed in #12327

from __future__ import annotations

from sqlalchemy import create_engine
from sqlalchemy import ForeignKey
from sqlalchemy import update
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column
from sqlalchemy.orm import relationship
from sqlalchemy.orm import Session


class Base(DeclarativeBase):
    pass


class A(Base):
    __tablename__ = "a"

    id: Mapped[int] = mapped_column(primary_key=True)
    data: Mapped[str]
    bs: Mapped[list[B]] = relationship("B")


class B(Base):
    __tablename__ = "b"
    id: Mapped[int] = mapped_column(primary_key=True)
    a_id: Mapped[int] = mapped_column(ForeignKey("a.id"))
    data: Mapped[str]


e = create_engine("postgresql://scott:tiger@localhost/test", echo=True)
Base.metadata.create_all(e)

s = Session(e)
s.add(
    A(data='a1', bs=[B(data='b2')])
)
s.flush()

result = s.execute(
    update(A).values(data='foo').where(A.id == B.a_id).returning(A.data, B.a_id, B.data)
)
print(result.all())

renders:

UPDATE a SET data=%(data)s FROM b WHERE a.id = b.a_id RETURNING a.id, a.data

and fails

sqlalchemy.exc.NoSuchColumnError: Could not locate column in row for column 'b.a_id'

[sqla-tester]

Mike Bayer has proposed a fix for this issue in the main branch:

only use _DMLReturningColFilter for "bulk insert", not other DML https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/5689

[sqla-tester]

Mike Bayer has proposed a fix for this issue in the rel_2_0 branch:

only use _DMLReturningColFilter for "bulk insert", not other DML https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/5690