Integrating LangGraph's AsyncPostgresStore with Alembic for Robust Schema Management

[mttcnnff]

I'm curious if there's been any consideration given to enabling the Postgres-backed Store to integrate seamlessly with migration frameworks like Alembic. Requiring the setup() function to be called at application startup—potentially modifying database schemas—is a significant divergence from conventional schema management practices. Moreover, this approach eliminates the possibility of rollback in case an issue arises.

I'd love to know if there's an existing roadmap or discussion about addressing this concern, or if others in the community share this perspective or have identified alternative solutions.

[ThomasAlbrt]

Absolutely!

[gmyrianthous]

@mttcnnff Do we have to call setup() at application startup or only once, during the very first deployment?
For instance, if you re-deploy the app, say for a second, or third time, you shouldn't call setup() again. Is my understanding correct?

[IhorPeresunko]

I have solved it by using a separate scheme from my main application's scheme.
For example: alembic app is using default "public" scheme in postgres, and for langgraph I have created "langgraph_memory" scheme. In this way they are not conflicting, and each is using their own connection.

def _with_schema_search_path(db_url: str, schema: str = "langgraph_memory") -> str:
    parsed = urlparse(db_url)
    query_pairs = parse_qsl(parsed.query, keep_blank_values=True)
    options_value = quote(f"-c search_path={schema},public", safe="")
    base_query = urlencode(query_pairs)
    new_query = f"{base_query}&options={options_value}" if base_query else f"options={options_value}"
    return urlunparse(
        (
            parsed.scheme,
            parsed.netloc,
            parsed.path,
            parsed.params,
            new_query,
            parsed.fragment,
        )
    )

def create_postgres_checkpointer() -> AsyncPostgresSaver:
    db_url = _with_schema_search_path(settings.db_url, schema="langgraph_memory")
    return AsyncPostgresSaver.from_conn_string(conn_string=db_url)

async def setup_checkpoint_tables():
    async with await AsyncConnection.connect(settings.db_url, autocommit=True) as conn:
        async with conn.cursor() as cur:
            await cur.execute("CREATE SCHEMA IF NOT EXISTS langgraph_memory;")

    db_url_with_schema = _with_schema_search_path(settings.db_url, schema="langgraph_memory")
    async with AsyncPostgresSaver.from_conn_string(db_url_with_schema) as checkpointer:
        await checkpointer.setup()
        return checkpointer