sdpy-kit 0.1.0__py3-none-any.whl

kit/db/crud.py

@@ -0,0 +1,361 @@
+"""
+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
+"""
+
+# date: 2026-08-23
+"""通用 CRUD 操作基类
+
+核心设计原则：CRUD 方法内部不调用 commit/rollback，仅执行 add/flush 或
+Core DML（update/delete）。事务的提交/回滚交给外层 session_scope 或
+@transactional 装饰器统一管理。这保证了多个 CRUD 操作可以在同一事务中安全组合。
+
+性能要点：
+- upsert 使用各数据库原生 ``INSERT ... ON CONFLICT``（原子，单语句）。
+- update_multi / delete_multi 使用单条批量 SQL，不再逐对象载入修改。
+- exists 只查主键列，不载入整对象。
+"""
+
+from typing import Any, cast
+
+from sqlalchemy import Column, Table, delete, func, update
+from sqlalchemy.dialects import mysql as mysql_dialect
+from sqlalchemy.dialects import postgresql as pg_dialect
+from sqlalchemy.dialects import sqlite as sqlite_dialect
+from sqlmodel import Session, SQLModel, select
+from sqlmodel.sql._expression_select_cls import SelectOfScalar
+
+from kit.config.log_config import get_default_logger
+from kit.db.exceptions import DBError, NotFoundError
+
+logger = get_default_logger(__name__)
+
+
+class CRUDBase[
+    ModelType: SQLModel,
+    CreateSchemaType: SQLModel,
+    UpdateSchemaType: SQLModel,
+]:
+    """通用 CRUD 操作基类（不在内部提交事务）
+
+    所有写操作（create/update/delete）仅执行 session.add() + flush() 或 Core DML，
+    不调用 session.commit()，由外层 session_scope 统一管理事务边界。
+
+    用法::
+
+        user_crud = CRUDBase(User)
+
+        with session_scope() as session:
+            user = user_crud.create(session, UserCreate(name="test"))
+            # 退出 with 时统一提交
+    """
+
+    def __init__(self, model: type[ModelType]):
+        self.model = model
+
+    def _select_all(self) -> SelectOfScalar[ModelType]:
+        """返回 sqlmodel.select 的标量查询（SQLModel 类型桥接）"""
+        return cast(SelectOfScalar[ModelType], select(self.model))
+
+    def _id_col(self) -> Column[Any]:
+        """约定的主键列（避免 SQLModel 泛型无 id 字段的类型缺陷）"""
+        return cast(Column[Any], getattr(self.model, "id"))
+
+    def _table(self) -> Table:
+        """模型对应的 Table（避免 SQLModel 泛型无 __table__ 的类型缺陷）"""
+        return cast(Table, getattr(self.model, "__table__"))
+
+    # ==================== 查询方法 ====================
+
+    def get(self, session: Session, id: Any) -> ModelType | None:
+        """根据主键获取单条记录"""
+        obj = session.get(self.model, id)
+        logger.info(
+            f"🔍 查询 {self.model.__name__} id={id} → {'命中' if obj else '未找到'}"
+        )
+        return obj
+
+    def get_or_raise(self, session: Session, id: Any) -> ModelType:
+        """根据主键获取记录，不存在时抛出 NotFoundError"""
+        obj = session.get(self.model, id)
+        if obj is None:
+            logger.warning(
+                f"❌ {self.model.__name__} id={id} 不存在，抛出 NotFoundError"
+            )
+            raise NotFoundError(f"{self.model.__name__} id={id} 不存在")
+        logger.info(f"🔍 查询 {self.model.__name__} id={id} → 命中")
+        return obj
+
+    def get_multi(
+        self,
+        session: Session,
+        *,
+        skip: int = 0,
+        limit: int = 100,
+        order_by: str | None = None,
+        desc: bool = False,
+    ) -> list[ModelType]:
+        """获取多条记录（分页 + 排序）"""
+        statement = self._select_all()
+
+        if order_by:
+            column = getattr(self.model, order_by, None)
+            if column is None:
+                raise DBError(f"模型 {self.model.__name__} 不存在字段: {order_by}")
+            statement = statement.order_by(column.desc() if desc else column.asc())
+
+        statement = statement.offset(skip).limit(limit)
+        results = list(session.exec(statement).all())
+        logger.info(f"🔍 查询 {self.model.__name__} 多条记录 → 返回 {len(results)} 条")
+        return results
+
+    def get_by_field(
+        self, session: Session, field_name: str, field_value: Any
+    ) -> ModelType | None:
+        """根据字段值获取单条记录"""
+        column = getattr(self.model, field_name, None)
+        if column is None:
+            raise DBError(f"模型 {self.model.__name__} 不存在字段: {field_name}")
+        result = session.exec(
+            self._select_all().where(column == field_value)
+        ).first()
+        logger.info(
+            f"🔍 按字段查询 {self.model.__name__}({field_name}={field_value}) "
+            f"→ {'命中' if result else '未找到'}"
+        )
+        return result
+
+    def get_multi_by_field(
+        self,
+        session: Session,
+        field_name: str,
+        field_value: Any,
+        *,
+        skip: int = 0,
+        limit: int = 100,
+    ) -> list[ModelType]:
+        """根据字段值获取多条记录"""
+        column = getattr(self.model, field_name, None)
+        if column is None:
+            raise DBError(f"模型 {self.model.__name__} 不存在字段: {field_name}")
+        statement = (
+            self._select_all().where(column == field_value).offset(skip).limit(limit)
+        )
+        return list(session.exec(statement).all())
+
+    # ==================== 写入方法（不提交） ====================
+
+    def create(
+        self, session: Session, obj_in: CreateSchemaType | dict[str, Any]
+    ) -> ModelType:
+        """创建记录（仅 add + flush，不 commit）"""
+        data = (
+            obj_in
+            if isinstance(obj_in, dict)
+            else obj_in.model_dump(exclude_unset=True)
+        )
+        db_obj = self.model(**data)
+        session.add(db_obj)
+        session.flush()  # 刷新以获取主键，但不提交
+        logger.info(
+            f"✅ 创建 {self.model.__name__} 记录 id={getattr(db_obj, 'id', '?')}"
+        )
+        return db_obj
+
+    def create_multi(
+        self, session: Session, objs_in: list[CreateSchemaType | dict[str, Any]]
+    ) -> list[ModelType]:
+        """批量创建记录（仅 add_all + flush，不 commit）"""
+        db_objs = []
+        for obj_in in objs_in:
+            data = (
+                obj_in
+                if isinstance(obj_in, dict)
+                else obj_in.model_dump(exclude_unset=True)
+            )
+            db_objs.append(self.model(**data))
+        session.add_all(db_objs)
+        session.flush()
+        logger.info(f"✅ 批量创建 {len(db_objs)} 条 {self.model.__name__} 记录")
+        return db_objs
+
+    def update(
+        self,
+        session: Session,
+        db_obj: ModelType,
+        obj_in: UpdateSchemaType | dict[str, Any],
+    ) -> ModelType:
+        """更新记录（仅 add + flush，不 commit）"""
+        update_data = (
+            obj_in
+            if isinstance(obj_in, dict)
+            else obj_in.model_dump(exclude_unset=True)
+        )
+        for field, value in update_data.items():
+            setattr(db_obj, field, value)
+        session.add(db_obj)
+        session.flush()
+        logger.info(
+            f"✅ 更新 {self.model.__name__} id={getattr(db_obj, 'id', '?')} "
+            f"→ 字段: {update_data}"
+        )
+        return db_obj
+
+    def update_multi(
+        self,
+        session: Session,
+        ids: list[Any],
+        obj_in: UpdateSchemaType | dict[str, Any],
+    ) -> int:
+        """批量更新记录（单条 UPDATE SQL，不 commit）"""
+        update_data = (
+            obj_in
+            if isinstance(obj_in, dict)
+            else obj_in.model_dump(exclude_unset=True)
+        )
+        id_col = self._id_col()
+        statement = update(self.model).where(id_col.in_(ids)).values(**update_data)
+        result = session.execute(statement)
+        count = getattr(result, "rowcount", 0) or 0
+        logger.info(
+            f"✅ 批量更新 {count} 条 {self.model.__name__} 记录 "
+            f"ids={ids} → 字段: {update_data}"
+        )
+        return count
+
+    def delete(self, session: Session, id: Any) -> ModelType | None:
+        """删除记录（仅 delete，不 commit）"""
+        db_obj = session.get(self.model, id)
+        if db_obj:
+            session.delete(db_obj)
+            logger.info(f"✅ 删除 {self.model.__name__} id={id} → 已删除")
+        else:
+            logger.info(f"🔍 删除 {self.model.__name__} id={id} → 记录不存在")
+        return db_obj
+
+    def delete_multi(self, session: Session, ids: list[Any]) -> int:
+        """批量删除记录（单条 DELETE SQL，不 commit）"""
+        id_col = self._id_col()
+        statement = delete(self.model).where(id_col.in_(ids))
+        result = session.execute(statement)
+        count = getattr(result, "rowcount", 0) or 0
+        logger.info(
+            f"✅ 批量删除 {count} 条 {self.model.__name__} 记录 ids={ids}"
+        )
+        return count
+
+    # ==================== 聚合查询 ====================
+
+    def count(self, session: Session) -> int:
+        """获取记录总数"""
+        total = session.exec(select(func.count()).select_from(self.model)).one()
+        logger.info(f"🔍 统计 {self.model.__name__} 记录总数 → {total}")
+        return total
+
+    def exists(self, session: Session, id: Any) -> bool:
+        """检查记录是否存在（仅查询主键列，不载入整对象）"""
+        id_col = self._id_col()
+        found = (
+            session.exec(
+                cast(SelectOfScalar[object], select(id_col).where(id_col == id))
+            ).first()
+            is not None
+        )
+        logger.info(
+            f"🔍 检查 {self.model.__name__} id={id} 是否存在 → "
+            f"{'是' if found else '否'}"
+        )
+        return found
+
+    def upsert(
+        self,
+        session: Session,
+        obj_in: CreateSchemaType | dict[str, Any],
+        conflict_fields: list[str],
+    ) -> ModelType:
+        """插入或更新记录（原生 upsert，原子单语句）
+
+        根据 conflict_fields 对应的唯一约束判断冲突，存在则更新，不存在则插入。
+        注意：conflict_fields 必须对应表上的唯一索引/约束，否则数据库会报错。
+        不 commit，由外层事务管理。
+
+        Args:
+            session: 数据库会话
+            obj_in: 创建/更新数据
+            conflict_fields: 用于判断冲突的字段列表（需为唯一约束）
+        """
+        data = (
+            obj_in
+            if isinstance(obj_in, dict)
+            else obj_in.model_dump(exclude_unset=True)
+        )
+        if not conflict_fields:
+            raise DBError("upsert 必须提供 conflict_fields")
+
+        # 更新集：排除冲突字段与主键（避免覆盖主键/冲突键）
+        pk_cols = {c.name for c in self._table().primary_key.columns}
+        update_cols = {
+            k: v
+            for k, v in data.items()
+            if k not in conflict_fields and k not in pk_cols
+        }
+        if not update_cols:
+            # 无更新字段时退化为空更新（保持合法 SQL）
+            update_cols = {conflict_fields[0]: data[conflict_fields[0]]}
+
+        # 各数据库方言的 INSERT ... ON CONFLICT / ON DUPLICATE KEY（原子单语句）
+        dialect = session.get_bind().dialect.name
+        if dialect == "postgresql":
+            session.execute(
+                pg_dialect.insert(self.model)
+                .values(**data)
+                .on_conflict_do_update(
+                    index_elements=conflict_fields, set_=update_cols
+                )
+            )
+        elif dialect == "mysql":
+            session.execute(
+                mysql_dialect.insert(self.model)
+                .values(**data)
+                .on_duplicate_key_update(**update_cols)
+            )
+        elif dialect == "sqlite":
+            session.execute(
+                sqlite_dialect.insert(self.model)
+                .values(**data)
+                .on_conflict_do_update(
+                    index_elements=conflict_fields, set_=update_cols
+                )
+            )
+        else:
+            raise DBError(f"不支持的数据库类型: {dialect}")
+
+        # 按冲突字段回查返回实例（populate_existing 强制刷新，避免身份映射返回旧值）
+        conditions = [
+            getattr(self.model, field) == data.get(field) for field in conflict_fields
+        ]
+        re_stmt = self._select_all().where(*conditions).execution_options(
+            populate_existing=True
+        )
+        obj = session.exec(re_stmt).first()
+        if obj is None:
+            raise DBError(f"upsert {self.model.__name__} 后回查失败")
+        obj_id = getattr(obj, "id", "?")
+        logger.info(
+            f"🔄 upsert {self.model.__name__} 冲突字段 {conflict_fields} → id={obj_id}"
+        )
+        return obj

