toms-fast 0.2.1__py3-none-any.whl

tomskit/sqlalchemy/sqlalchemy.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from sqlalchemy.orm.relationships import _RelationshipDeclared
+from typing import Any, Optional
+
+from sqlalchemy import CHAR as sa_CHAR
+from sqlalchemy import JSON as sa_JSON
+from sqlalchemy import BigInteger as sa_BigInteger
+from sqlalchemy import Boolean as sa_Boolean
+from sqlalchemy import Column as sa_Column
+from sqlalchemy import DateTime as sa_DateTime
+from sqlalchemy import Float as sa_Float
+from sqlalchemy import ForeignKey as sa_ForeignKey
+from sqlalchemy import Index as sa_Index
+from sqlalchemy import Integer as sa_Integer
+from sqlalchemy import LargeBinary as sa_LargeBinary
+from sqlalchemy import MetaData as sa_MetaData
+from sqlalchemy import Numeric as sa_Numeric
+from sqlalchemy import PickleType as sa_PickleType
+from sqlalchemy import PrimaryKeyConstraint as sa_PrimaryKeyConstraint
+from sqlalchemy import Sequence as sa_Sequence
+from sqlalchemy import String as sa_String
+from sqlalchemy import Text as sa_Text
+from sqlalchemy import UniqueConstraint as sa_UniqueConstraint
+from sqlalchemy import and_ as sa_and_
+from sqlalchemy import delete as sa_delete
+from sqlalchemy import func as sa_func
+from sqlalchemy import insert as sa_insert
+from sqlalchemy import select as sa_select
+from sqlalchemy import text as sa_text
+from sqlalchemy import update as sa_update
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
+from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.orm import relationship as sa_relationship
+from sqlalchemy.ext.asyncio import AsyncAttrs
+from sqlalchemy.sql import Select
+
+from tomskit.sqlalchemy.pagination import Pagination, SelectPagination
+
+__all__ = [
+    "SQLAlchemy",
+    "Pagination",
+    "SelectPagination"
+]
+# Define a naming convention for indexes and constraints in MySQL
+DB_INDEXES_NAMING_CONVENTION = {
+    "ix": "%(column_0_label)s_idx",
+    "uq": "%(table_name)s_%(column_0_name)s_key",
+    "ck": "%(table_name)s_%(constraint_name)s_check",
+    "fk": "%(table_name)s_%(column_0_name)s_fkey",
+    "pk": "%(table_name)s_pkey",
+}
+
+# Define metadata with the naming convention
+metadata = sa_MetaData(naming_convention=DB_INDEXES_NAMING_CONVENTION)
+
+# Base model class for all models
+# Base = declarative_base(metadata=metadata)
+
+
+# Define the SQLAlchemy class
+class SQLAlchemy(ABC):
+    # Define all the common SQLAlchemy constructs
+    class Model(AsyncAttrs, DeclarativeBase):
+        metadata = metadata
+
+    Column = sa_Column
+    CHAR = sa_CHAR
+    BigInteger = sa_BigInteger
+    Boolean = sa_Boolean
+    DateTime = sa_DateTime
+    Float = sa_Float
+    Integer = sa_Integer
+    JSON = sa_JSON
+    LargeBinary = sa_LargeBinary
+    Numeric = sa_Numeric
+    PickleType = sa_PickleType
+    Sequence = sa_Sequence
+    String = sa_String
+    Text = sa_Text
+    text = staticmethod(sa_text)
+    ForeignKey = sa_ForeignKey
+    Index = sa_Index
+    uuid = sa_CHAR(36)  # Define a UUID column
+    PrimaryKeyConstraint = sa_PrimaryKeyConstraint
+    UniqueConstraint = sa_UniqueConstraint
+    select = staticmethod(sa_select)
+    delete = staticmethod(sa_delete)
+    update = staticmethod(sa_update)
+    insert = staticmethod(sa_insert)
+    func = sa_func
+    relationship = staticmethod[..., _RelationshipDeclared[Any]](sa_relationship)
+    and_ = staticmethod(sa_and_)
+    def __init__(self) -> None:
+        self._engine: Optional[AsyncEngine] = None
+        self._SessionLocal: Optional[async_sessionmaker] = None  
+
+
+    @abstractmethod
+    async def paginate(self,
+        select: Select[Any], 
+        *, 
+        page: int | None = None,
+        per_page: int | None = None,
+        max_per_page: int | None = None,
+        error_out: bool = True,
+        count: bool = True
+    ) -> Pagination:
+        raise NotImplementedError
+
+    @property
+    def session(self) -> AsyncSession:
+        raise NotImplementedError
+    
+    @abstractmethod
+    def create_session(self) -> AsyncSession:
+        raise NotImplementedError
+    
+    @abstractmethod
+    async def close_session(self):
+        raise NotImplementedError
+    
+    @abstractmethod
+    def initialize_session_pool(self, db_url: str):
+        raise NotImplementedError
+    
+    @abstractmethod
+    async def close_session_pool(self):
+        raise NotImplementedError
+
+

