refactor: Refactor storage system (#937)

dbgpt/app/knowledge/document_db.py

@@ -2,19 +2,13 @@ from datetime import datetime
 
 from sqlalchemy import Column, String, DateTime, Integer, Text, func
 
-from dbgpt.storage.metadata import BaseDao
-from dbgpt.storage.metadata.meta_data import (
-    Base,
-    engine,
-    session,
-    META_DATA_DATABASE,
-)
+from dbgpt.storage.metadata import BaseDao, Model
 from dbgpt._private.config import Config
 
 CFG = Config()
 
 
-class KnowledgeDocumentEntity(Base):
+class KnowledgeDocumentEntity(Model):
     __tablename__ = "knowledge_document"
     __table_args__ = {
         "mysql_charset": "utf8mb4",
@@ -39,16 +33,8 @@ class KnowledgeDocumentEntity(Base):
 
 
 class KnowledgeDocumentDao(BaseDao):
-    def __init__(self):
-        super().__init__(
-            database=META_DATA_DATABASE,
-            orm_base=Base,
-            db_engine=engine,
-            session=session,
-        )
-
     def create_knowledge_document(self, document: KnowledgeDocumentEntity):
-        session = self.get_session()
+        session = self.get_raw_session()
         knowledge_document = KnowledgeDocumentEntity(
             doc_name=document.doc_name,
             doc_type=document.doc_type,
@@ -69,7 +55,7 @@ class KnowledgeDocumentDao(BaseDao):
         return doc_id
 
     def get_knowledge_documents(self, query, page=1, page_size=20):
-        session = self.get_session()
+        session = self.get_raw_session()
         print(f"current session:{session}")
         knowledge_documents = session.query(KnowledgeDocumentEntity)
         if query.id is not None:
@@ -104,7 +90,7 @@ class KnowledgeDocumentDao(BaseDao):
         return result
 
     def get_documents(self, query):
-        session = self.get_session()
+        session = self.get_raw_session()
         print(f"current session:{session}")
         knowledge_documents = session.query(KnowledgeDocumentEntity)
         if query.id is not None:
@@ -136,7 +122,7 @@ class KnowledgeDocumentDao(BaseDao):
         return result
 
     def get_knowledge_documents_count_bulk(self, space_names):
-        session = self.get_session()
+        session = self.get_raw_session()
         """
         Perform a batch query to count the number of documents for each knowledge space.
 
@@ -161,7 +147,7 @@ class KnowledgeDocumentDao(BaseDao):
         return docs_count
 
     def get_knowledge_documents_count(self, query):
-        session = self.get_session()
+        session = self.get_raw_session()
         knowledge_documents = session.query(func.count(KnowledgeDocumentEntity.id))
         if query.id is not None:
             knowledge_documents = knowledge_documents.filter(
@@ -188,14 +174,14 @@ class KnowledgeDocumentDao(BaseDao):
         return count
 
     def update_knowledge_document(self, document: KnowledgeDocumentEntity):
-        session = self.get_session()
+        session = self.get_raw_session()
         updated_space = session.merge(document)
         session.commit()
         return updated_space.id
 
     #
     def delete(self, query: KnowledgeDocumentEntity):
-        session = self.get_session()
+        session = self.get_raw_session()
         knowledge_documents = session.query(KnowledgeDocumentEntity)
         if query.id is not None:
             knowledge_documents = knowledge_documents.filter(