如何保持SQLAlchemy查询时内存占用稳定？

我编写了如下Python脚本用于测试SQLAlchemy的内存占用：

from typing import Any, Tuple
import sys

def find_memory_usage(insertion: bool):
    from sqlalchemy import create_engine, select, String
    from sqlalchemy.orm import MappedAsDataclass, DeclarativeBase, declared_attr
    from sqlalchemy.orm import Mapped, mapped_column
    import random
    import string
    from sqlalchemy.orm import sessionmaker
    from datetime import datetime
    import gc
    import psutil

    class Base(MappedAsDataclass, DeclarativeBase):
        __abstract__ = True

        @declared_attr
        def __tablename__(self) -> Any:
            return self.__name__

    class User(Base):
        id: Mapped[int] = mapped_column(init=False, name="ID", primary_key=True)
        name: Mapped[str] = mapped_column(String(16), name="Name", nullable=False)
        description: Mapped[str] = mapped_column(
            String(256), name="Description", nullable=False
        )

    engine = create_engine("postgresql+psycopg2://postgres:password@localhost/postgres")
    Base.metadata.create_all(engine, tables=[User.__table__])

    # Function to generate a random string
    def random_string(length):
        letters = string.ascii_letters
        return "".join(random.choice(letters) for i in range(length))

    # Set up the session maker
    Session = sessionmaker(bind=engine)

    with Session() as session:
        if insertion:
            # Insert 100,000 User rows with random data
            for _ in range(100000):
                user = User(
                    name=random_string(16),  # 10 characters long names
                    description=random_string(128),  # 128 characters long descriptions
                )
                session.add(user)

            session.commit()
            print("Data insertion complete.")
            return

        print(datetime.now(), psutil.Process().memory_info().rss / (1024 * 1024))
        number = (int(sys.argv[1]) if len(sys.argv) > 1 else 100000)
        for row in session.scalars(select(User).fetch(number)):
            pass
        session.commit()
        session.expire_all()
        session.expunge_all()
        session.close()
        del session
    gc.collect()
    print(datetime.now(), psutil.Process().memory_info().rss / (1024 * 1024))

# find_memory_usage(True) # I have done this once already to insert data.
find_memory_usage(False)

执行不同参数时的结果：

• 执行python test.py 1000：

2024-01-21 22:59:37.291944 48.23828125
2024-01-21 22:59:37.321317 49.86328125

• 执行python test.py 100000：

2024-01-21 22:59:51.152666 47.8984375
2024-01-21 22:59:52.477458 73.06640625

可见内存占用增长与查询的行数直接相关，问题出在这段代码：

number = (int(sys.argv[1]) if len(sys.argv) > 1 else 100000)
for row in session.scalars(select(User).fetch(number)):
    pass

期望逐行处理查询结果时内存占用保持稳定，请问如何修改上述查询代码？

核心问题原因 #

fetch(number)会一次性将所有匹配的行加载到内存中，并且SQLAlchemy的Session会缓存这些对象，导致内存随查询行数线性上升。要实现稳定内存的逐行/分批处理，需要使用服务器端游标（Server-Side Cursor）配合分批获取数据的方式，避免一次性加载全量数据。

具体修改步骤 #

1. 修改引擎配置，启用服务器端游标
   在创建create_engine时添加server_side_cursors=True参数，让PostgreSQL使用服务器端游标处理大结果集：

engine = create_engine(
    "postgresql+psycopg2://postgres:password@localhost/postgres",
    server_side_cursors=True
)

2. 替换fetch()为yield_per()，分批获取数据
   将select(User).fetch(number)改为select(User).yield_per(1000)（这里的1000是每次从数据库获取的行数，可根据实际情况调整），同时在循环中及时从Session移除已处理的对象，避免缓存积累：

print(datetime.now(), psutil.Process().memory_info().rss / (1024 * 1024))
# number参数仅用于控制总处理行数，不再影响一次性加载量
total_rows = int(sys.argv[1]) if len(sys.argv) > 1 else 100000
processed = 0
# 每次从数据库获取1000行，处理后释放
for row in session.scalars(select(User).yield_per(1000)):
    if processed >= total_rows:
        break
    # 这里添加你的业务处理逻辑
    session.expunge(row)  # 从Session中移除对象，释放内存
    processed += 1

修改后效果 #

修改后，查询会分批从数据库获取数据，每处理完一批（或一行）就释放对应的内存，内存占用会保持在一个稳定的区间，不会随总处理行数的增加而大幅上升。

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