rainycode 1.0.0__py3-none-any.whl

core/databases/aiodb.py

@@ -0,0 +1,360 @@
+import datetime
+from typing import Dict, List, Optional, Any, Tuple
+from fastapi import FastAPI
+from fastapi.encoders import jsonable_encoder
+from tortoise import Tortoise
+from tortoise.models import Model
+from tortoise.contrib.fastapi import RegisterTortoise
+from tortoise.queryset import QuerySetSingle, QuerySet
+from contextlib import asynccontextmanager
+
+from core.base_config import base_config
+
+db_models = [
+    'aerich.models',
+    'models'
+]
+
+db_config = {
+    'connections': {
+        'main': base_config.mysql_url,  # 默认数据库连接
+    },
+    'apps': {
+        'main': {'models': db_models, 'default_connection': 'main'},
+    },
+    'use_tz': False,
+    'timezone': 'Asia/Shanghai'
+}
+
+class AioDb:
+
+    ### 事务的两种方式：
+    # async with in_transaction() as connection:  
+    #     pass
+    #
+    # @atomic('main')  （推荐使用）
+    # def func():
+    #     pass
+    #
+
+    @classmethod
+    async def connect(cls):
+        await Tortoise.init(config=db_config)
+
+    @classmethod
+    async def close(cls):
+        await Tortoise.close_connections()
+
+    @classmethod
+    async def register_db(cls, app: FastAPI) -> RegisterTortoise:
+        # 注册数据库，FastAPI用
+        tortoise = RegisterTortoise(
+            app,
+            config=db_config,
+            generate_schemas=False,  # 自动创建不存在的表（生产环境建议关闭，手动管理）
+            add_exception_handlers=True  # 添加 Tortoise 异常处理器
+        )
+        return tortoise
+        
+
+    @classmethod
+    async def init_tables(cls):
+        """初始化数据库表，如果表不存在则创建"""
+        try:
+            await Tortoise.generate_schemas()
+        except Exception as e:
+            raise e
+
+    @classmethod
+    @asynccontextmanager
+    async def execute(cls):
+        """
+        :example:
+            async with AioDb.execute():
+                await async_func()
+        """
+        await cls.connect()
+        yield
+        await cls.close()
+
+    @classmethod
+    def get_model_by_name(cls, name, app='main'):
+        return Tortoise.apps[app][name]
+
+    @classmethod
+    async def run_func(cls, func, *args, **kwargs):
+        async with AioDb.execute():
+            await func(*args, **kwargs)
+
+
+class MySQLUtils(Model):
+
+    async def model_dump(self, **kwargs):
+        return await MySQLUtils.model2dict(self, **kwargs)
+
+    @staticmethod
+    def json_encoder(
+            obj: Any,
+            include: Optional[List[str]] = None,
+            exclude: Optional[List[str]] = None,
+            **kwargs
+    ):
+        return jsonable_encoder(
+            obj,
+            include=set(include) if include else None,
+            exclude=set(exclude) if exclude else None,
+            custom_encoder={
+                datetime.datetime: lambda dt: datetime.datetime.strftime(dt, '%Y-%m-%d %H:%M:%S')
+            }
+        )
+
+    @classmethod
+    async def model2dict(
+            cls,
+            model: Model,
+            include: List[str] | None = None,
+            exclude: List[str] | None = None,
+            alias: List[Tuple[str, str]] | None = None,
+            fetch_related: List[str] | None = None,
+            flat: bool = False,  # 扁平化
+            **kwargs
+    ) -> Dict[str, Any]:
+        meta = getattr(cls, 'PydanticMeta', getattr(cls, 'Meta', None))
+        if meta:
+            _include = getattr(meta, 'include', None)
+            _exclude = getattr(meta, 'exclude', None)
+            if _include:
+                include = list(set(include + _include)) if include else _include
+            if _exclude:
+                exclude = list(set(exclude + _exclude)) if exclude else _exclude
+
+        if include and exclude:
+            for include_field in include:
+                if include_field in exclude:
+                    exclude.remove(include_field)
+
+        alias_dict = dict()
+        if alias:
+            for item in alias:
+                alias_dict[item[0]] = item[1]
+        
+        # 默认过滤掉deleted字段
+        if exclude:
+            exclude.append('deleted')
+        else:
+            exclude = ['deleted']
+
+        result = cls.json_encoder(
+            obj=model,
+            include=include,
+            exclude=exclude,
+            **kwargs
+        )
+
+        if alias_dict:
+            for attr in result.keys():
+                if attr in alias_dict:
+                    result[alias_dict[attr]] = result.pop(attr)
+
+        fix_data = getattr(model, 'fix_data', None)
+        if fix_data and callable(fix_data):
+            fix_data(result)
+
+        if fetch_related:
+            for fr in fetch_related:
+                fr_include = kwargs.get(f'fr_{fr}_include')
+                fr_exclude = kwargs.get(f'fr_{fr}_exclude')
+                fr_field = getattr(model, fr)
+                if not fr_field:
+                    result[fr] = None
+                    continue
+                fr_model = await fr_field
+
+                fr_meta = getattr(fr_model, 'PydanticMeta', None)
+                if fr_meta:
+                    _fr_include = getattr(fr_meta, 'include', None)
+                    _fr_exclude = getattr(fr_meta, 'exclude', None)
+                    if _fr_include:
+                        fr_include = list(set(fr_include + _fr_include)) if fr_include else _fr_include
+                    if _fr_exclude:
+                        fr_exclude = list(set(fr_exclude + _fr_exclude)) if fr_exclude else _fr_exclude
+
+                if fr_include and fr_exclude:
+                    for fr_include_field in fr_include:
+                        if fr_include_field in fr_exclude:
+                            fr_exclude.remove(fr_include_field)
+
+                fix_data = None
+                if isinstance(fr_model, list) and fr_model:
+                    fix_data = getattr(fr_model[0], 'fix_data', None)
+                elif isinstance(fr_model, Model):
+                    fix_data = getattr(fr_model, 'fix_data', None)
+
+                # 默认过滤掉deleted字段
+                if fr_exclude:
+                    fr_exclude.append('deleted')
+                else:
+                    fr_exclude = ['deleted']
+
+                fr_result = cls.json_encoder(
+                    obj=fr_model,
+                    include=fr_include,
+                    exclude=fr_exclude,
+                    **kwargs
+                )
+
+                if alias_dict:
+                    for attr in result.keys():
+                        key = f'{fr}.{attr}'
+                        if key in alias_dict:
+                            fr_result[alias_dict[key]] = fr_result.pop(attr)
+
+                if fix_data and callable(fix_data):
+                    if isinstance(fr_result, list):
+                        for frr in fr_result:
+                            fix_data(frr)
+                    elif isinstance(fr_result, dict):
+                        fix_data(fr_result)
+
+                if flat and isinstance(fr_result, dict):
+                    for k, v in fr_result.items():
+                        result[fr + '_' + k] = v
+                else:
+                    result[fr] = fr_result
+        return result
+
+    @classmethod
+    async def get_data(
+            cls,
+            *args,
+            include: List[str] | None = None,  # 包含哪些字段，默认全部
+            exclude: List[str] | None = None,  # 排除哪些字段，默认全部
+            alias: List[Tuple[str, str]] | None = None,  # 字段别名，如[('id', 'uid')]
+            fetch_related: List[str] | None = None,  # 对应关系表
+            flat: bool = False,  # 扁平化
+            deleted: bool | None = None,
+            queryset_model: QuerySetSingle | Model | None = None,
+            **kwargs
+    ) -> Dict[str, Any] | None:
+        if 'deleted' in cls._meta.fields and deleted is None:
+            kwargs['deleted'] = False
+
+        queryKwargs = {}
+        if fetch_related:
+            for k, v in kwargs.items():
+                if k.startswith('fr_'):
+                    continue
+                queryKwargs[k] = v
+        else:
+            queryKwargs = kwargs
+
+        qm: QuerySetSingle | Model = queryset_model if queryset_model else cls.get(*args, **queryKwargs)
+        if isinstance(qm, Model):
+            if fetch_related:
+                await qm.fetch_related(*fetch_related)
+            model = qm
+        else:
+            if fetch_related:
+                qm = qm.prefetch_related(*fetch_related)
+            model = await qm
+
+        if not model:
+            return None
+        result = await cls.model2dict(model, include, exclude, alias, fetch_related, flat, **kwargs)
+        return result
+
+    @classmethod
+    async def get_or_none_data(
+            cls,
+            *args,
+            include: List[str] | None = None,  # 包含哪些字段，默认全部
+            exclude: List[str] | None = None,  # 排除哪些字段，默认全部
+            alias: List[Tuple[str, str]] | None = None,  # 字段别名，如[('id', 'uid')]
+            fetch_related: List[str] | None = None,  # 对应关系表
+            flat: bool = False,  # 扁平化
+            deleted: bool | None = None,
+            queryset_model: QuerySetSingle | Model | None = None,
+            **kwargs
+    ) -> Dict[str, Any] | None:
+        if 'deleted' in cls._meta.fields and deleted is None:
+            kwargs['deleted'] = False
+
+        queryKwargs = {}
+        if fetch_related:
+            for k, v in kwargs.items():
+                if k.startswith('fr_'):
+                    continue
+                queryKwargs[k] = v
+        else:
+            queryKwargs = kwargs
+
+        qm: QuerySetSingle | Model = queryset_model if queryset_model else cls.get_or_none(*args, **kwargs)
+        if isinstance(qm, Model):
+            if fetch_related:
+                await qm.fetch_related(*fetch_related)
+            model = qm
+        else:
+            if fetch_related:
+                qm = qm.prefetch_related(*fetch_related)
+            model = await qm
+
+        if not model:
+            return None
+        result = await cls.model2dict(model, include, exclude, alias, fetch_related, flat, **kwargs)
+        return result
+
+    @classmethod
+    async def query_data(
+            cls,
+            *args,
+            page_no: int | None = 0,
+            page_size: int | None = 0,
+            nocount: bool = False,
+            include: List[str] | None = None,   # 包含哪些字段，默认全部
+            exclude: List[str] | None = None,   # 排除哪些字段，默认全部
+            alias: List[Tuple[str, str]] | None = None,   # 字段别名，如[('id', 'uid')]
+            fetch_related: List[str] | None = None,  # 对应关系表
+            flat: bool = False,  # 扁平化
+            deleted: bool | None = None,
+            queryset: QuerySet | None = None,
+            order_by: List[str] | None = None,
+            **kwargs
+    ) -> Dict[str, Any] | List[Dict[str, Any]]:
+        if 'deleted' in cls._meta.fields and deleted is None:
+            kwargs['deleted'] = False
+
+        queryKwargs = {}
+        if fetch_related:
+            for k, v in kwargs.items():
+                if k.startswith('fr_'):
+                    continue
+                queryKwargs[k] = v
+        else:
+            queryKwargs = kwargs
+
+        qs: QuerySet = queryset if queryset else cls.filter(*args, **kwargs)
+        result = {'items': [], 'count': 0}
+        if not nocount:
+            result['count'] = await qs.count()
+        else:
+            del result['count']
+
+        if order_by:
+            qs = qs.order_by(*order_by)
+        if page_no and page_size:
+            qs = qs.limit(page_size).offset((page_no - 1) * page_size)
+        if fetch_related:
+            qs = qs.prefetch_related(*fetch_related)
+
+        model_list = await qs
+        for model in model_list:
+            item = await cls.model2dict(model, include, exclude, alias, fetch_related, flat, **kwargs)
+            result['items'].append(item)
+
+        if nocount:
+            return result['items']
+            
+        return result
+
+    class Meta:
+        abstract = True

