107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
# !/usr/bin/env python3
|
|
# -*- coding:utf-8 -*-
|
|
|
|
# @Time : 2024/5/21 17:12
|
|
# @Author : fanen.lhy
|
|
# @Email : fanen.lhy@antgroup.com
|
|
# @FileName: sqldb_wrapper.py
|
|
|
|
import json
|
|
from typing import Optional, Sequence, Any, Dict
|
|
|
|
from langchain_community.utilities.sql_database import SQLDatabase
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from ..base.config.application_configer.application_config_manager import (
|
|
ApplicationConfigManager
|
|
)
|
|
from ..base.config.component_configer.configers.sqldb_wrapper_config import SQLDBWrapperConfiger
|
|
from ..base.component.component_base import ComponentBase
|
|
from ..base.component.component_enum import ComponentEnum
|
|
|
|
|
|
class SQLDBWrapper(ComponentBase):
|
|
"""A sql DB wrapper based on sqlalchemy and langchain sql database, """
|
|
|
|
# Basic attributes of the service class.
|
|
component_type: ComponentEnum = ComponentEnum.SQLDB_WRAPPER
|
|
name: Optional[str] = None
|
|
description: Optional[str] = None
|
|
__sql_database: Optional[SQLDatabase] = None
|
|
db_session: Optional[sessionmaker] = None
|
|
db_wrapper_configer: Optional[SQLDBWrapperConfiger] = None
|
|
|
|
class Config:
|
|
arbitrary_types_allowed = True
|
|
|
|
def get_instance_code(self) -> str:
|
|
"""Generate the full instance code from sql db wrapper name. """
|
|
app_cfg_manager: ApplicationConfigManager = ApplicationConfigManager()
|
|
appname = app_cfg_manager.app_configer.base_info_appname
|
|
return f"{appname}.{self.component_type.value.lower()}.{self.name}"
|
|
|
|
def initialize_by_component_configer(self,
|
|
db_wrapper_configer: SQLDBWrapperConfiger
|
|
) -> 'SQLDBWrapper':
|
|
"""Initialize the SQLDBWrapper by the ComponentConfiger object.
|
|
|
|
Args:
|
|
db_wrapper_configer(SQLDBWrapperConfiger): A configer contains service
|
|
basic info.
|
|
Returns:
|
|
SQLDBWrapper: A SQLDBWrapper instance.
|
|
"""
|
|
self.name = db_wrapper_configer.name
|
|
self.description = db_wrapper_configer.description
|
|
self.db_wrapper_configer = db_wrapper_configer
|
|
return self
|
|
|
|
def run(self, command: str) -> Sequence[Dict[str, Any]]:
|
|
"""
|
|
Execute given sql command and return a result sequence.
|
|
"""
|
|
return self.sql_database._execute(command=command)
|
|
|
|
|
|
def run_with_str_return(self, command: str) -> str:
|
|
"""
|
|
Execute given sql command and return a str result, intended to be used
|
|
as a part of llm input. If db wrapper's 'max_string_length' property is
|
|
not negative, result str will be truncated.
|
|
"""
|
|
return self.sql_database.run(command, fetch="all")
|
|
|
|
@property
|
|
def sql_database(self):
|
|
"""
|
|
Lazy init, to ensure that database engine can be init correctly in
|
|
separate processes like gunicorn.
|
|
"""
|
|
if not self.__sql_database:
|
|
self.db_wrapper_configer.engine_args.setdefault(
|
|
"json_serializer", lambda x: json.dumps(x, ensure_ascii=False)
|
|
)
|
|
self.db_wrapper_configer.sql_database_args.setdefault(
|
|
"sample_rows_in_table_info", 3
|
|
)
|
|
self.db_wrapper_configer.sql_database_args.setdefault(
|
|
"max_string_length", -1
|
|
)
|
|
self.__sql_database = SQLDatabase.from_uri(
|
|
engine_args=self.db_wrapper_configer.engine_args,
|
|
database_uri=self.db_wrapper_configer.db_uri,
|
|
**self.db_wrapper_configer.sql_database_args
|
|
)
|
|
return self.__sql_database
|
|
|
|
|
|
def get_session(self):
|
|
"""
|
|
Get a sqlalchemy session, used for operating with orm.
|
|
"""
|
|
if self.db_session:
|
|
return self.db_session
|
|
# Create database engine
|
|
self.db_session = sessionmaker(bind=self.sql_database._engine)
|
|
return self.db_session
|