github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-07 20:10:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: Refactor datasource module (#1309)

Browse Source
This commit is contained in:
Fangyin Cheng
2024-03-18 18:06:40 +08:00
committed by GitHub
parent 84bedee306
commit 4970c9f813
108 changed files with 1194 additions and 1066 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
dbgpt/rag/summary/rdbms_db_summary.py
View File

@@ -1,11 +1,14 @@
"""Summary for rdbms database."""
from typing import List
from typing import TYPE_CHECKING, List, Optional
from dbgpt._private.config import Config
from dbgpt.datasource.rdbms.base import RDBMSDatabase
from dbgpt.datasource.rdbms.base import RDBMSConnector
from dbgpt.rag.summary.db_summary import DBSummary
if TYPE_CHECKING:
from dbgpt.datasource.manages import ConnectorManager
CFG = Config()
@@ -17,7 +20,9 @@ class RdbmsSummary(DBSummary):
column3(column3 comment) and index keys, and table comment is {table_comment})
"""
def __init__(self, name: str, type: str):
def __init__(
self, name: str, type: str, manager: Optional["ConnectorManager"] = None
):
"""Create a new RdbmsSummary."""
self.name = name
self.type = type
@@ -26,10 +31,11 @@ class RdbmsSummary(DBSummary):
# self.tables_info = []
# self.vector_tables_info = []
if not CFG.LOCAL_DB_MANAGE:
raise ValueError("Local db manage is not initialized.")
# TODO: Don't use the global variable.
self.db = CFG.LOCAL_DB_MANAGE.get_connect(name)
db_manager = manager or CFG.local_db_manager
if not db_manager:
raise ValueError("Local db manage is not initialized.")
self.db = db_manager.get_connector(name)
self.metadata = """user info :{users}, grant info:{grant}, charset:{charset},
collation:{collation}""".format(
@@ -58,12 +64,12 @@ class RdbmsSummary(DBSummary):
def _parse_db_summary(
conn: RDBMSDatabase, summary_template: str = "{table_name}({columns})"
conn: RDBMSConnector, summary_template: str = "{table_name}({columns})"
) -> List[str]:
"""Get db summary for database.
Args:
conn (RDBMSDatabase): database connection
conn (RDBMSConnector): database connection
summary_template (str): summary template
"""
tables = conn.get_table_names()
@@ -75,12 +81,12 @@ def _parse_db_summary(
def _parse_table_summary(
conn: RDBMSDatabase, summary_template: str, table_name: str
conn: RDBMSConnector, summary_template: str, table_name: str
) -> str:
"""Get table summary for table.
Args:
conn (RDBMSDatabase): database connection
conn (RDBMSConnector): database connection
summary_template (str): summary template
table_name (str): table name