diff --git a/examples/pg_vectorstore.ipynb b/examples/pg_vectorstore.ipynb index 340abe1..f41809a 100644 --- a/examples/pg_vectorstore.ipynb +++ b/examples/pg_vectorstore.ipynb @@ -493,6 +493,26 @@ "source": [ "await vectorstore.adrop_vector_index() # Drop index using default name" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Clean up\n", + "\n", + "**⚠️ WARNING: this can not be undone**\n", + "\n", + "Drop the vector store table." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "await pg_engine.adrop_table(TABLE_NAME)" + ] } ], "metadata": { diff --git a/examples/pg_vectorstore_how_to.ipynb b/examples/pg_vectorstore_how_to.ipynb index aad661e..de1dc83 100644 --- a/examples/pg_vectorstore_how_to.ipynb +++ b/examples/pg_vectorstore_how_to.ipynb @@ -691,12 +691,18 @@ "source": [ "## Clean up\n", "\n", - "#### ⚠️ WARNING: this can not be undone\n", - "In a [`psql`](https://www.postgresql.org/docs/current/app-psql.html) client, run:\n", + "**⚠️ WARNING: this can not be undone**\n", "\n", - "```\n", - "DROP TABLE ;\n", - "```" + "Drop the vector store table." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "await pg_engine.adrop_table(TABLE_NAME)" ] } ], diff --git a/langchain_postgres/v2/engine.py b/langchain_postgres/v2/engine.py index dfd6c3b..c2a0d93 100644 --- a/langchain_postgres/v2/engine.py +++ b/langchain_postgres/v2/engine.py @@ -346,3 +346,35 @@ def init_vectorstore_table( store_metadata=store_metadata, ) ) + + async def _adrop_table( + self, + table_name: str, + *, + schema_name: str = "public", + ) -> None: + """Drop the vector store table""" + query = f'DROP TABLE "{schema_name}"."{table_name}";' + async with self._pool.connect() as conn: + await conn.execute(text(query)) + await conn.commit() + + async def adrop_table( + self, + table_name: str, + *, + schema_name: str = "public", + ) -> None: + await self._run_as_async( + self._adrop_table(table_name=table_name, schema_name=schema_name) + ) + + async def drop_table( + self, + table_name: str, + *, + schema_name: str = "public", + ) -> None: + self._run_as_sync( + self._adrop_table(table_name=table_name, schema_name=schema_name) + ) diff --git a/tests/unit_tests/v2/test_async_pg_vectorstore.py b/tests/unit_tests/v2/test_async_pg_vectorstore.py index bcf37e8..acad83f 100644 --- a/tests/unit_tests/v2/test_async_pg_vectorstore.py +++ b/tests/unit_tests/v2/test_async_pg_vectorstore.py @@ -50,8 +50,8 @@ async def engine(self) -> AsyncIterator[PGEngine]: engine = PGEngine.from_connection_string(url=CONNECTION_STRING) yield engine - await aexecute(engine, f'DROP TABLE IF EXISTS "{DEFAULT_TABLE}"') - await aexecute(engine, f'DROP TABLE IF EXISTS "{CUSTOM_TABLE}"') + await engine.adrop_table(DEFAULT_TABLE) + await engine.adrop_table(CUSTOM_TABLE) await engine.close() @pytest_asyncio.fixture(scope="class") diff --git a/tests/unit_tests/v2/test_async_pg_vectorstore_search.py b/tests/unit_tests/v2/test_async_pg_vectorstore_search.py index 8e5e371..72f91d8 100644 --- a/tests/unit_tests/v2/test_async_pg_vectorstore_search.py +++ b/tests/unit_tests/v2/test_async_pg_vectorstore_search.py @@ -66,9 +66,9 @@ class TestVectorStoreSearch: async def engine(self) -> AsyncIterator[PGEngine]: engine = PGEngine.from_connection_string(url=CONNECTION_STRING) yield engine - await aexecute(engine, f"DROP TABLE IF EXISTS {DEFAULT_TABLE}") - await aexecute(engine, f"DROP TABLE IF EXISTS {CUSTOM_TABLE}") - await aexecute(engine, f"DROP TABLE IF EXISTS {CUSTOM_FILTER_TABLE}") + await engine.adrop_table(DEFAULT_TABLE) + await engine.adrop_table(CUSTOM_TABLE) + await engine.adrop_table(CUSTOM_FILTER_TABLE) await engine.close() @pytest_asyncio.fixture(scope="class")