使用pytest_postgresql测试PostgreSQL时表创建失败问题排查

运行pytest测试时出现错误：

FAILED tests/test_model_with_test_db.py::test_authors - sqlalchemy.exc.ProgrammingError: (psycopg2.errors.UndefinedTable) relation "authors" does not exist

执行model.Base.metadata.create_all(con)后并未创建预期的表，相关代码如下：

测试代码 #

import pytest
from pytest_postgresql import factories
from pytest_postgresql.janitor import DatabaseJanitor
from sqlalchemy import create_engine, select
from sqlalchemy.orm.session import sessionmaker

import model

test_db = factories.postgresql_proc(port=None, dbname="test_db")


@pytest.fixture(scope="session")
def db_session(test_db):
    pg_host = test_db.host
    pg_port = test_db.port
    pg_user = test_db.user
    pg_password = test_db.password
    pg_db = test_db.dbname

    with DatabaseJanitor(pg_user, pg_host, pg_port, pg_db, test_db.version, pg_password):
        connection_str = f"postgresql+psycopg2://{pg_user}:@{pg_host}:{pg_port}/{pg_db}"
        engine = create_engine(connection_str)
        with engine.connect() as con:
            model.Base.metadata.create_all(con)
            yield sessionmaker(bind=engine, expire_on_commit=False)


@pytest.fixture(scope="module")
def create_test_data():
    authors = [
        ["John", "Smith", "john@gmail.com"],
        ["Bill", "Miles", "bill@gmail.com"],
        ["Frank", "James", "frank@gmail.com"]
    ]
    return [model.Author(firstname=firstname, lastname=lastname, email=email)
            for firstname, lastname, email in authors]


def test_persons(db_session, create_test_data):
    s = db_session()
    for obj in create_test_data:
        s.add(obj)
    s.commit()

    query_result = s.execute(select(model.Author)).all()
    s.close()

    assert len(query_result) == len(create_test_data)

model.py代码 #

from sqlalchemy import create_engine, Column, Integer, String, DateTime, Text, ForeignKey
from sqlalchemy.engine import URL
from sqlalchemy.orm import declarative_base, relationship, sessionmaker
from datetime import datetime

Base = declarative_base()


class Author(Base):
    __tablename__ = 'authors'

    id = Column(Integer(), primary_key=True)
    firstname = Column(String(100))
    lastname = Column(String(100))
    email = Column(String(255), nullable=False)
    joined = Column(DateTime(), default=datetime.now)

    articles = relationship('Article', backref='author')


class Article(Base):
    __tablename__ = 'articles'

    id = Column(Integer(), primary_key=True)
    slug = Column(String(100), nullable=False)
    title = Column(String(100), nullable=False)
    created_on = Column(DateTime(), default=datetime.now)
    updated_on = Column(DateTime(), default=datetime.now, onupdate=datetime.now)
    content = Column(Text)
    author_id = Column(Integer(), ForeignKey('authors.id'))


url = URL.create(
    drivername="postgresql",
    username="postgres",
    host="localhost",
    port=5433,
    database="andy"
)

engine = create_engine(url)
Session = sessionmaker(bind=engine)

问题原因 #

核心问题是DDL操作未提交事务：SQLAlchemy通过engine.connect()获取的连接默认处于事务模式，create_all执行建表DDL后，若未手动提交事务，当with engine.connect() as con上下文结束时，连接关闭，事务自动回滚，导致创建的表被撤销，后续测试自然找不到对应表。

解决方案 #

以下两种方式均可解决问题，确保DDL操作被持久化：

方式1：直接使用引擎执行create_all

create_all可直接接收SQLAlchemy引擎作为参数，此时会自动处理事务提交：

@pytest.fixture(scope="session")
def db_session(test_db):
    pg_host = test_db.host
    pg_port = test_db.port
    pg_user = test_db.user
    pg_password = test_db.password
    pg_db = test_db.dbname

    with DatabaseJanitor(pg_user, pg_host, pg_port, pg_db, test_db.version, pg_password):
        connection_str = f"postgresql+psycopg2://{pg_user}:@{pg_host}:{pg_port}/{pg_db}"
        engine = create_engine(connection_str)
        # 直接用引擎执行create_all，自动提交DDL
        model.Base.metadata.create_all(engine)
        yield sessionmaker(bind=engine, expire_on_commit=False)

方式2：手动提交连接事务

若需保留连接上下文，在create_all后手动提交事务：

@pytest.fixture(scope="session")
def db_session(test_db):
    pg_host = test_db.host
    pg_port = test_db.port
    pg_user = test_db.user
    pg_password = test_db.password
    pg_db = test_db.dbname

    with DatabaseJanitor(pg_user, pg_host, pg_port, pg_db, test_db.version, pg_password):
        connection_str = f"postgresql+psycopg2://{pg_user}:@{pg_host}:{pg_port}/{pg_db}"
        engine = create_engine(connection_str)
        with engine.connect() as con:
            model.Base.metadata.create_all(con)
            # 手动提交事务，持久化建表操作
            con.commit()
            yield sessionmaker(bind=engine, expire_on_commit=False)

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