Postgres docker "server closed the connection unexpectedly"

Question

I want to run and connect to the postgresql Docker image in Python using SQLModel. Here's my attempt

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

I'm seeing

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.

If instead of using the Python Docker SDK I use the CLI with

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

I don't see an error.

Answer 1

You have a race condition in your code. Take a look at the logs when you start up a Postgres container:

$ docker run --rm -e POSTGRES_PASSWORD=secret postgres:14
The files belonging to this database system will be owned by user "postgres".
This user must also own the server process.

The database cluster will be initialized with locale "en_US.utf8".
The default database encoding has accordingly been set to "UTF8".
The default text search configuration will be set to "english".
.
.
.

2023-01-28 02:02:54.575 UTC [1] LOG:  database system is ready to accept connections

All of that initialization takes a second or so to run. On the other hand, your code immediately tries to interact with the database after creating the container:

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

        SQLModel.metadata.create_all(engine)

In other words, it's trying to connect to the database before the database is ready to handle connections. The best solution is to block until you are able to communicate successfully with the Postgres server. Something like this:

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

        while True:
            try:
                with engine.connect() as conn:
                   conn.execute(sqlalchemy.select(1))
            except sqlalchemy.exc.OperationalError:
                print('retrying connection')
                time.sleep(1)
                continue
            else:
                break

        SQLModel.metadata.create_all(engine)

With this code in place, running the script produces this output:

retrying connection
retrying connection

...and then exits without error. The complete script, with the above modifications, looks like:

from contextlib import contextmanager
import time

import docker
import sqlalchemy
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 = "secret"
    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"
        )

        while True:
            try:
                with engine.connect() as conn:
                   conn.execute(sqlalchemy.select(1))
            except sqlalchemy.exc.OperationalError:
                print('retrying connection')
                time.sleep(1)
                continue
            else:
                break

        SQLModel.metadata.create_all(engine)

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


with postgres_engine() as engine:
    pass