chore: Drop vector indexes before dropping table/database (#436)

Fixes a bug where we were dropping vector indexes when the connection
pool is disposed, which lead to missing vector indexes.

---------

Co-authored-by: JULIA OFFERMAN <juliaofferman@google.com>
Co-authored-by: Jack Wotherspoon <jackwoth@google.com>

retrieval_service/datastore/providers/cloudsql_mysql.py

@@ -81,13 +81,21 @@ async def create(cls, config: Config) -> "Client":
         pool = await loop.run_in_executor(None, cls.create_sync, config)
         return pool
 
+    def drop_vector_indexes(self):
+        with self.__pool.connect() as conn:
+            s = text("""SELECT index_name FROM mysql.vector_indexes""")
+            results = (conn.execute(s)).mappings().fetchall()
+            for r in results:
+                conn.execute(text(f"CALL mysql.drop_vector_index('{r['index_name']}')"))
+
     def initialize_data_sync(
         self,
         airports: list[models.Airport],
         amenities: list[models.Amenity],
         flights: list[models.Flight],
         policies: list[models.Policy],
     ) -> None:
+        self.drop_vector_indexes()
         with self.__pool.connect() as conn:
             # If the table already exists, drop it to avoid conflicts
             conn.execute(text("DROP TABLE IF EXISTS airports"))
@@ -815,17 +823,4 @@ async def policies_search(
         return res
 
     async def close(self):
-        # Vector indexes must be dropped before any DDLs on the base table are permitted
-        with self.__pool.connect() as conn:
-            s = text(
-                """
-                CALL mysql.drop_vector_index(:index_name)
-                """
-            )
-            params = [
-                {"index_name": f"{self.__db_name}.amenities_index"},
-                {"index_name": f"{self.__db_name}.policies_index"},
-            ]
-
-            conn.execute(s, parameters=params)
         self.__pool.dispose()

retrieval_service/datastore/providers/cloudsql_mysql_test.py

@@ -133,6 +133,9 @@ async def ds(
     if ds is None:
         raise TypeError("datastore creation failure")
     yield ds
+
+    if isinstance(ds, cloudsql_mysql.Client):
+        ds.drop_vector_indexes()
     await ds.close()
 
 

retrieval_service/requirements.txt

@@ -18,4 +18,4 @@ langchain-google-vertexai==1.0.5
 asyncio==3.4.3
 datetime==5.5
 pymysql==1.1.1
-types-PyMySQL==1.1.0.20240524
+types-PyMySQL==1.1.0.20240524