SQLAlchemy 高级技巧：通过虚拟连接创建带延迟加载列的持久化ORM对象

本文深入探讨了在sqlalchemy中，如何高效地从现有数据创建持久化orm对象，并确保未加载的列实现延迟加载。针对同时获取多种类型实体并保持其orm特性的需求，文章介绍了一种利用postgresql的`unnest()`函数和`left join`构建虚拟查询的先进方法。该方案通过单次数据库查询，巧妙地生成带延迟加载列的orm实例，显著提升了数据检索效率和orm对象管理的灵活性。

背景与挑战

在SQLAlchemy应用中，我们经常需要根据特定条件从数据库中检索数据，并将其映射为ORM对象。通常情况下，这通过直接查询ORM模型来完成。然而，在某些场景下，我们可能已经通过非ORM方式（例如，使用UNION ALL或复杂的原生SQL查询）获取了部分数据，并希望基于这些数据手动创建持久化的ORM对象，同时确保那些未在初始查询中加载的列能够被延迟加载（lazy loading）。

传统的UNION ALL方法虽然能高效地在一个查询中获取多类型数据，但其结果通常是原始的元组（tuples），而非ORM对象。将这些元组转换为具有延迟加载特性的持久化ORM对象，并且不改变现有调用代码对ORM对象的预期，是一个挑战。直接实例化User(id=..., name=...)只会创建一个游离（detached）状态的对象，它不与会话关联，也无法自动触发其他属性的延迟加载。session.merge()可以帮助将游离对象重新附加到会话并使其持久化，但它通常需要一个完整的对象或主键，并且可能仍会触发额外的数据库查询来填充缺失的属性。

解决方案：基于虚拟连接的单次查询策略

为了解决上述挑战，我们可以采用一种高级的查询策略，它利用数据库的特定函数（如PostgreSQL的unnest()）结合LEFT JOIN来构建一个“虚拟根表”，从而在一个查询中同时获取不同类型的ORM对象，并确保它们是持久化且支持延迟加载的。

该策略的核心思想是：

1. 创建虚拟根表： 使用一个序列或一组常量值作为查询的起点，形成一个“虚拟”的行集合。
2. 条件性外部连接： 将需要加载的ORM模型通过LEFT JOIN（或OUTER JOIN）连接到这个虚拟根表上。每个ORM模型的连接条件都与虚拟根表中的特定标识符关联。
3. ORM对象映射： 利用SQLAlchemy的add_columns()方法，指示会话将连接结果中的特定部分映射为ORM对象。

下面通过一个具体的例子来演示如何实现：

示例模型定义

首先，我们定义两个简单的ORM模型：Country 和 User。

from sqlalchemy import create_engine, Column, Integer, String, BigInteger, func, cast, ARRAY, and_
from sqlalchemy.orm import sessionmaker, declarative_base, Session, load_only

Base = declarative_base()

class Country(Base):
    __tablename__ = 'countries'
    id = Column(Integer, primary_key=True)
    code = Column(String, unique=True, nullable=False)
    name = Column(String, nullable=False)
    population = Column(BigInteger) # 示例：一个可能需要延迟加载的列

    def __repr__(self):
        return f""

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    email = Column(String, unique=True, nullable=False)
    name = Column(String, nullable=False)
    age = Column(Integer) # 示例：一个可能需要延迟加载的列

    def __repr__(self):
        return f""

# 数据库初始化（仅用于演示）
# engine = create_engine('postgresql://user:password@host:port/database')
# Base.metadata.create_all(engine)
# SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 假设我们已经有了一个session实例
# session: Session = SessionLocal()

实现虚拟连接查询

现在，我们将重写原始问题中的my_func函数，使用虚拟连接策略来获取Country和User对象。

Shakespeare

一款人工智能文案软件，能够创建几乎任何类型的文案。

下载

from sqlalchemy import select, func, cast, ARRAY, BigInteger, and_
from sqlalchemy.orm import Session

def my_func_optimized(
    session: Session, country_code: str, user_email: str
) -> tuple[Country | None, User | None]:
    """
    通过虚拟连接策略，在一个查询中获取Country和User对象。

    Args:
        session: SQLAlchemy会话实例。
        country_code: 国家代码。
        user_email: 用户邮箱。

    Returns:
        一个包含Country对象和User对象的元组。
    """
    # 1. 创建虚拟根表：使用func.unnest生成一个包含0和1的虚拟列
    # 0 将用于User，1 将用于Country
    col_ids = func.unnest(cast([0, 1], ARRAY(BigInteger))).alias("col_ids_alias")

    # 2. 构建基础查询，从虚拟根表开始
    q = select(col_ids.column).select_from(col_ids)

    # 3. 条件性外部连接User表，并将其ORM对象添加到结果中
    # 当虚拟ID为0时，尝试连接User表
    q = q.outerjoin(
        User,
        and_(col_ids.column == 0, User.email == user_email)
    ).add_columns(User) # 告知SQLAlchemy将User行映射为User对象

    # 4. 条件性外部连接Country表，并将其ORM对象添加到结果中
    # 当虚拟ID为1时，尝试连接Country表
    q = q.outerjoin(
        Country,
        and_(col_ids.column == 1, Country.code == country_code)
    ).add_columns(Country) # 告知SQLAlchemy将Country行映射为Country对象

    # 5. 执行查询并处理结果
    # session.execute(q) 返回的每一行包含：
    # (虚拟ID, User对象或None, Country对象或None)

    # 示例结果行：
    # (0, , None)
    # (1, None, )

    # 将结果转换为一个字典，以便按虚拟ID轻松访问ORM对象
    result_map = {}
    for key, user_obj, country_obj in session.execute(q):
        # 筛选出非None的ORM对象
        if user_obj:
            result_map[key] = user_obj
        elif country_obj:
            result_map[key] = country_obj

    # 根据虚拟ID返回对应的ORM对象
    # 注意：这里0对应User，1对应Country，所以返回顺序是 (Country, User)
    return result_map.get(1), result_map.get(0)