core/databases/aioredis.py

@@ -0,0 +1,279 @@
+import time
+import asyncio
+from contextlib import asynccontextmanager
+from typing import Optional, Union, Callable, Mapping
+from redis.asyncio import Redis
+from redis import asyncio as aioredis
+from redis.typing import KeyT, EncodableT, ExpiryT, ZScoreBoundT, AnyKeyT
+from common.logging import log_exc
+from core.base_config import base_config
+
+
+class AioRedis:
+    """异步 Redis 客户端封装"""
+
+    __redis_pool: Redis | None = None
+    __connection_lock: asyncio.Lock = asyncio.Lock()  # 确保连接池初始化线程安全
+
+    @classmethod
+    async def connect(cls) -> Redis:
+        """获取 Redis 连接（单例连接池）"""
+        if cls.__redis_pool:
+            return cls.__redis_pool
+
+        # 使用锁确保连接池初始化的线程安全
+        async with cls.__connection_lock:
+            if not cls.__redis_pool:  # 双重检查
+                try:
+                    pool = aioredis.ConnectionPool.from_url(
+                        base_config.redis_url,
+                        encoding='utf-8',
+                        decode_responses=True,
+                        retry_on_timeout=True  # 增加超时重试
+                    )
+                    cls.__redis_pool = Redis(connection_pool=pool)
+                    # 测试连接是否有效
+                    cls.__redis_pool.ping()
+                except Exception as e:
+                    log_exc(e)
+                    raise
+        return cls.__redis_pool
+
+    @classmethod
+    async def close(cls) -> None:
+        """关闭 Redis 连接池"""
+        if cls.__redis_pool:
+            # 安全关闭连接池
+            try:
+                await cls.__redis_pool.connection_pool.disconnect()
+            except Exception as e:
+                log_exc(e)
+            finally:
+                cls.__redis_pool = None
+
+    @classmethod
+    @asynccontextmanager
+    async def session(cls):
+        """上下文管理器，自动管理连接生命周期"""
+        try:
+            redis = await cls.connect()
+            yield redis
+        finally:
+            # 连接池模式下不需要每次关闭，由外部统一管理
+            pass
+
+    @classmethod
+    async def get(cls, name: KeyT) -> Optional[bytes]:
+        redis = await cls.connect()
+        result = await redis.get(name)
+        return result
+
+    @classmethod
+    async def set(
+            cls,
+            name: KeyT,
+            value: EncodableT,
+            ex: Union[ExpiryT, None] = None,  # 超时秒数
+            px: Union[ExpiryT, None] = None,  # 超时毫秒
+            nx: bool = False,
+            xx: bool = False
+    ) -> Optional[bool]:
+        redis = await cls.connect()
+        result = await redis.set(name, value, ex=ex, px=px, nx=nx, xx=xx)
+        return result
+
+    @classmethod
+    async def exists(
+        cls,
+        *names: KeyT
+    ) -> Optional[bool]:
+        redis = await cls.connect()
+        result = await redis.exists(*names)
+        return result
+
+    @classmethod
+    async def zadd(
+            cls,
+            name: KeyT,
+            mapping: Mapping[AnyKeyT, EncodableT],
+            nx: bool = False,
+            xx: bool = False,
+            ch: bool = False,
+            incr: bool = False,
+            gt: bool = False,
+            lt: bool = False
+    ) -> Union[int, float]:
+        redis = await cls.connect()
+        result = await redis.zadd(name, mapping, nx, xx, ch, incr, gt, lt)
+        return result
+
+    @classmethod
+    async def zrem(cls, name: KeyT, values) -> int:
+        redis = await cls.connect()
+        result = await redis.zrem(name, values)
+        return result
+
+    @classmethod
+    async def zscore(
+            cls,
+            name: KeyT,
+            value: EncodableT
+    ) -> Optional[float]:
+        redis = await cls.connect()
+        result = await redis.zscore(name, value)
+        return result
+
+    @classmethod
+    async def zcard(cls, name: KeyT) -> int:
+        redis = await cls.connect()
+        result = await redis.zcard(name)
+        return result
+
+    @classmethod
+    async def zrevrangebyscore(
+            cls,
+            name: KeyT,
+            max: ZScoreBoundT,
+            min: ZScoreBoundT,
+            start: Union[int, None] = None,
+            num: Union[int, None] = None,
+            withscores: bool = False,
+            score_cast_func: Union[type, Callable] = float,
+    ):
+        redis = await cls.connect()
+        result = await redis.zrevrangebyscore(name, max, min, start, num, withscores, score_cast_func)
+        return result
+
+    @classmethod
+    async def zremrangebyscore(
+            cls,
+            name: KeyT,
+            min: ZScoreBoundT,
+            max: ZScoreBoundT
+    ) -> int:
+        redis = await cls.connect()
+        result = await redis.zremrangebyscore(name, min, max)
+        return result
+
+    @classmethod
+    async def delete(cls, *names: KeyT) -> int:
+        redis = await cls.connect()
+        result = await redis.delete(*names)
+        return result
+
+    @classmethod
+    async def ttl(cls, name: KeyT) -> int:
+        redis = await cls.connect()
+        result = await redis.ttl(name)
+        return result
+
+    @classmethod
+    async def expire(cls, name: KeyT, expire: ExpiryT) -> int:
+        redis = await cls.connect()
+        return await redis.expire(name, expire)
+
+    @classmethod
+    async def incrby(cls, name: KeyT, amount: int = 1) -> int:
+        redis = await cls.connect()
+        return await redis.incrby(name, amount)
+
+    @classmethod
+    async def decrby(cls, name: KeyT, amount: int = 1) -> int:
+        redis = await cls.connect()
+        return await redis.decrby(name, amount)
+
+    # 分布式锁实现优化
+    @classmethod
+    @asynccontextmanager
+    async def lock(
+        cls,
+        name: str,
+        value: EncodableT = 1,
+        timeout: int = 3,  # 等待锁的时间（秒）
+        expire: int = 3    # 锁的有效期（秒）
+    ):
+        if isinstance(value, int):
+            value = str(value)
+        
+        lock_key = f"lock:{name}"
+        acquired = False
+        renewal_task: asyncio.Task | None = None
+        
+        try:
+            # 获取锁
+            timeout_at = time.time() + timeout
+            while time.time() < timeout_at:
+                # 使用 set 命令的 nx 和 ex 参数原子性获取锁
+                result = await cls.set(lock_key, value, nx=True, ex=expire)
+                if result:
+                    acquired = True
+                    break
+                await asyncio.sleep(0.1)  # 短暂等待后重试
+            
+            if not acquired:
+                raise Exception(f'获取分布式锁失败：{name}')
+            
+            # 启动锁续期任务（只在需要时启动）
+            async def renewal():
+                # 续期间隔设为锁有效期的 1/3，确保在锁过期前完成续期
+                renewal_interval = max(0.1, expire / 3)
+                while acquired:
+                    # 检查剩余时间，不足 1 秒时续期
+                    if await cls.ttl(lock_key) < 1:
+                        await cls.expire(lock_key, expire)
+                    await asyncio.sleep(renewal_interval)
+            
+            renewal_task = asyncio.create_task(renewal())
+            yield acquired
+        
+        except Exception as e:
+            log_exc(e)
+            raise
+        finally:
+            # 清理工作
+            acquired = False  # 停止续期任务
+            if renewal_task:
+                # 等待续期任务结束
+                renewal_task.cancel()
+                try:
+                    await renewal_task
+                except asyncio.CancelledError:
+                    pass
+            
+            # 只删除自己持有的锁
+            if await cls.get(lock_key) == value:
+                await cls.delete(lock_key)
+
+    @classmethod
+    async def pass_fix_window(cls, name: str, limit: int, expire: int) -> bool:
+        """固定窗口限流算法"""
+        key = f"ratelimit:fixed:{name}"
+        async with cls.lock(name):
+            # 使用 incr 原子操作替代 get + set，减少 Redis 调用
+            count = await cls.incrby(key, 1)
+            # 首次设置过期时间
+            if count == 1:
+                await cls.expire(key, expire)
+            # 将 count 转换为整数后再进行比较
+            return count <= limit
+
+
+    @classmethod
+    async def pass_slide_window(cls, name, limit, expire):
+        """限流：滑动窗口"""
+        key = f"slidelimit:fixed:{name}"
+        async with cls.lock(name):
+            count = await cls.zcard(key)
+            if not count or count < limit:
+                now_timestamp = int(time.time() * 1000)
+                old_timestamp = now_timestamp - expire * 1000
+                await cls.zadd(key, {str(now_timestamp): now_timestamp})
+                await cls.zremrangebyscore(key, 0, old_timestamp)
+                await cls.expire(key, expire)
+                return True
+            else:
+                now_timestamp = int(time.time() * 1000)
+                old_timestamp = now_timestamp - expire * 1000
+                await cls.zremrangebyscore(key, 0, old_timestamp)
+
+            return False

core/middleware/http_middleware.py

@@ -0,0 +1,28 @@
+import time
+from fastapi import Request
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.types import ASGIApp
+from starlette.responses import Response
+from common.logging import logger
+
+class RequestHttpMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: ASGIApp) -> None:
+        super().__init__(app)
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        start_time = time.time()
+        response = await call_next(request)
+        process_time = time.time() - start_time
+        response.headers['X-Process-Time'] = str(process_time)
+        self.write_request_log(request, response)
+        return response
+
+    @staticmethod
+    def write_request_log(request: Request, response: Response) -> None:
+        http_version = f"http/{request.scope['http_version']}"
+        process_time = response.headers["X-Process-Time"]
+        # 获取客户端IP和端口
+        client_host = request.client.host if request.client else "unknown"
+        client_port = request.client.port if request.client else "unknown"
+        content = f"{client_host}:{client_port} - {request.method} {request.url} {http_version} {response.status_code} {round(float(process_time), 2)}ms"
+        logger.info(content)