Using the Postgres checkpointer with psycopg_pool results in Connection Closed error

[GaryFail]

I am experiencing the scenario where the max_lifetime on a connection pool has been exceeded by one of the connections i.e. a simple example here where the max lifetime is set to 15 seconds and the pool size set to 1:

database_connection_pool = AsyncConnectionPool(
    conninfo=os.environ["DATABASE_URL"],
    min_size=1,
    max_size=1,
    max_lifetime=15.0,
    open=False,
    kwargs=connection_kwargs)

I compile the graph like so:

async def get_ai_assistant() -> CompiledStateGraph:
    async with database_connection_pool.connection() as conn:
        saver = AsyncPostgresSaver(conn)
        graph = builder.compile(checkpointer=saver)
        return graph

When this lifetime has been exceeded I receive a connection closed error. I can see from the psycopg logs that the connection is discarded when returned back to the pool and a new one is added. If I then do the same operation again I get a response back from the graph. Of course, once the 15 seconds have elapsed a second time I get the connection closed error again.

The strange thing is, if I simply query the database via a cursor in psycopg then even after the max_lifetime has surpassed it will still work, it's as though the psycopg code is checking the state of the connection and prior to using it and if closed then returning it to the pool and pulling a new one:

async def query_database() -> int:
    async with database_connection_pool.connection() as conn:
        print("connection acquired")
        async with conn.cursor() as cur:
            print("cursor acquired")
            await cur.execute('SELECT COUNT(*) FROM public."checkpoints"')
            print("query executed")
            result = await cur.fetchone()
            print("result fetched")
            if result is None:
                return 0
            else:
                count = result["count"]
                return count

The output of the logs is as follows:

I never get any connection closed errors when simply using the cursor.

I have tried calling check() on the connection pool but get errors when I do so:

async def get_ai_assistant() -> CompiledStateGraph:
    await database_connection_pool.check()
    async with database_connection_pool.connection() as conn:
        saver = AsyncPostgresSaver(conn)
        graph = builder.compile(checkpointer=saver)
        return graph

[vbarda]

@GaryFail the problem is here:

async with database_connection_pool.connection() as conn:
        saver = AsyncPostgresSaver(conn)
        graph = builder.compile(checkpointer=saver)
        return graph

as soon as you leave the async with statement the connection is closed. you might want to implement something similar to this instead

in the meantime, we'll be adding proper support to just passing connection pool directly to PostgresSaver/AsyncPostgresSaver

[vbarda]

also, the pool.check() is failing because you need to pass connection to it (like pool.check(conn)), but i don't think that would solve your problem

[GaryFail]

Hi @vbarda thanks for the reply. The connection isn't closed at that point, the connection will stay open for as long as it is set in the max_lifetime so if it left to the default setting the connection will remain open for 1 hour. It will then close and if I have say four connections then I will need to retry 4 times in order to clear out the stale connections from the pool and then it works again for another hour.

If I follow the example you have linked to, as far as I can see it is only pulling one connection from the pool at the time of the app starting and therefore every hit to that endpoint will use the same connection.

This is what I currently have in my lifespan function:

@asynccontextmanager
async def lifespan(app: FastAPI):
    await database_connection_pool.open(wait = True)
    async with database_connection_pool.connection() as conn:
        await AsyncPostgresSaver(conn).setup()
    yield
    await database_connection_pool.close()

I appreciate I may be wrong here, more than likely I am as I am new to Python but I can say that with the way I have it configured it works fine for an hour (or however long the max_lifetime is set) and then I get connection failed errors whilst the pool is replacing them and then it works fine again for the next period of max_lifetime.

I hope that makes sense.

Also the pool check doesn't have a parameter for the connection, it's the check_connection function that does which is called from the check function.

[vbarda]

@GaryFail we just released langgraph-checkpoint-postgres==1.0.4 which supports passing connection pool to the postgres saver -- could you try it and see if it resolves your issues?

[GaryFail]

Hi @vbarda that's great thanks. I've tested it this morning and it looks to be working really well. I did have to install langgraph-checkpoint==1.0.6 explicitly though, otherwise when I imported the AsyncPostgresSaver I was getting an error about failing to import WRITES_IDX_MAP from langgraph.checkpoint.base

[vbarda]

thanks for the heads up -- will fix that up

[TonySimonovsky]

same problem with 1.0.8