kit/db/engine.py

@@ -0,0 +1,230 @@
+"""
+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
+"""
+
+# date: 2025-04-19
+"""数据库引擎管理（线程安全单例）
+
+支持 PostgreSQL、MySQL、SQLite 三种数据库。
+使用 threading.Lock 保护 _engines 字典，避免 TOCTOU 竞态。
+"""
+
+import threading
+import urllib.parse
+from typing import Any, Type, cast
+
+from sqlalchemy import Engine, create_engine, event, inspect
+from sqlmodel import SQLModel
+
+from kit.config.kit_config import Settings, get_settings
+from kit.config.log_config import get_default_logger
+
+logger = get_default_logger(__name__)
+
+# 数据库类型常量
+POSTGRESQL = "postgresql"
+MYSQL = "mysql"
+SQLITE = "sqlite"
+
+# 线程安全的引擎缓存
+_engines: dict[str, Engine] = {}
+_lock = threading.Lock()
+
+
+def _build_pg_url(settings: Settings) -> str:
+    """构建 PostgreSQL 连接字符串"""
+    password = urllib.parse.quote_plus(settings.DB_PASSWORD)
+    return (
+        f"postgresql://{settings.DB_USER}:{password}"
+        f"@{settings.DB_HOST}:{settings.DB_PORT}"
+        f"/{settings.DB_DATABASE}"
+    )
+
+
+def _build_mysql_url(settings: Settings) -> str:
+    """构建 MySQL 连接字符串
+
+    MySQL 端口/用户/库名通过 MYSQL_* 覆盖，密码共用 DB_PASSWORD。
+    """
+    password = urllib.parse.quote_plus(settings.DB_PASSWORD)
+    port = getattr(settings, "MYSQL_PORT", settings.DB_PORT)
+    database = getattr(settings, "MYSQL_DATABASE", settings.DB_DATABASE)
+    user = getattr(settings, "MYSQL_USER", settings.DB_USER)
+    charset = getattr(settings, "MYSQL_CHARSET", "utf8mb4")
+    return (
+        f"mysql+pymysql://{user}:{password}"
+        f"@{settings.DB_HOST}:{port}"
+        f"/{database}?charset={charset}"
+    )
+
+
+def _build_sqlite_url(settings: Settings) -> str:
+    """构建 SQLite 连接字符串"""
+    db_path = getattr(settings, "SQLITE_DB_PATH", "./data/sqlite.db")
+    return f"sqlite:///{db_path}"
+
+
+def _resolve_db_type(db_type: str | None) -> str:
+    """解析数据库类型，回退到配置默认值"""
+    if db_type:
+        return db_type
+    return get_settings().DEFAULT_DB_TYPE
+
+
+def _receive_connect(dbapi_connection: Any, connection_record: Any) -> None:
+    """每次 SQLite 连接建立时自动加载 sqlite-vector 扩展"""
+    try:
+        import importlib.resources
+
+        dbapi_connection.enable_load_extension(True)
+        ext_path = (
+            importlib.resources.files("sqlite_vector.binaries") / "vector"
+        )
+        dbapi_connection.load_extension(str(ext_path))
+        dbapi_connection.enable_load_extension(False)
+        logger.debug("🔌 sqlite-vector 扩展已加载")
+    except Exception as e:
+        logger.warning(f"⚠️ 加载 sqlite-vector 扩展失败: {e}")
+
+
+def _register_vector_events(eng: Engine) -> None:
+    """注册 sqlite-vector 扩展加载事件"""
+    event.listen(eng, "connect", _receive_connect)
+
+
+def _create_engine(db_type: str, **kwargs: object) -> Engine:
+    """根据数据库类型创建引擎
+
+    SQLite 使用 check_same_thread=False 以支持多线程；
+    SQLite 内存模式（:memory:）额外使用 StaticPool 保持单连接共享。
+    MySQL / PostgreSQL 使用 QueuePool 连接池（pool_size=10, max_overflow=20）。
+    """
+    settings = get_settings()
+
+    if db_type == POSTGRESQL:
+        url = _build_pg_url(settings)
+    elif db_type == MYSQL:
+        url = _build_mysql_url(settings)
+    elif db_type == SQLITE:
+        url = _build_sqlite_url(settings)
+    else:
+        raise ValueError(f"不支持的数据库类型: {db_type}")
+
+    # 默认连接池参数（仅对支持连接池的数据库有效）
+    default_config: dict[str, object] = {
+        "pool_pre_ping": True,
+        "pool_recycle": 3600,
+        "echo": False,
+    }
+
+    if db_type == SQLITE:
+        # SQLite 使用线程安全的连接配置
+        default_config["connect_args"] = {"check_same_thread": False}
+        sqlite_path = getattr(settings, "SQLITE_DB_PATH", "./data/sqlite.db")
+        if ":memory:" in sqlite_path:
+            from sqlalchemy.pool import StaticPool
+
+            default_config["poolclass"] = StaticPool
+
+    else:
+        # MySQL / PostgreSQL 使用 QueuePool
+        default_config["pool_size"] = 10
+        default_config["max_overflow"] = 20
+
+    default_config.update(kwargs)
+
+    try:
+        # create_engine 接受 dict[str, Any] 风格 kwargs，此处为安全子集 cast
+        engine = create_engine(url, **cast(dict[str, Any], default_config))
+        if db_type == SQLITE:
+            _register_vector_events(engine)
+        logger.info(f"✅ {db_type} 数据库引擎已创建")
+        return engine
+    except Exception as e:
+        logger.error(f"❌ 创建 {db_type} 数据库引擎失败: {e}")
+        raise
+
+
+def get_engine(db_type: str | None = None, **kwargs: object) -> Engine:
+    """获取数据库引擎（线程安全单例）
+
+    Args:
+        db_type: 数据库类型（postgresql/mysql/sqlite），None 时使用配置默认值
+        **kwargs: 传递给 create_engine 的额外参数（仅首次创建时生效）
+
+    Returns:
+        SQLAlchemy Engine 实例
+    """
+    resolved = _resolve_db_type(db_type)
+
+    # 双重检查锁，避免已创建时仍获取锁
+    if resolved not in _engines:
+        with _lock:
+            if resolved not in _engines:
+                _engines[resolved] = _create_engine(resolved, **kwargs)
+
+    return _engines[resolved]
+
+
+def create_tables(models: list[Type[SQLModel]], db_type: str | None = None) -> None:
+    """创建数据库表（单次调用 metadata.create_all）
+
+    Args:
+        models: 需要创建表的 SQLModel 模型类列表（仅用于确认模型已注册）
+        db_type: 数据库类型，None 时使用配置默认值
+    """
+    engine = get_engine(db_type)
+    try:
+        # 所有 SQLModel 模型共享同一个 metadata，只需调用一次
+        SQLModel.metadata.create_all(engine)
+        logger.info(f"✅ 成功创建表（模型数: {len(models)}）")
+    except Exception as e:
+        logger.error(f"❌ 创建表失败: {e}")
+        raise
+
+
+def list_tables(db_type: str | None = None) -> list[str]:
+    """列出数据库中的所有表"""
+    engine = get_engine(db_type)
+    return inspect(engine).get_table_names()
+
+
+def table_exists(table_name: str, db_type: str | None = None) -> bool:
+    """检查表是否存在"""
+    return table_name in list_tables(db_type)
+
+
+def close_all_engines() -> None:
+    """关闭并清理所有数据库引擎"""
+    with _lock:
+        for db_type, engine in _engines.items():
+            try:
+                engine.dispose()
+                logger.info(f"🔌 {db_type} 数据库引擎已关闭")
+            except Exception as e:
+                logger.error(f"❌ 关闭 {db_type} 引擎失败: {e}")
+        _engines.clear()
+
+
+def close_engine(db_type: str) -> None:
+    """关闭指定类型的数据库引擎"""
+    with _lock:
+        engine = _engines.pop(db_type, None)
+        if engine:
+            engine.dispose()
+            logger.info(f"🔌 {db_type} 数据库引擎已关闭")

kit/db/exceptions.py

@@ -0,0 +1,34 @@
+"""
+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
+"""
+
+# date: 2026-08-23
+"""统一数据库异常体系
+
+所有异常继承 SQLAlchemyError，确保 @transactional 装饰器能统一捕获。
+"""
+
+from sqlalchemy.exc import SQLAlchemyError
+
+
+class DBError(SQLAlchemyError):
+    """所有数据库异常的基类"""
+
+
+class NotFoundError(DBError):
+    """记录不存在"""

kit/db/models/__init__.py

@@ -0,0 +1,33 @@
+"""
+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
+"""
+
+# date: 2026-08-23
+"""模型导出"""
+
+from kit.db.models.base import (
+    AuditMixin,
+    SoftDeleteMixin,
+    TimestampMixin,
+)
+
+__all__ = [
+    "TimestampMixin",
+    "SoftDeleteMixin",
+    "AuditMixin",
+]