From e41af15d75dd5164d68803b6ecb1af4298f7a9b9 Mon Sep 17 00:00:00 2001 From: Mrxh Date: Mon, 21 Sep 2026 05:06:09 +0800 Subject: [PATCH] fix(v2): guard CREATE EXTENSION with the vector extension advisory lock `PGEngine.ainit_vectorstore_table` ran `CREATE EXTENSION IF NOT EXISTS vector` without the advisory lock that the legacy `PGVector` path takes in `_create_vector_extension`. Postgres lets two concurrent sessions both pass the existence check and then collide on `pg_extension_name_index`, so concurrent initializations could fail with a duplicate key error. Take the same lock key as the legacy path so both code paths serialize against each other. --- langchain_postgres/v2/engine.py | 20 ++++++++++++++ tests/unit_tests/v2/test_engine.py | 43 ++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/langchain_postgres/v2/engine.py b/langchain_postgres/v2/engine.py index 53d5c77b..9a899f85 100644 --- a/langchain_postgres/v2/engine.py +++ b/langchain_postgres/v2/engine.py @@ -13,6 +13,15 @@ T = TypeVar("T") +# Advisory lock key guarding `CREATE EXTENSION IF NOT EXISTS vector`. +# Postgres has no `IF NOT EXISTS` protection against concurrent sessions: two +# sessions can both pass the existence check and then collide on +# `pg_extension_name_index` with a duplicate key error. This key is +# intentionally the same one used by the legacy `_create_vector_extension` +# path in `langchain_postgres/vectorstores.py`, so the legacy and v2 code +# paths serialize against each other instead of only within themselves. +VECTOR_EXTENSION_ADVISORY_LOCK_KEY = 1573678846307946496 + class ColumnDict(TypedDict): name: str @@ -212,6 +221,17 @@ async def _ainit_vectorstore_table( id_column["name"] = self._escape_postgres_identifier(id_column["name"]) async with self._pool.connect() as conn: + # Take the same transaction-level advisory lock as the legacy + # `_create_vector_extension` path before creating the extension, so + # concurrent initializations (including one here racing a legacy + # `PGVector` init) are serialized. `pg_advisory_xact_lock` is held + # until the surrounding transaction ends, i.e. across the + # `CREATE EXTENSION` below and the `commit()` that follows it. + await conn.execute( + text( + f"SELECT pg_advisory_xact_lock({VECTOR_EXTENSION_ADVISORY_LOCK_KEY})" + ) + ) await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) await conn.commit() diff --git a/tests/unit_tests/v2/test_engine.py b/tests/unit_tests/v2/test_engine.py index 66f299aa..41eacfc7 100644 --- a/tests/unit_tests/v2/test_engine.py +++ b/tests/unit_tests/v2/test_engine.py @@ -1,3 +1,4 @@ +import asyncio import os import uuid from typing import AsyncIterator, Sequence @@ -24,6 +25,7 @@ HYBRID_SEARCH_TABLE_SYNC = "hybrid_sync" + str(uuid.uuid4()).replace("-", "_") CUSTOM_TYPEDDICT_TABLE_SYNC = "custom_td_sync" + str(uuid.uuid4()).replace("-", "_") INT_ID_CUSTOM_TABLE_SYNC = "custom_int_id_sync" + str(uuid.uuid4()).replace("-", "_") +LOCK_TABLE = "lock" + str(uuid.uuid4()).replace("-", "_") VECTOR_SIZE = 768 embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE) @@ -88,6 +90,47 @@ async def test_init_table(self, engine: PGEngine) -> None: stmt = f"INSERT INTO {DEFAULT_TABLE} (langchain_id, content, embedding) VALUES ('{id}', '{content}','{embedding_string}');" await aexecute(engine, stmt) + async def test_init_table_waits_for_vector_extension_lock( + self, engine: PGEngine + ) -> None: + """Creating the table must serialize on the vector-extension lock. + + Postgres can let two concurrent sessions both pass the + `IF NOT EXISTS` existence check for `vector` and then collide on + `pg_extension_name_index`. A concurrent initializer that already holds + the advisory lock therefore has to block this call until it commits; + without the lock the call would run straight through. + """ + # Same key as the advisory lock taken by the legacy + # `_create_vector_extension` path in `langchain_postgres/vectorstores.py`, + # so the two code paths serialize against each other. + lock_key = 1573678846307946496 + + async with engine._pool.connect() as concurrent_init: + await concurrent_init.execute( + text(f"SELECT pg_advisory_xact_lock({lock_key})") + ) + init_task = asyncio.create_task( + engine.ainit_vectorstore_table(LOCK_TABLE, VECTOR_SIZE) + ) + # The lock is held, so the table init cannot get past CREATE EXTENSION. + # `shield` keeps the task alive when `wait_for` times out. + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(asyncio.shield(init_task), timeout=1.0) + # Releasing the lock lets the blocked init proceed. + await concurrent_init.rollback() + + await asyncio.wait_for(init_task, timeout=20.0) + try: + result = await afetch( + engine, + "SELECT count(*) AS count FROM information_schema.tables" + f" WHERE table_name = '{LOCK_TABLE}';", + ) + assert result[0]["count"] == 1 + finally: + await aexecute(engine, f'DROP TABLE IF EXISTS "{LOCK_TABLE}"') + async def test_engine_args(self, engine: PGEngine) -> None: assert "Pool size: 3" in engine._pool.pool.status()