refactor: Refactor datasource module (#1309)

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

View File

@@ -148,7 +148,7 @@ async def app_resources(
results = []
match type:
case ResourceType.DB.value:
dbs = CFG.LOCAL_DB_MANAGE.get_db_list()
dbs = CFG.local_db_manager.get_db_list()
results = [db["db_name"] for db in dbs]
if name:
results = [r for r in results if name in r]

View File

@@ -22,7 +22,7 @@ class DatasourceLoadClient(ResourceDbClient):
).create()
def get_data_type(self, resource: AgentResource) -> str:
conn = CFG.LOCAL_DB_MANAGE.get_connect(resource.value)
conn = CFG.local_db_manager.get_connector(resource.value)
return conn.db_type
async def a_get_schema_link(self, db: str, question: Optional[str] = None) -> str:
@@ -44,7 +44,7 @@ class DatasourceLoadClient(ResourceDbClient):
except Exception as e:
print("db summary find error!" + str(e))
if not table_infos:
conn = CFG.LOCAL_DB_MANAGE.get_connect(db)
conn = CFG.local_db_manager.get_connector(db)
table_infos = await blocking_func_to_async(
self._executor, conn.table_simple_info
)
@@ -52,13 +52,13 @@ class DatasourceLoadClient(ResourceDbClient):
return table_infos
async def a_query_to_df(self, db: str, sql: str):
conn = CFG.LOCAL_DB_MANAGE.get_connect(db)
conn = CFG.local_db_manager.get_connector(db)
return conn.run_to_df(sql)
async def a_query(self, db: str, sql: str):
conn = CFG.LOCAL_DB_MANAGE.get_connect(db)
conn = CFG.local_db_manager.get_connector(db)
return conn.query_ex(sql)
async def a_run_sql(self, db: str, sql: str):
conn = CFG.LOCAL_DB_MANAGE.get_connect(db)
conn = CFG.local_db_manager.get_connector(db)
return conn.run(sql)

View File

@@ -1,17 +1,11 @@
import json
import logging
from typing import Any, Dict, List, Optional, Tuple, Union
from typing import Any, List, Optional
from dbgpt._private.config import Config
from dbgpt.agent.plugin.generator import PluginPromptGenerator
from dbgpt.agent.resource.resource_api import AgentResource, ResourceType
from dbgpt.agent.resource.resource_api import AgentResource
from dbgpt.agent.resource.resource_knowledge_api import ResourceKnowledgeClient
from dbgpt.component import ComponentType
from dbgpt.rag.chunk import Chunk
from dbgpt.serve.agent.hub.controller import ModulePlugin
from dbgpt.core import Chunk
from dbgpt.serve.rag.retriever.knowledge_space import KnowledgeSpaceRetriever
from dbgpt.util.executor_utils import ExecutorFactory, blocking_func_to_async
from dbgpt.util.tracer import root_tracer, trace
CFG = Config()