agentic-kit-common 0.0.8__py3-none-any.whl → 0.0.9__py3-none-any.whl

agentic_kit_common/orm/execution.py

@@ -0,0 +1,53 @@
+from typing import List, Union
+
+import sqlparse
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+
+def is_readonly_sql(sql: str) -> bool:
+    """
+    只允许：
+      SELECT / WITH / VALUES / EXPLAIN / DESCRIBE / SHOW
+    拒绝：
+      INSERT/UPDATE/DELETE/CREATE/ALTER/DROP/TRUNCATE/LOAD/REPLACE/LOCK/UNLOCK/GRANT/REVOKE/EXECUTE/CALL
+    """
+    sql_clean = sqlparse.format(sql.strip(), strip_comments=True)
+    tokens = [t.normalized for t in sqlparse.parse(sql_clean)[0].flatten()
+              if t.ttype in (sqlparse.tokens.Keyword, sqlparse.tokens.DML, sqlparse.tokens.DDL)]
+
+    forbidden = {"INSERT", "UPDATE", "DELETE", "CREATE", "ALTER", "DROP", "TRUNCATE",
+                 "LOAD", "REPLACE", "LOCK", "UNLOCK", "GRANT", "REVOKE", "EXEC", "CALL"}
+    allowed_root = {"SELECT", "WITH", "EXPLAIN", "DESCRIBE", "SHOW", "VALUES"}
+
+    root = tokens[0].upper() if tokens else ""
+    if root not in allowed_root:
+        return False
+    if any(k in forbidden for k in tokens):
+        return False
+    return True
+
+
+def session_sql_execute(db_session: Session, sql_text: Union[str, List], query_only: bool = True):
+    if isinstance(sql_text, str):
+        # 处理单条SQL语句
+        if query_only:
+            if not is_readonly_sql(sql_text):
+                return []
+        result = db_session.execute(text(f"{sql_text}"))
+        columns = list(result.keys())
+        rows = [dict(zip(columns, r)) for r in result.fetchall()]
+        return rows
+    elif isinstance(sql_text, list):
+        results = []
+        for sub_sql_text in sql_text:
+            if query_only:
+                if not is_readonly_sql(sub_sql_text):
+                    continue
+            result = db_session.execute(text(f"{sub_sql_text}"))
+            columns = list(result.keys())
+            rows = [dict(zip(columns, r)) for r in result.fetchall()]
+            results.append(rows)
+            return results
+    else:
+        return []

agentic_kit_common/orm/multi_session.py