tomskit/sqlalchemy/types.py

@@ -0,0 +1,32 @@
+import uuid
+
+from sqlalchemy import CHAR, TypeDecorator
+from sqlalchemy.dialects.postgresql import UUID
+
+
+class StringUUID(TypeDecorator):
+    impl = CHAR
+    cache_ok = True
+
+    def process_bind_param(self, value, dialect):
+        if value is None:
+            return value
+        elif dialect.name == "postgresql":
+            return str(value)
+        else:
+            return value.hex
+
+    def load_dialect_impl(self, dialect):
+        if dialect.name == "postgresql":
+            return dialect.type_descriptor(UUID())
+        else:
+            return dialect.type_descriptor(CHAR(36))
+
+    def process_result_value(self, value, dialect):
+        if value is None:
+            return value
+        return str(value)
+
+
+def uuid_generate_v4():
+    return uuid.uuid4().hex

tomskit/task/README.md

@@ -0,0 +1,67 @@
+# AsyncTaskManager 使用指南
+
+`AsyncTaskManager` 是一个基于 `asyncio.TaskGroup` 的简单异步任务管理器，支持批量和单次任务执行，并提供可选的数据库会话支持和丰富的日志功能。
+
+## 功能概述
+
+- **异步任务管理**：支持添加和并发执行多个异步任务。
+- **数据库会话支持**：可选的数据库会话管理，适用于需要数据库交互的任务。
+- **调试日志**：提供详细的任务执行日志，便于调试和性能分析。
+
+## 类定义
+
+### `AsyncTaskManager`
+
+#### 初始化
+
+```python
+def __init__(self, task_name: str = __name__, db: bool = False, debug: bool = False)
+```
+
+- `task_name`：任务管理器的名称，用于日志记录。
+- `db`：是否启用数据库会话。
+- `debug`：是否启用调试日志。
+
+#### 方法
+
+- `add_task(target: TaskTarget, args: tuple = (), kwargs: Optional[dict[str, Any]] = None) -> None`
+  - 添加一个异步任务。
+  - `target`：异步任务函数。
+  - `args`：传递给任务函数的位置参数。
+  - `kwargs`：传递给任务函数的关键字参数。
+
+- `add_tasks(targets: list[tuple[TaskTarget, tuple, Optional[dict[str, Any]]]]) -> None`
+  - 批量添加多个异步任务。
+
+- `async def run_all() -> None`
+  - 并发执行所有添加的任务，捕获并收集异常。
+
+- `async def run_task(target: TaskTarget, args: tuple = (), kwargs: Optional[dict[str, Any]] = None) -> Any`
+  - 快捷单任务调用，执行单个任务并返回结果。
+
+## 使用示例
+
+```python
+import asyncio
+
+async def sample_task(x):
+    await asyncio.sleep(1)
+    return x * 2
+
+async def main():
+    manager = AsyncTaskManager(debug=True)
+    manager.add_task(sample_task, args=(5,))
+    await manager.run_all()
+    print(manager.results)  # 输出: [10]
+
+asyncio.run(main())
+```
+
+## 日志
+
+- 启用 `debug` 模式后，任务的开始和结束时间将被记录。
+- 异常将被记录并可供后续分析。
+
+## 数据库支持
+
+- 如果 `db` 参数为 `True`，任务执行时将自动创建和关闭数据库会话。 

tomskit/task/__init__.py

@@ -0,0 +1,4 @@
+from .task_manager import AsyncTaskManager
+
+
+__all__ = ["AsyncTaskManager"]

tomskit/task/task_manager.py

