feat: add drop table (#195)

Author: averikitsch

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": {

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 <TABLE_NAME>;\n",
-    "```"
+    "Drop the vector store table."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "await pg_engine.adrop_table(TABLE_NAME)"
    ]
   }
  ],

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)
+        )

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")

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")