代码解释

1. func.unnest(cast([0, 1], ARRAY(BigInteger))).alias("col_ids_alias"):

   • cast([0, 1], ARRAY(BigInteger))：将Python列表[0, 1]转换为PostgreSQL的BIGINT[]数组类型。
   • func.unnest(...)：PostgreSQL的unnest函数将数组展开为一系列行。这里它会生成两行，一行包含0，一行包含1。
   • .alias("col_ids_alias")：为这个虚拟列起一个别名，方便后续引用。
   • 这一步创建了一个包含两行（值分别为0和1）的“虚拟表”，作为我们后续连接的起点。

2. select(col_ids.column).select_from(col_ids):

   • 构建基础查询，从我们创建的虚拟表中选择其唯一的列。

3. q.outerjoin(User, and_(col_ids.column == 0, User.email == user_email)).add_columns(User):

   • outerjoin(User, ...)：使用OUTER JOIN（外连接）来连接User表。这意味着即使没有匹配的User，虚拟根表的行也会保留。
   • and_(col_ids.column == 0, User.email == user_email)：这是连接条件。只有当虚拟ID为0（我们约定用于User）且User.email匹配时，才会连接User表。
   • add_columns(User)：这是关键。它告诉SQLAlchemy，将连接结果中与User模型对应的所有列数据，映射成一个User对象，并将其作为结果集的一部分。如果User表没有匹配，则该位置为None。

4. q.outerjoin(Country, and_(col_ids.column == 1, Country.code == country_code)).add_columns(Country):

   • 与User表的连接类似，但这里使用虚拟ID 1 来连接Country表，并将其映射为Country对象。

5. 结果处理:

   • session.execute(q)执行查询后，每一行结果将是一个元组，例如 (虚拟ID, User对象或None, Country对象或None)。
   • 我们通过遍历结果并根据虚拟ID（key）将实际的ORM对象存入result_map字典中。
   • 最后，根据约定的虚拟ID（1对应Country，0对应User）从result_map中取出并返回对应的ORM对象。

优势与注意事项

优势

• 单次数据库查询： 无论需要获取多少种不同类型的ORM对象，只要它们能通过LEFT JOIN连接到虚拟根表，就可以在一个数据库往返中完成数据获取，显著提高效率。
• 持久化的ORM对象： 返回的Country和User对象是完全持久化的，它们与当前会话关联，可以像通过session.get()或session.scalar(select(...))获取的对象一样进行操作。
• 自动延迟加载： 未在add_columns()中明确加载的ORM属性（例如Country.population或User.age）将自动支持延迟加载。当首次访问这些属性时，SQLAlchemy会触发额外的查询来获取它们。
• 接口兼容性： 这种方法允许我们重构内部的数据获取逻辑，而无需改变外部调用代码对返回ORM对象的期望。

注意事项

• 数据库特异性： func.unnest()是PostgreSQL特有的函数。对于其他数据库（如MySQL或SQLite），可能需要使用其他方法来创建虚拟根表，例如：
  • 通用方法： 使用CTE（Common Table Expressions）结合UNION ALL来生成虚拟ID，然后LEFT JOIN。
  • MySQL/SQLite： 可以创建一个包含数字序列的临时表或使用一系列SELECT ... UNION ALL SELECT ...来模拟。
• 结果集宽度： 当连接的表数量非常多时，session.execute(q)返回的结果行可能会包含大量的None值，导致结果集在逻辑上显得“宽”且“稀疏”。虽然PostgreSQL等现代数据库在处理这类查询时通常表现良好，但在极端情况下仍需考虑其对内存和网络带宽的潜在影响。
• 复杂性： 相比直接查询单一模型，这种多模型虚拟连接查询的构建逻辑更为复杂，需要仔细设计连接条件和结果处理逻辑。

总结

通过利用unnest()函数和LEFT JOIN构建虚拟根表的策略，我们可以在SQLAlchemy中实现高效、灵活的多类型ORM对象加载。这种方法不仅解决了从部分数据创建持久化、带延迟加载列的ORM对象的难题，还通过单次数据库查询优化了性能。尽管它对数据库类型有一定要求，并且查询构建稍显复杂，但对于需要精细控制数据加载和保持ORM对象完整性的高级应用场景，这无疑是一种强大而实用的技术。