@@ -0,0 +1,124 @@
+import asyncio
+import logging
+import time
+import uuid
+from typing import Any, Awaitable, Callable, Optional
+from tomskit.sqlalchemy.database import db as db_instance
+
+logger = logging.getLogger(__name__)
+
+# 异步任务目标类型
+TaskTarget = Callable[..., Awaitable[Any]]
+
+class AsyncTaskManager:
+    """
+    基于 asyncio.TaskGroup 的简单异步任务管理器，支持批量/单次任务执行，
+    可选 db 会话，丰富日志
+    """
+    def __init__(
+        self,
+        task_name: str = __name__,
+        db: bool = False,
+        debug: bool = False,
+    ):
+        self.task_name = task_name
+        self.db = db
+        self.debug = debug
+        # 存放任务生成器
+        self.tasks: list[Callable[[], Awaitable[Any]]] = []
+        # 存放结果和异常
+        self.results: list[Any] = []
+        self.exceptions: list[Exception] = []
+
+    def add_task(
+        self,
+        target: TaskTarget,
+        args: tuple = (),
+        kwargs: Optional[dict[str, Any]] = None,
+    ) -> None:
+        """
+        添加一个异步任务，参数同 target 函数签名
+        """
+        if kwargs is None:
+            kwargs = {}
+        func_name = getattr(target, '__name__', repr(target))
+        task_id = uuid.uuid4().hex[:8]
+
+        async def wrapper() -> Any:
+            session = None
+            try:
+                # 建会话
+                if self.db:
+                    session = db_instance.create_session()
+                    logger.debug(f"[{self.task_name}] db session connected, task {task_id}")
+
+                # 调试日志: 开始
+                if self.debug:
+                    logger.info(f"[{self.task_name}] Start {func_name}({task_id})")
+                    start = time.perf_counter()
+
+                # 执行目标协程
+                result = await target(*args, **kwargs)
+                self.results.append(result)
+
+                # 调试日志: 完成
+                if self.debug:
+                    elapsed = time.perf_counter() - start
+                    logger.info(f"[{self.task_name}] Finish {func_name}({task_id}) in {elapsed:.3f}s")
+
+            except Exception:
+                # 只日志，不立即收集，留给外层 TaskGroup 收集
+                logger.exception(f"[{self.task_name}] Error in {func_name}({task_id})")
+                raise
+
+            finally:
+                # 关闭会话
+                if self.db and session is not None:
+                    await db_instance.close_session(session)
+                    logger.debug(f"[{self.task_name}] db session closed, task {task_id}")
+
+        self.tasks.append(wrapper)
+
+    def add_tasks(
+        self,
+        targets: list[tuple[TaskTarget, tuple, Optional[dict[str, Any]]]]
+    ) -> None:
+        """
+        批量添加任务: [(func, args, kwargs), ...]
+        """
+        for target, args, kwargs in targets:
+            self.add_task(target, args=args, kwargs=kwargs)
+
+    async def run_all(self) -> None:
+        """
+        并发执行所有添加的任务，使用 asyncio.TaskGroup，
+        捕获异常并收集，不向外抛出。
+        """
+        self.results.clear()
+        self.exceptions.clear()
+        try:
+            async with asyncio.TaskGroup() as tg:
+                for wrapper in self.tasks:
+                    tg.create_task(wrapper())  # type: ignore
+        except* Exception as eg:
+            for e in eg.exceptions:
+                if not isinstance(e, asyncio.CancelledError):
+                    self.exceptions.append(e)
+
+    async def run_task(
+        self,
+        target: TaskTarget,
+        args: tuple = (),
+        kwargs: Optional[dict[str, Any]] = None,
+    ) -> Any:
+        """
+        快捷单任务调用
+        """
+        self.tasks.clear()
+        self.results.clear()
+        self.exceptions.clear()
+        self.add_task(target, args=args, kwargs=kwargs)
+        await self.run_all()
+        if self.exceptions:
+            raise self.exceptions[0]
+        return self.results[0]

tomskit/tools/README.md

