sdpy-kit 0.1.0__py3-none-any.whl

kit/db/base_db_kit.py

@@ -0,0 +1,406 @@
+"""
+Copyright (c) 2021-2026 Clark Chang. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Project: sdpy
+Author: Clark Chang
+"""
+
+from __future__ import annotations
+
+# date: 2026-08-22
+"""数据库 Kit 共享基类（DRY 收敛）
+
+三个数据库 Kit（PgKit / MySqlKit / SQLiteKit）在重构前拥有约 400 行几乎
+逐行相同的代码（连接管理 / ORM / 原生 SQL / 事务 / 表操作），仅 DB_TYPE
+与 get_table_info 的 SQL 不同。本基类将全部共享实现收敛，子类只声明
+DB_TYPE 并覆写 get_table_info（SQLite 另含向量操作）。
+
+设计要点：
+- 所有公开方法保留原签名与语义，外部调用零影响。
+- 使用 PEP 695 方法级泛型 `T: SQLModel` 表达「模型类 → 实例」。
+- 数据参数用 `dict[str, object]`（比 `dict[str, Any]` 严格）。
+- 三层边界不变：本层方法自动提交（内部 session_scope）；CRUDBase 不提交；
+  transaction() 上下文管理器管理多操作原子性。
+- 类中存在公开方法 `list`，会遮蔽内置 `list` 类型；因此所有返回/参数中
+  的列表类型统一显式写作 `builtins.list[...]`，避免 mypy 将 `list[T]`
+  解析为方法名。
+"""
+
+import builtins
+import threading
+from abc import ABC, abstractmethod
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import cast
+
+from sqlalchemy import text
+from sqlmodel import Session, SQLModel
+
+from kit.config.log_config import get_default_logger
+from kit.db.crud import CRUDBase
+from kit.db.engine import (
+    close_engine,
+    create_tables,
+    get_engine,
+    list_tables,
+    table_exists,
+)
+from kit.db.session import session_scope
+
+logger = get_default_logger(__name__)
+
+# 字段数据：列名到值的映射（object 覆盖全部 SQLModel 支持类型）
+type FieldData = dict[str, object]
+
+# CRUD 实例缓存：CRUDBase 无状态（仅存 model），按 model 复用避免重复实例化
+# 值类型取宽基类 SQLModel，规避 PEP 695 类型参数上界约束
+_crud_cache: dict[type, CRUDBase[SQLModel, SQLModel, SQLModel]] = {}
+_crud_lock = threading.Lock()
+
+
+def _get_crud[T: SQLModel](model: type[T]) -> CRUDBase[T, T, T]:
+    """按模型获取（并缓存）CRUD 实例"""
+    crud = _crud_cache.get(model)
+    if crud is None:
+        with _crud_lock:
+            crud = _crud_cache.get(model)
+            if crud is None:
+                new_crud: CRUDBase[T, T, T] = CRUDBase(model)
+                crud = cast(CRUDBase[SQLModel, SQLModel, SQLModel], new_crud)
+                _crud_cache[model] = crud
+    return cast(CRUDBase[T, T, T], crud)
+
+
+class BaseDbKit(ABC):
+    """数据库工具类共享基类
+
+    子类需声明 `DB_TYPE` 并实现 `get_table_info`。所有方法均为类方法，
+    业务代码无需实例化。
+    """
+
+    # 数据库类型标识（postgresql / mysql / sqlite），由子类覆写
+    DB_TYPE: str
+
+    # ==================== 连接管理 ====================
+
+    @classmethod
+    def connect(cls) -> bool:
+        """初始化数据库连接（延迟创建引擎并验证连通性）
+
+        Returns:
+            True 表示连接成功
+        """
+        try:
+            engine = get_engine(cls.DB_TYPE)
+            with engine.connect() as conn:
+                conn.execute(text("SELECT 1"))
+            logger.info("✅ 数据库连接成功")
+            return True
+        except Exception as e:
+            logger.error(f"❌ 数据库连接失败: {e}")
+            return False
+
+    @classmethod
+    def close(cls) -> None:
+        """关闭数据库连接"""
+        close_engine(cls.DB_TYPE)
+        logger.info("🔌 数据库连接已关闭")
+
+    @classmethod
+    def is_connected(cls) -> bool:
+        """检查数据库是否可连接"""
+        try:
+            engine = get_engine(cls.DB_TYPE)
+            with engine.connect() as conn:
+                conn.execute(text("SELECT 1"))
+            return True
+        except Exception:
+            return False
+
+    # ==================== ORM 操作（自动管理 session） ====================
+
+    @classmethod
+    def create[T: SQLModel](cls, model: type[T], data: FieldData) -> T:
+        """创建单条记录
+
+        Args:
+            model: SQLModel 模型类
+            data: 字段名到值的映射
+
+        Returns:
+            创建后的模型实例（含主键）
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            crud = _get_crud(model)
+            obj = crud.create(session, data)
+            session.expunge(obj)  # 脱离 session，使返回值在 session 关闭后仍可用
+            return obj
+
+    @classmethod
+    def create_multi[T: SQLModel](
+        cls, model: type[T], data_list: builtins.list[FieldData]
+    ) -> builtins.list[T]:
+        """批量创建记录
+
+        Args:
+            model: SQLModel 模型类
+            data_list: 字段字典列表
+
+        Returns:
+            创建后的模型实例列表
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            crud = _get_crud(model)
+            objs = crud.create_multi(
+                session, cast(builtins.list[T | FieldData], data_list)
+            )
+            for obj in objs:
+                session.expunge(obj)
+            return objs
+
+    @classmethod
+    def get[T: SQLModel](cls, model: type[T], id: object) -> T | None:
+        """根据主键获取单条记录
+
+        Args:
+            model: SQLModel 模型类
+            id: 主键值
+
+        Returns:
+            模型实例，不存在时返回 None
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            obj = session.get(model, id)
+            if obj:
+                session.expunge(obj)
+            return obj
+
+    @classmethod
+    def get_by[T: SQLModel](
+        cls, model: type[T], field_name: str, field_value: object
+    ) -> T | None:
+        """根据字段值获取单条记录
+
+        Args:
+            model: SQLModel 模型类
+            field_name: 字段名
+            field_value: 字段值
+
+        Returns:
+            模型实例，不存在时返回 None
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            crud = _get_crud(model)
+            obj = crud.get_by_field(session, field_name, field_value)
+            if obj:
+                session.expunge(obj)
+            return obj
+
+    @classmethod
+    def list[T: SQLModel](
+        cls,
+        model: type[T],
+        *,
+        skip: int = 0,
+        limit: int = 100,
+        order_by: str | None = None,
+        desc: bool = False,
+    ) -> builtins.list[T]:
+        """获取多条记录（分页 + 排序）
+
+        Args:
+            model: SQLModel 模型类
+            skip: 跳过记录数
+            limit: 返回记录数上限
+            order_by: 排序字段名
+            desc: 是否降序排列
+
+        Returns:
+            模型实例列表
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            crud = _get_crud(model)
+            results = crud.get_multi(
+                session, skip=skip, limit=limit, order_by=order_by, desc=desc
+            )
+            for obj in results:
+                session.expunge(obj)
+            return results
+
+    @classmethod
+    def update[T: SQLModel](cls, model: type[T], id: object, data: FieldData) -> T:
+        """更新单条记录
+
+        Args:
+            model: SQLModel 模型类
+            id: 主键值
+            data: 需要更新的字段字典
+
+        Returns:
+            更新后的模型实例
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            crud = _get_crud(model)
+            obj = crud.get_or_raise(session, id)
+            crud.update(session, obj, data)
+            session.expunge(obj)
+            return obj
+
+    @classmethod
+    def delete[T: SQLModel](cls, model: type[T], id: object) -> bool:
+        """删除单条记录
+
+        Args:
+            model: SQLModel 模型类
+            id: 主键值
+
+        Returns:
+            True 表示删除成功，False 表示记录不存在
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            crud = _get_crud(model)
+            obj = crud.delete(session, id)
+            return obj is not None
+
+    @classmethod
+    def upsert[T: SQLModel](
+        cls, model: type[T], data: FieldData, conflict_fields: builtins.list[str]
+    ) -> T:
+        """插入或更新记录
+
+        根据 conflict_fields 判断记录是否已存在，存在则更新，不存在则插入。
+
+        Args:
+            model: SQLModel 模型类
+            data: 字段字典
+            conflict_fields: 用于判断冲突的字段列表
+
+        Returns:
+            创建或更新后的模型实例
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            crud = _get_crud(model)
+            obj = crud.upsert(session, data, conflict_fields)
+            session.expunge(obj)
+            return obj
+
+    @classmethod
+    def count(cls, model: type[SQLModel]) -> int:
+        """获取记录总数
+
+        Args:
+            model: SQLModel 模型类
+
+        Returns:
+            记录总数
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            crud = _get_crud(model)
+            return crud.count(session)
+
+    # ==================== 原生 SQL 操作 ====================
+
+    @classmethod
+    def query(
+        cls, sql: str, params: FieldData | None = None
+    ) -> builtins.list[dict[str, object]]:
+        """执行原生查询 SQL，返回字典列表
+
+        Args:
+            sql: SQL 查询语句（使用命名参数 :param）
+            params: 参数字典
+
+        Returns:
+            查询结果列表，每行为一个字典
+        """
+        with session_scope(cls.DB_TYPE, autocommit=False) as session:
+            result = session.execute(text(sql), params or {})
+            rows = result.fetchall()
+            columns = result.keys()
+            data = [dict(zip(columns, row)) for row in rows]
+            logger.info(f"🔍 执行查询 → 返回 {len(data)} 行")
+            return data
+
+    @classmethod
+    def execute(cls, sql: str, params: FieldData | None = None) -> int:
+        """执行原生更新 SQL（INSERT/UPDATE/DELETE）
+
+        Args:
+            sql: SQL 语句（使用命名参数 :param）
+            params: 参数字典
+
+        Returns:
+            受影响的行数
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            result = session.execute(text(sql), params or {})
+            affected = getattr(result, "rowcount", 0) or 0
+            logger.info(f"✅ 执行更新 → 影响 {affected} 行")
+            return affected
+
+    # ==================== 事务管理 ====================
+
+    @classmethod
+    @contextmanager
+    def transaction(cls) -> Generator[Session, None, None]:
+        """事务上下文管理器
+
+        正常退出时自动提交，异常时自动回滚。
+        适用于多个操作需要原子性的场景。
+
+        Yields:
+            SQLModel Session 实例
+
+        用法::
+
+            with PgKit.transaction() as session:
+                crud = CRUDBase(UserModel)
+                user = crud.create(session, {"name": "test"})
+        """
+        with session_scope(cls.DB_TYPE) as session:
+            yield session
+
+    # ==================== 表操作 ====================
+
+    @classmethod
+    def list_tables(cls) -> builtins.list[str]:
+        """列出数据库中的所有表"""
+        return list_tables(cls.DB_TYPE)
+
+    @classmethod
+    def table_exists(cls, table_name: str) -> bool:
+        """检查表是否存在"""
+        return table_exists(table_name, cls.DB_TYPE)
+
+    @classmethod
+    def create_tables(cls, models: builtins.list[type[SQLModel]]) -> None:
+        """创建数据库表
+
+        Args:
+            models: 需要创建表的 SQLModel 模型类列表
+        """
+        create_tables(models, cls.DB_TYPE)
+
+    @classmethod
+    @abstractmethod
+    def get_table_info(cls, table_name: str) -> builtins.list[dict[str, object]]:
+        """获取表的列信息
+
+        Args:
+            table_name: 表名
+
+        Returns:
+            列信息字典列表（column_name, data_type, is_nullable 等）
+        """