如何手动创建持久化状态的ORM对象？（Union All改写场景）

原代码通过两次独立查询加载指定列的ORM对象：

def my_func(
    session: Session, country_code: str, user_email: str
) -> tuple[Country | None, User | None]:
    country = session.scalar(
        select(Country)
        .where(Country.code == country_code)
        .options(load_only(Country.id))
    )

    user = session.scalar(
        select(User)
        .where(User.email == user_email)
        .options(load_only(User.id, User.name))
    )

    return country, user

改用UNION ALL合并查询后，要保持返回对象的结构和延迟加载特性不变，实现代码如下：

from sqlalchemy import union_all, select
from sqlalchemy.orm import Session, InstanceState

def my_func(
    session: Session, country_code: str, user_email: str
) -> tuple[Country | None, User | None]:
    # 构造UNION ALL查询，用第一个字段标记数据类型，统一列名保证结构一致
    country_query = (
        select(1, Country.id.label("id"), None.label("name"))
        .where(Country.code == country_code)
    )
    user_query = (
        select(2, User.id.label("id"), User.name.label("name"))
        .where(User.email == user_email)
    )
    # 执行查询获取所有结果行
    result_rows = session.execute(union_all(country_query, user_query)).all()

    country = None
    user = None

    for row in result_rows:
        type_flag, id_val, name_val = row
        if type_flag == 1 and id_val:
            # 创建Country实例并设置已加载的id字段
            country = Country(id=id_val)
            # 将实例加入session，转为持久化状态
            session.add(country)
            # 获取实例状态，手动标记字段加载状态
            state = InstanceState(country)
            state.attrs['id'].loaded = True
            # 其余列设置为延迟加载
            for col_key in Country.__mapper__.columns.keys():
                if col_key != 'id':
                    state.attrs[col_key].loaded = False
                    state.attrs[col_key].deferred = True

        elif type_flag == 2 and id_val:
            # 创建User实例并设置已加载的id、name字段
            user = User(id=id_val, name=name_val)
            session.add(user)
            state = InstanceState(user)
            # 标记已加载的字段
            state.attrs['id'].loaded = True
            state.attrs['name'].loaded = True
            # 其余列设置为延迟加载
            for col_key in User.__mapper__.columns.keys():
                if col_key not in ['id', 'name']:
                    state.attrs[col_key].loaded = False
                    state.attrs[col_key].deferred = True

    return country, user

核心要点 #

• 用type_flag区分Country和User数据，查询时统一列名避免结构冲突
• 手动构造ORM实例后，通过session.add()将其纳入session的持久化上下文
• 借助InstanceState精确控制每个字段的加载状态：已获取的字段标记为已加载，其余字段设为延迟加载，和原load_only的行为完全一致
• 返回的对象和原接口完全兼容，调用方无需修改任何代码

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