@@ -0,0 +1,63 @@
+# FastAI Toolkit - Worker 模块
+
+该模块提供了与 Redis 交互的功能，用于管理和监控 `uvicorn.workers.UvicornWorker` 的进程信息。以下是模块中可用的函数及其用途：
+
+## 函数
+
+### `worker_register_to_redis(redis: Redis, hostname: str, pid: int)`
+
+- **描述**: 
+  - 在 `gunicorn` 启动时，将 `uvicorn.workers.UvicornWorker` 的进程信息注册到 Redis。
+  
+- **参数**:
+  - `redis`: Redis 客户端实例。
+  - `hostname`: 主机名。
+  - `pid`: 进程 ID。
+
+- **功能**:
+  - 将进程信息存储在 Redis 中，以便后续管理和监控。
+
+### `worker_delete_from_redis(redis: Redis, hostname: str, pid: int)`
+
+- **描述**: 
+  - 在 `gunicorn` 关闭时，从 Redis 中删除 `uvicorn.workers.UvicornWorker` 的进程信息。
+  
+- **参数**:
+  - `redis`: Redis 客户端实例。
+  - `hostname`: 主机名。
+  - `pid`: 进程 ID。
+
+- **功能**:
+  - 从 Redis 中移除进程信息，释放资源。
+
+### `async worker_update_to_redis(hostname: str, pid: int, update_info: dict)`
+
+- **描述**: 
+  - 在 `uvicorn.workers.UvicornWorker` 中，更新进程信息到 Redis。
+  
+- **参数**:
+  - `hostname`: 主机名。
+  - `pid`: 进程 ID。
+  - `update_info`: 包含更新信息的字典，例如请求计数、异常计数等。
+
+- **功能**:
+  - 更新 Redis 中的进程信息，保持数据的实时性。
+
+### `async get_all_worker_info_from_redis(hostname: str) -> dict`
+
+- **描述**: 
+  - 从 Redis 中获取所有的进程信息，并返回一个字典。
+  
+- **参数**:
+  - `hostname`: 主机名。
+
+- **返回**:
+  - 包含所有进程信息的字典。
+
+- **功能**:
+  - 提供对所有注册进程的全面监控。
+
+## 注意事项
+
+- 确保在使用这些函数前，已正确初始化 Redis 客户端。
+- 这些函数主要用于管理和监控 `uvicorn` 工作进程，适用于需要实时更新和获取进程状态的场景。 

tomskit/tools/__init__.py

@@ -0,0 +1,18 @@
+
+from tomskit.tools.config import GunicornSettings
+from tomskit.tools.warnings import enable_unawaited_warning
+from tomskit.tools.woker import (
+    worker_register_to_redis, 
+    worker_delete_from_redis, 
+    worker_update_to_redis, 
+    get_all_worker_info_from_redis )
+
+
+__all__ = [
+    "worker_register_to_redis", 
+    "worker_delete_from_redis", 
+    "worker_update_to_redis", 
+    "get_all_worker_info_from_redis",
+    "GunicornSettings",
+    "enable_unawaited_warning"
+]

tomskit/tools/config.py

@@ -0,0 +1,70 @@
+import os
+from typing import List, Optional
+from pydantic import Field, field_validator
+from pydantic_settings import BaseSettings
+
+class GunicornSettings(BaseSettings):
+    """
+    Configuration settings for Gunicorn
+    """
+
+    bind: str = Field(
+        default="0.0.0.0:5001",
+        description="Bind address and port to listen on"
+    )
+
+    pidfile: Optional[str] = Field(
+        default=None,
+        description="File to write the PID of the main process"
+    )
+
+    proc_name: Optional[str] = Field(
+        default=None,
+        description="Name of the process"
+    )
+
+    workers: int = Field(
+        default=0,
+        description="Number of workers to run (0 = auto by CPU cores)"
+    )
+
+    cpu_affinity: List[int] = Field(
+        default_factory=list,
+        description="List of CPU core IDs to bind workers to"
+    )
+
+    daemon: bool = Field(
+        default=False,
+        description="Run as a daemon"
+    )
+
+    @field_validator("workers", mode="before")
+    @classmethod
+    def default_workers(cls, v):
+        if not v or int(v) <= 0:
+            return os.cpu_count() or 1
+        return int(v)
+
+    @field_validator("cpu_affinity", mode="before")
+    @classmethod
+    def parse_cpu_affinity(cls, v):
+        if v is None or (isinstance(v, str) and not v.strip()):
+            return []
+        if isinstance(v, str):
+            try:
+                return [int(x.strip()) for x in v.split(",") if x.strip()]
+            except Exception as e:
+                raise ValueError(f"Invalid CPU_AFFINITY string: {v}") from e
+        if isinstance(v, list):
+            return v
+        return []
+
+    @field_validator("pidfile", mode="after")
+    @classmethod
+    def ensure_pidfile_dir_exists(cls, v: str):
+        if not v:
+            return None
+        dir_path = os.path.dirname(v)
+        if dir_path and not os.path.exists(dir_path):
+            os.makedirs(dir_path, exist_ok=True)
+        return v