@@ -0,0 +1,245 @@
+import logging
+from contextlib import contextmanager
+from threading import Lock
+from typing import Dict, Optional, Any, List
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.orm import sessionmaker, scoped_session
+
+from .schema import DatabaseEngineModel, _DEFAULT_CONFIG
+
+# 配置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class DatabaseEngineManager:
+    """数据库引擎管理器"""
+
+    def __init__(self, init_engines: Optional[List[DatabaseEngineModel]] = None):
+        self._engines: Dict[str, Any] = {}
+        self._sessions: Dict[str, Any] = {}
+        self._engine_configs: Dict[str, Dict] = {}
+        self._default_engine_name = ""
+        self._lock = Lock()
+
+        # 初始化默认配置
+        self._default_config = _DEFAULT_CONFIG.model_dump()
+
+        # 初始化引擎
+        if init_engines:
+            self._init_engines(init_engines=init_engines)
+
+    def _init_engines(self, init_engines: List[DatabaseEngineModel]):
+        for engine in init_engines:
+            self.add_engine(engine)
+
+    def _create_engine(self, database_uri: str, database_name: str, config: Optional[Dict] = None) -> Any:
+        """创建数据库引擎"""
+        if config is None:
+            config = self._default_config.copy()
+
+        # 构建完整的数据库URL
+        if database_uri.endswith('/'):
+            full_url = f"{database_uri}{database_name}"
+        else:
+            full_url = f"{database_uri}/{database_name}"
+
+        try:
+            engine = create_engine(
+                url=full_url,
+                echo=config.get("echo", False),
+                pool_size=config.get("pool_size", 10),
+                max_overflow=config.get("max_overflow", 5),
+                pool_pre_ping=config.get("pool_pre_ping", True),
+                pool_recycle=config.get("pool_recycle", 1800),
+            )
+            logger.info(f"成功创建引擎: {database_name} -> {database_uri}")
+            return engine
+        except Exception as e:
+            logger.error(f"创建引擎失败: {database_name}, 错误: {e}")
+            raise
+
+    def add_engine(self, engine_info: DatabaseEngineModel) -> bool:
+        """添加数据库引擎"""
+        with self._lock:
+            engine_name = engine_info.engine_name
+            database_uri = engine_info.database_uri
+            database_name = engine_info.database_name
+            config = engine_info.config.model_dump().copy()
+
+            if engine_name in self._engines:
+                logger.warning(f"引擎 '{engine_name}' 已存在")
+                return False
+
+            try:
+                # 创建引擎
+                engine = self._create_engine(database_uri, database_name, config)
+
+                # 创建会话工厂
+                session_factory = scoped_session(
+                    sessionmaker(
+                        bind=engine,
+                        expire_on_commit=False,
+                        autocommit=False,
+                        autoflush=False
+                    )
+                )
+
+                # 存储配置和实例
+                self._engines[engine_name] = engine
+                self._sessions[engine_name] = session_factory
+                self._engine_configs[engine_name] = {
+                    "database_uri": database_uri,
+                    "database_name": database_name,
+                    "config": config or self._default_config.copy()
+                }
+
+                logger.info(f"成功添加引擎: {engine_name}")
+                return True
+
+            except Exception as e:
+                logger.error(f"添加引擎失败: {engine_name}, 错误: {e}")
+                return False
+
+    def remove_engine(self, engine_name: str) -> bool:
+        """移除数据库引擎"""
+        with self._lock:
+            if engine_name not in self._engines:
+                logger.warning(f"引擎 '{engine_name}' 不存在")
+                return False
+
+            try:
+                # 关闭所有会话
+                if engine_name in self._sessions:
+                    self._sessions[engine_name].close_all()
+                    del self._sessions[engine_name]
+
+                # 处置引擎
+                if engine_name in self._engines:
+                    self._engines[engine_name].dispose()
+                    del self._engines[engine_name]
+
+                # 移除配置
+                if engine_name in self._engine_configs:
+                    del self._engine_configs[engine_name]
+
+                # 如果移除的是默认引擎
+                if engine_name == self._default_engine_name and self._engines:
+                    self._default_engine_name = ''
+
+                logger.info(f"成功移除引擎: {engine_name}")
+                return True
+
+            except Exception as e:
+                logger.error(f"移除引擎失败: {engine_name}, 错误: {e}")
+                return False
+
+    def get_engine(self, engine_name: Optional[str] = None) -> Any:
+        """获取数据库引擎"""
+        if engine_name is None:
+            engine_name = self._default_engine_name
+
+        if engine_name not in self._engines:
+            raise ValueError(f"数据库引擎 '{engine_name}' 不存在")
+
+        return self._engines[engine_name]
+
+    def get_session_factory(self, engine_name: Optional[str] = None) -> Any:
+        """获取会话工厂"""
+        if engine_name is None:
+            engine_name = self._default_engine_name
+
+        if engine_name not in self._sessions:
+            raise ValueError(f"数据库引擎 '{engine_name}' 不存在")
+
+        return self._sessions[engine_name]
+
+    @contextmanager
+    def get_db_session(self,
+                       engine_name: Optional[str] = None,
+                       auto_commit_by_exit: bool = False,
+                       auto_close: bool = True):
+        """获取数据库会话（上下文管理器）"""
+        if engine_name is None:
+            engine_name = self._default_engine_name
+
+        if engine_name not in self._sessions:
+            raise ValueError(f"数据库引擎 '{engine_name}' 不存在")
+
+        session_factory = self._sessions[engine_name]
+        session = session_factory()
+
+        try:
+            yield session
+            if auto_commit_by_exit:
+                session.commit()
+        except OperationalError as e:
+            logger.warning(f"数据库连接异常，尝试重新连接: {e}")
+            session.rollback()
+            # 重新创建会话
+            session.close()
+            session = session_factory()
+            yield session
+            if auto_commit_by_exit:
+                session.commit()
+        except Exception as e:
+            session.rollback()
+            logger.error(f"数据库操作异常: {e}")
+            raise e
+        finally:
+            if auto_close:
+                session.close()
+
+    def set_default_engine(self, engine_name: str) -> bool:
+        """设置默认引擎"""
+        with self._lock:
+            if engine_name not in self._engines:
+                logger.warning(f"引擎 '{engine_name}' 不存在")
+                return False
+
+            self._default_engine_name = engine_name
+            logger.info(f"设置默认引擎为: {engine_name}")
+            return True
+
+    def list_engines(self) -> List[str]:
+        """列出所有引擎名称"""
+        return list(self._engines.keys())
+
+    def get_engine_info(self, engine_name: str) -> Optional[Dict]:
+        """获取引擎信息"""
+        if engine_name not in self._engine_configs:
+            return None
+
+        info = self._engine_configs[engine_name].copy()
+        info["is_default"] = (engine_name == self._default_engine_name)
+        return info
+
+    def health_check(self, engine_name: Optional[str] = None) -> Dict[str, bool]:
+        """健康检查"""
+        results = {}
+
+        if engine_name:
+            engines_to_check = [engine_name] if engine_name in self._engines else []
+        else:
+            engines_to_check = list(self._engines.keys())
+
+        for name in engines_to_check:
+            try:
+                with self.get_db_session(name, auto_close=True) as session:
+                    session.execute(text("SELECT 1"))
+                results[name] = True
+                logger.debug(f"健康检查通过: {name}")
+            except Exception as e:
+                results[name] = False
+                logger.error(f"健康检查失败: {name}, 错误: {e}")
+
+        return results
+
+
+# 初始化数据库引擎管理器
+def initialize_database_engine_manager(init_engines: Optional[List[DatabaseEngineModel]] = None):
+    """初始化默认数据库"""
+    manager = DatabaseEngineManager(init_engines=init_engines)
+    return manager

