fix oracle to store collection description (#130)

* fix oracle tables

use two tables to save collection information and vector items

* fix lint

* fix lint

* fix lint

* fix lint

Author: jin38324

deepsearcher/vector_db/oracle.py

@@ -48,6 +48,7 @@ def __init__(
                 increment=increment,
             )
             log.color_print(f"Connected to Oracle database at {dsn}")
+            self.check_table()
         except Exception as e:
             log.critical(f"Failed to connect to Oracle database at {dsn}")
             log.critical(f"Oracle database error in init: {e}")
@@ -148,30 +149,34 @@ def has_collection(self, collection: str = "deepsearcher"):
         else:
             return False
 
-    def has_table(self, collection: str = "deepsearcher"):
+    def check_table(self):
         SQL = SQL_TEMPLATES["has_table"]
-        res = self.query(SQL)
-        if res:
-            if res[0]["rowcnt"] > 0:
-                return True
-            else:
-                return False
-        else:
-            return False
+        try:
+            res = self.query(SQL)
+            if len(res) < 2:
+                missing_table = TABLES.keys() - set([i["table_name"] for i in res])
+                for table in missing_table:
+                    self.create_tables(table)
+        except Exception as e:
+            log.critical(f"Failed to check table in Oracle database, error info: {e}")
+            raise
 
-    def create_tables(self):
-        SQL = SQL_TEMPLATES["ddl"]
+    def create_tables(self, table_name):
+        SQL = TABLES[table_name]
         try:
             self.execute(SQL)
-            log.color_print("Created table DEEPSEARCHER in Oracle database")
-        except Exception:
-            log.critical("Failed to create table DEEPSEARCHER in Oracle database")
+            log.color_print(f"Created table {table_name} in Oracle database")
+        except Exception as e:
+            log.critical(f"Failed to create table {table_name} in Oracle database, error info: {e}")
             raise
 
     def drop_collection(self, collection: str = "deepsearcher"):
         try:
-            SQL = SQL_TEMPLATES["drop_collection"]
             params = {"collection": collection}
+            SQL = SQL_TEMPLATES["drop_collection"]
+            self.execute(SQL, params)
+
+            SQL = SQL_TEMPLATES["drop_collection_item"]
             self.execute(SQL, params)
             log.color_print(f"Collection {collection} dropped")
         except Exception as e:
@@ -227,22 +232,18 @@ def init_collection(
             collection = self.default_collection
         if description is None:
             description = ""
-        try:
-            has_table = self.has_table()
-            if not has_table:
-                self.create_table()
-        except Exception as e:
-            log.critical(f"fail to init db for oracle, error info: {e}")
-
         try:
             has_collection = self.has_collection(collection)
             if force_new_collection and has_collection:
                 self.drop_collection(collection)
             elif has_collection:
                 return
+            # insert collection info
+            SQL = SQL_TEMPLATES["insert_collection"]
+            params = {"collection": collection, "description": description}
+            self.execute(SQL, params)
         except Exception as e:
-            log.critical(f"fail to init db for milvus, error info: {e}")
-            raise
+            log.critical(f"fail to init_collection for oracle, error info: {e}")
 
     def insert_data(
         self,
@@ -315,16 +316,13 @@ def list_collections(self, *args, **kwargs) -> List[CollectionInfo]:
             collections = self.query(SQL)
             if collections:
                 for collection in collections:
-                    description = ""
                     collection_infos.append(
                         CollectionInfo(
                             collection_name=collection["collection"],
-                            description=description,
+                            description=collection["description"],
                         )
                     )
-                return collection_infos
-            else:
-                log.critical("No collections found")
+            return collection_infos
         except Exception as e:
             log.critical(f"fail to list collections, error info: {e}")
             raise
@@ -339,8 +337,15 @@ def clear_db(self, collection: str = "deepsearcher", *args, **kwargs):
             raise
 
 
-SQL_TEMPLATES = {
-    "ddl": """CREATE TABLE DEEPSEARCHER (    
+TABLES = {
+    "DEEPSEARCHER_COLLECTION_INFO": """CREATE TABLE DEEPSEARCHER_COLLECTION_INFO (    
+        id INT generated by default as identity primary key,
+        collection varchar(256),
+        description CLOB,
+        status NUMBER DEFAULT 1,
+        createtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+        updatetime TIMESTAMP DEFAULT NULL)""",
+    "DEEPSEARCHER_COLLECTION_ITEM": """CREATE TABLE DEEPSEARCHER_COLLECTION_ITEM (    
         id INT generated by default as identity primary key,
         collection varchar(256),
         embedding VECTOR,
@@ -350,15 +355,24 @@ def clear_db(self, collection: str = "deepsearcher", *args, **kwargs):
         status NUMBER DEFAULT 1,
         createtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
         updatetime TIMESTAMP DEFAULT NULL)""",
-    "has_table": "SELECT count(*) as rowcnt FROM all_tables WHERE table_name='DEEPSEARCHER'",
-    "has_collection": "select count(*) as rowcnt from DEEPSEARCHER where collection=:collection and status=1",
-    "list_collections": "select distinct collection from DEEPSEARCHER where status=1",
-    "drop_collection": "update DEEPSEARCHER set status=0 where collection=:collection and status=1",
-    "insert": """INSERT INTO DEEPSEARCHER (collection,embedding,text,reference,metadata) 
+}
+
+SQL_TEMPLATES = {
+    "has_table": f"""SELECT table_name FROM all_tables 
+        WHERE table_name in ({",".join([f"'{k}'" for k in TABLES.keys()])})""",
+    "has_collection": "select count(*) as rowcnt from DEEPSEARCHER_COLLECTION_INFO where collection=:collection and status=1",
+    "list_collections": "select collection,description from DEEPSEARCHER_COLLECTION_INFO where status=1",
+    "drop_collection": "update DEEPSEARCHER_COLLECTION_INFO set status=0 where collection=:collection and status=1",
+    "drop_collection_item": "update DEEPSEARCHER_COLLECTION_ITEM set status=0 where collection=:collection and status=1",
+    "insert_collection": """INSERT INTO DEEPSEARCHER_COLLECTION_INFO (collection,description) 
+        values (:collection,:description)""",
+    "insert": """INSERT INTO DEEPSEARCHER_COLLECTION_ITEM (collection,embedding,text,reference,metadata) 
         values (:collection,:embedding,:text,:reference,:metadata)""",
     "search": """SELECT * FROM 
         (SELECT t.*,
             VECTOR_DISTANCE(t.embedding,vector(:embedding_string,{dimension},{dtype}),COSINE) as distance
-        FROM DEEPSEARCHER t WHERE t.collection=:collection)
+        FROM DEEPSEARCHER_COLLECTION_ITEM t 
+        JOIN DEEPSEARCHER_COLLECTION_INFO c ON t.collection=c.collection 
+        WHERE t.collection=:collection AND t.status=1 AND c.status=1)
         WHERE distance<:max_distance ORDER BY distance ASC FETCH FIRST :top_k ROWS ONLY""",
 }