diff --git a/examples/pg_vectorstore_how_to.ipynb b/examples/pg_vectorstore_how_to.ipynb index d3b35bf9..1fd63a7e 100644 --- a/examples/pg_vectorstore_how_to.ipynb +++ b/examples/pg_vectorstore_how_to.ipynb @@ -208,6 +208,22 @@ "```" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Optional Tip: 💡\n", + "By default, `init_vectorstore_table()` runs `CREATE EXTENSION IF NOT EXISTS vector` to ensure the `pgvector` extension is available. If the extension is already installed, or the database user does not have permission to create extensions (common on managed databases such as Cloud SQL or AlloyDB where an admin enables it beforehand), you can skip this step by passing `create_extension=False`:\n", + "\n", + "```python\n", + "await pg_engine.ainit_vectorstore_table(\n", + " table_name=TABLE_NAME,\n", + " vector_size=VECTOR_SIZE,\n", + " create_extension=False, # Default: True\n", + ")\n", + "```" + ] + }, { "cell_type": "markdown", "metadata": {}, diff --git a/langchain_postgres/v2/engine.py b/langchain_postgres/v2/engine.py index 53d5c77b..61a73301 100644 --- a/langchain_postgres/v2/engine.py +++ b/langchain_postgres/v2/engine.py @@ -159,6 +159,7 @@ async def _ainit_vectorstore_table( overwrite_existing: bool = False, store_metadata: bool = True, hybrid_search_config: Optional[HybridSearchConfig] = None, + create_extension: bool = True, ) -> None: """ Create a table for saving of vectors to be used with PGVectorStore. @@ -183,6 +184,9 @@ async def _ainit_vectorstore_table( Default: True. hybrid_search_config (HybridSearchConfig): Hybrid search configuration. Default: None. + create_extension (bool): Whether to create the pgvector extension if it + does not exist. Default: True. Disable if the extension is already + installed or if the database user lacks the required permissions. Raises: :class:`DuplicateTableError `: if table already exists. @@ -211,9 +215,10 @@ async def _ainit_vectorstore_table( self._validate_column_dict(id_column) id_column["name"] = self._escape_postgres_identifier(id_column["name"]) - async with self._pool.connect() as conn: - await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) - await conn.commit() + if create_extension: + async with self._pool.connect() as conn: + await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + await conn.commit() if overwrite_existing: async with self._pool.connect() as conn: @@ -277,6 +282,7 @@ async def ainit_vectorstore_table( overwrite_existing: bool = False, store_metadata: bool = True, hybrid_search_config: Optional[HybridSearchConfig] = None, + create_extension: bool = True, ) -> None: """ Create a table for saving of vectors to be used with PGVectorStore. @@ -303,6 +309,9 @@ async def ainit_vectorstore_table( Note that queries might be slow if the hybrid search column does not exist. For best hybrid search performance, consider creating a TSV column and adding GIN index. Default: None. + create_extension (bool): Whether to create the pgvector extension if it + does not exist. Default: True. Disable if the extension is already + installed or if the database user lacks the required permissions. """ await self._run_as_async( self._ainit_vectorstore_table( @@ -317,6 +326,7 @@ async def ainit_vectorstore_table( overwrite_existing=overwrite_existing, store_metadata=store_metadata, hybrid_search_config=hybrid_search_config, + create_extension=create_extension, ) ) @@ -334,6 +344,7 @@ def init_vectorstore_table( overwrite_existing: bool = False, store_metadata: bool = True, hybrid_search_config: Optional[HybridSearchConfig] = None, + create_extension: bool = True, ) -> None: """ Create a table for saving of vectors to be used with PGVectorStore. @@ -360,6 +371,9 @@ def init_vectorstore_table( Note that queries might be slow if the hybrid search column does not exist. For best hybrid search performance, consider creating a TSV column and adding GIN index. Default: None. + create_extension (bool): Whether to create the pgvector extension if it + does not exist. Default: True. Disable if the extension is already + installed or if the database user lacks the required permissions. """ self._run_as_sync( self._ainit_vectorstore_table( @@ -374,6 +388,7 @@ def init_vectorstore_table( overwrite_existing=overwrite_existing, store_metadata=store_metadata, hybrid_search_config=hybrid_search_config, + create_extension=create_extension, ) ) diff --git a/tests/unit_tests/v2/test_engine.py b/tests/unit_tests/v2/test_engine.py index 66f299aa..01843bf0 100644 --- a/tests/unit_tests/v2/test_engine.py +++ b/tests/unit_tests/v2/test_engine.py @@ -19,11 +19,13 @@ HYBRID_SEARCH_TABLE = "hybrid" + str(uuid.uuid4()).replace("-", "_") CUSTOM_TYPEDDICT_TABLE = "custom_td" + str(uuid.uuid4()).replace("-", "_") INT_ID_CUSTOM_TABLE = "custom_int_id" + str(uuid.uuid4()).replace("-", "_") +NO_CREATE_EXT_TABLE = "no_ext" + str(uuid.uuid4()).replace("-", "_") DEFAULT_TABLE_SYNC = "default_sync" + str(uuid.uuid4()).replace("-", "_") CUSTOM_TABLE_SYNC = "custom_sync" + str(uuid.uuid4()).replace("-", "_") 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("-", "_") +NO_CREATE_EXT_TABLE_SYNC = "no_ext_sync" + str(uuid.uuid4()).replace("-", "_") VECTOR_SIZE = 768 embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE) @@ -76,6 +78,7 @@ async def engine(self) -> AsyncIterator[PGEngine]: await aexecute(engine, f'DROP TABLE IF EXISTS "{CUSTOM_TYPEDDICT_TABLE}"') await aexecute(engine, f'DROP TABLE IF EXISTS "{DEFAULT_TABLE}"') await aexecute(engine, f'DROP TABLE IF EXISTS "{INT_ID_CUSTOM_TABLE}"') + await aexecute(engine, f'DROP TABLE IF EXISTS "{NO_CREATE_EXT_TABLE}"') await engine.close() async def test_init_table(self, engine: PGEngine) -> None: @@ -251,6 +254,15 @@ async def test_column(self, engine: PGEngine) -> None: with pytest.raises(ValueError): Column(1, "INTEGER") # type: ignore + async def test_init_table_no_create_extension(self, engine: PGEngine) -> None: + """create_extension=False should not attempt CREATE EXTENSION but still create the table.""" + await engine.ainit_vectorstore_table( + NO_CREATE_EXT_TABLE, VECTOR_SIZE, create_extension=False + ) + stmt = f"SELECT column_name FROM information_schema.columns WHERE table_name = '{NO_CREATE_EXT_TABLE}';" + results = await afetch(engine, stmt) + assert len(results) > 0 + @pytest.mark.enable_socket @pytest.mark.asyncio @@ -264,6 +276,7 @@ async def engine(self) -> AsyncIterator[PGEngine]: await aexecute(engine, f'DROP TABLE IF EXISTS "{DEFAULT_TABLE_SYNC}"') await aexecute(engine, f'DROP TABLE IF EXISTS "{INT_ID_CUSTOM_TABLE_SYNC}"') await aexecute(engine, f'DROP TABLE IF EXISTS "{CUSTOM_TYPEDDICT_TABLE_SYNC}"') + await aexecute(engine, f'DROP TABLE IF EXISTS "{NO_CREATE_EXT_TABLE_SYNC}"') await engine.close() async def test_init_table(self, engine: PGEngine) -> None: @@ -410,3 +423,12 @@ async def test_engine_constructor_key( key = object() with pytest.raises(Exception): PGEngine(key, engine._pool, None, None) + + async def test_init_table_no_create_extension(self, engine: PGEngine) -> None: + """create_extension=False should not attempt CREATE EXTENSION but still create the table.""" + engine.init_vectorstore_table( + NO_CREATE_EXT_TABLE_SYNC, VECTOR_SIZE, create_extension=False + ) + stmt = f"SELECT column_name FROM information_schema.columns WHERE table_name = '{NO_CREATE_EXT_TABLE_SYNC}';" + results = await afetch(engine, stmt) + assert len(results) > 0