agentic_kit_common/orm/schema.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel, Field
+
+
+class DatabaseEngineConfigModel(BaseModel):
+    echo: bool = Field(False, description="数据库的 SQL 语句 + 参数 打印到标准输出")
+    pool_pre_ping: bool = Field(True, description="每次从池里拿连接前，先发一句做“心跳”")
+    pool_size: int = Field(10, description="连接池里 长期保持的“永久”连接 数量")
+    max_overflow: int = Field(5, description="当 pool_size 用光后，最多还能再新建多少条“临时”连接")
+    pool_recycle: int = Field(1800, description="一条连接被 复用多久之后强制回收（关闭并新建）")
+
+
+_DEFAULT_CONFIG = DatabaseEngineConfigModel()
+
+
+class DatabaseEngineModel(BaseModel):
+    engine_name: str = Field(..., description="数据库引擎名字")
+    database_uri: str = Field(..., description="数据库uir地址")
+    database_name: str = Field(..., description="数据库名称")
+    config: DatabaseEngineConfigModel = Field(_DEFAULT_CONFIG, description="配置信息")

{agentic_kit_common-0.0.8.dist-info → agentic_kit_common-0.0.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentic-kit-common
-Version: 0.0.8
+Version: 0.0.9
 Summary: Common utilities and tools for agentic kit ecosystem
 Home-page: 
 Author: manson

{agentic_kit_common-0.0.8.dist-info → agentic_kit_common-0.0.9.dist-info}/RECORD

@@ -3,7 +3,10 @@ agentic_kit_common/minio/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZ
 agentic_kit_common/minio/minio_manager.py,sha256=PzoUWn0YqXTHx1UClbkwLkIUmv5O4aSDc7eVf08qOzs,21494
 agentic_kit_common/orm/__init__.py,sha256=wqY81g3P7FftFvLK5SaxFJzNNxrQwtmcb4RCXOSAZa8,71
 agentic_kit_common/orm/base.py,sha256=QIura_i2nIY2XeA3-KkO2loLNbEAoJK2qx0hu_8nhYU,2277
+agentic_kit_common/orm/execution.py,sha256=sCEcInFGFZ9ZiIs0eHKbQPtN4Z02v7vGdFMyVEWygcU,1951
 agentic_kit_common/orm/manager.py,sha256=lWgFk5fUu_9m6yN_fskWHSYGaW30ty--KZKj8AuvIh0,4852
+agentic_kit_common/orm/multi_session.py,sha256=CAEnOLl0I-r77JAknKQ2sERF40OkHD7BOrYVYoHuQV4,8814
+agentic_kit_common/orm/schema.py,sha256=ukdVP71NE_JO5_HOe_FApJjBAKUsDluoElVn5vNSGJQ,1023
 agentic_kit_common/orm/session.py,sha256=LX4ZUKJNXdcQ0KqRt5L0pJX-QG-tDyKXEDijcPUkGD0,2716
 agentic_kit_common/vector/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 agentic_kit_common/vector/embedding/__init__.py,sha256=mp--cfCDzTn5hNxzcIAU4m0g4F0d0rrZErFVWsGDvx4,40
@@ -21,8 +24,9 @@ test/config.py,sha256=VMcvfWnuWZKr_p0rC57i90jTq-XVsRuVzfEIe1C7Drg,3804
 test/settings.py,sha256=vU1dqUvGzshi6MG7JbGfW4jKfUsfAObgUfky-LMjUNs,344
 test/test_embedding.py,sha256=78M6t9tt6-oCfdQZfOartCIkUFZMvogcu5UKJh1lJb8,609
 test/test_minio.py,sha256=TOkX8A2pPkjrwAIH88xBZmFpDc1ZgTk1QS6mtubOZ-Y,2308
+test/test_orm.py,sha256=8fGCU7BWaD5sDbg0fgYN0Saf_hi7t-q8svHCHLQDceo,1620
 test/test_vector.py,sha256=4sJ7bXt5iRhc6j1Ig0bZ1j0huPU-clnsnT7O0-uFbkw,1311
-agentic_kit_common-0.0.8.dist-info/METADATA,sha256=mLi0procHOdh1DiCFCstlsD1RM-wAMWEKFpHTRMbLxo,7496
-agentic_kit_common-0.0.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-agentic_kit_common-0.0.8.dist-info/top_level.txt,sha256=nEKDlp84vqKSVWssGcxyuIsTqWLhMo45xqMs2GK4Dgg,24
-agentic_kit_common-0.0.8.dist-info/RECORD,,
+agentic_kit_common-0.0.9.dist-info/METADATA,sha256=2QyqNG2MQrLYyQ6TDuPXkyBnPdf5La0g4-rloQ4jAZY,7496
+agentic_kit_common-0.0.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+agentic_kit_common-0.0.9.dist-info/top_level.txt,sha256=nEKDlp84vqKSVWssGcxyuIsTqWLhMo45xqMs2GK4Dgg,24
+agentic_kit_common-0.0.9.dist-info/RECORD,,

test/test_orm.py

@@ -0,0 +1,47 @@
+import sys
+import unittest
+from pathlib import Path
+
+from agentic_kit_common.orm.execution import session_sql_execute
+from agentic_kit_common.orm.multi_session import initialize_database_engine_manager
+from agentic_kit_common.orm.schema import DatabaseEngineModel
+
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+
+class MyTestCase(unittest.TestCase):
+    def test_multi_session(self):
+        engine = DatabaseEngineModel(
+            engine_name='default',
+            database_name='czailab_llm',
+            database_uri='mysql+mysqlconnector://ailab_dev:Qwert!%40%234@45.120.102.236'
+        )
+        # manager = initialize_database_engine_manager(init_engines=[engine])
+
+        manager = initialize_database_engine_manager()
+        manager.add_engine(engine_info=engine)
+
+        print(manager.get_engine('default'))
+        print(manager.get_engine_info('default'))
+        print(manager.list_engines())
+
+        # session = manager.get_db_session(engine_name='default')
+        with manager.get_db_session(engine_name='default') as session:
+            print(session)
+            print(type(session))
+
+            result = session_sql_execute(db_session=session, sql_text='select * from tenant; select * from rerank_model')
+            print(result)
+            print(type(result))
+            for item in result:
+                print(item)
+                print(item.keys())
+                print(item.values())
+                print(item.items())
+                for _item in item.items():
+                    print(_item)
+
+
+if __name__ == '__main__':
+    unittest.main()