tomskit/tools/warnings.py

@@ -0,0 +1,37 @@
+
+
+import sys
+import warnings
+import logging
+
+logger = logging.getLogger(__name__)
+
+def enable_unawaited_warning():
+    """
+    捕获所有未被 await 的协程产生的 RuntimeWarning，
+    并用 logger.critical 记录：包括文件名和行号。
+    """
+    # 原始的 showwarning，用于非 RuntimeWarning 恢复默认行为
+    _orig_showwarning = warnings.showwarning
+
+    def _custom_warning_handler(
+        message,
+        category,
+        filename,
+        lineno,
+        file=None,
+        line=None
+    ):
+        if issubclass(category, RuntimeWarning):
+            # 只处理 RuntimeWarning
+            logger.critical(f"(File: {filename}, Line: {lineno}) {message}")
+        else:
+            # 其他警告恢复原来的行为
+            _orig_showwarning(message, category, filename, lineno, file, line)
+
+    # 打开协程 origin tracking（需 Python 3.11+）
+    sys.set_coroutine_origin_tracking_depth(5)
+    # 始终发出 RuntimeWarning
+    warnings.simplefilter("always", RuntimeWarning)
+    # 覆写全局 showwarning
+    warnings.showwarning = _custom_warning_handler

tomskit/tools/woker.py

@@ -0,0 +1,81 @@
+
+from datetime import datetime
+from redis import Redis
+from tomskit.redis import redis_client
+
+def worker_register_to_redis(redis: Redis, hostname: str, pid: int):
+    """
+    gunicorn 启动时，把 uvicorn.workers.UvicornWorker 注册进程到 redis
+    """
+    process_info: dict = {
+        "pid": str(pid),
+        "start_at": int(datetime.now().timestamp()),
+        "uptime": 0,
+        "request_count": 0,
+        "exception_count": 0,
+        "server_error_count": 0,
+        "status": "starting",
+        "last_update_by" : "gunicorn",
+        "last_update_at": int(datetime.now().timestamp()),
+    }
+    redis.sadd(f"{hostname}:workers", str(pid))
+    redis.hset(f"{hostname}:worker:{str(pid)}", mapping=process_info)
+
+
+def worker_delete_from_redis(redis: Redis, hostname: str, pid: int):
+    """
+    gunicorn 关闭时，把 uvicorn.workers.UvicornWorker 从 redis 中删除
+    """
+    redis.srem(f"{hostname}:workers", str(pid))
+    redis.delete(f"{hostname}:worker:{str(pid)}")
+
+
+async def worker_update_to_redis(hostname: str, pid: int, update_info: dict):
+    """
+    在 uvicorn.workers.UvicornWorker 中，更新进程信息到 redis
+    update_info: dict = {
+        "request_count": int,
+        "exception_count": int,
+        "server_error_count": int,
+    }
+    """
+
+    redis = redis_client._client
+    if redis is None:
+        raise RuntimeError("Redis client is not initialized. Call initialize first.")
+
+    worker_key = f"{hostname}:worker:{str(pid)}"
+    
+    # 获取当前进程信息
+    current_info = await redis.hgetall(worker_key) # type: ignore
+    
+    # 更新进程信息
+    if current_info:
+        current_info["uptime"] = int(datetime.now().timestamp()) - int(current_info["start_at"])
+        current_info["last_update_at"] = int(datetime.now().timestamp())
+        current_info["status"] = "running"
+        current_info["last_update_by"] = "uvicorn.tomskit"
+        current_info.update(update_info)
+        # 更新到 redis
+        await redis.hset(worker_key, mapping=current_info) # type: ignore
+
+
+async def get_all_worker_info_from_redis(hostname: str) -> dict:
+    """
+    从 redis 中获取所有的进程信息，返回 dict
+    """
+    redis = redis_client._client
+    if redis is None:
+        raise RuntimeError("Redis client is not initialized. Call initialize first.")
+
+    # 获取所有的 worker pids
+    worker_pids = await redis.smembers(f"{hostname}:workers") # type: ignore
+
+    all_process_info = {}
+    for pid in worker_pids:
+        worker_key = f"{hostname}:worker:{pid}"
+        process_info = await redis.hgetall(worker_key) # type: ignore
+        if process_info:
+            all_process_info[pid] = {k: v for k, v in process_info.items()}
+
+    return all_process_info