使用SQLModel连接Docker PostgreSQL时遇连接意外关闭问题

实现代码 #

from contextlib import contextmanager

import docker
from sqlmodel import create_engine, SQLModel, Field

DEFAULT_POSTGRES_PORT = 5432


class Foo(SQLModel, table=True):
    id_: int = Field(primary_key=True)


@contextmanager
def postgres_engine():
    db_pass = "foo"
    host_port = 1234

    client = docker.from_env()

    container = client.containers.run(
        "postgres",
        ports={DEFAULT_POSTGRES_PORT: host_port},
        environment={"POSTGRES_PASSWORD": db_pass},
        detach=True,
    )

    try:
        engine = create_engine(
            f"postgresql://postgres:{db_pass}@localhost:{host_port}/postgres"
        )

        SQLModel.metadata.create_all(engine)

        yield engine
    finally:
        container.kill()
        container.remove()


with postgres_engine():
    pass

报错信息 #

sqlalchemy.exc.OperationalError: (psycopg2.OperationalError) connection to server at "localhost" (127.0.0.1), port 1234 failed: server closed the connection unexpectedly
    This probably means the server terminated abnormally
    before or while processing the request.

对比：CLI启动可正常连接 #

使用以下命令启动容器后，手动连接无问题：

docker run -it -e POSTGRES_PASSWORD=foo -p 1234:5432 postgres

问题原因与解决思路 #

容器启动成功不代表内部的PostgreSQL服务已经就绪，Python代码在容器启动后立刻发起连接，此时数据库服务还在初始化流程中，导致连接失败。而CLI启动时，你会等待容器日志输出database system is ready to accept connections后再操作，因此不会触发错误。

解决方法 #

核心是等待PostgreSQL服务完全就绪后再建立连接，两种常用实现方式：

1. 连接重试轮询：多次尝试连接，直到成功或超时

from contextlib import contextmanager
import time
from sqlalchemy.exc import OperationalError

import docker
from sqlmodel import create_engine, SQLModel, Field

DEFAULT_POSTGRES_PORT = 5432


class Foo(SQLModel, table=True):
    id_: int = Field(primary_key=True)


@contextmanager
def postgres_engine():
    db_pass = "foo"
    host_port = 1234

    client = docker.from_env()

    container = client.containers.run(
        "postgres",
        ports={DEFAULT_POSTGRES_PORT: host_port},
        environment={"POSTGRES_PASSWORD": db_pass},
        detach=True,
    )

    try:
        engine = None
        max_retries = 10
        retry_interval = 2
        # 循环尝试连接，直到服务就绪
        for _ in range(max_retries):
            try:
                engine = create_engine(
                    f"postgresql://postgres:{db_pass}@localhost:{host_port}/postgres"
                )
                # 发起实际连接验证服务状态
                with engine.connect():
                    break
            except OperationalError:
                time.sleep(retry_interval)
        else:
            raise RuntimeError("PostgreSQL容器启动后超时未就绪")

        SQLModel.metadata.create_all(engine)

        yield engine
    finally:
        container.kill()
        container.remove()


with postgres_engine():
    pass

2. 监听容器日志：检查日志中是否出现服务就绪的标识字符串

# 替换原代码中创建引擎前的逻辑
ready_flag = "database system is ready to accept connections"
while True:
    logs = container.logs().decode("utf-8")
    if ready_flag in logs:
        break
    time.sleep(1)

engine = create_engine(f"postgresql://postgres:{db_pass}@localhost:{host_port}/postgres")

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