fastapi-augment 0.1.0__py3-none-any.whl

fastapi_augment/db/sqlalchemy/crud_base.py

@@ -0,0 +1,426 @@
+"""
+@Author         : hangu
+@CreateDate     : 2026/8/31
+@Description    : Generic async CRUD base with create / read / update / delete operations.
+"""
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from math import ceil
+from typing import Any, Generic, TypeVar, cast
+
+from sqlalchemy import ColumnElement, delete, exists, func, select, update
+from sqlalchemy.engine import CursorResult
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import InstrumentedAttribute
+
+from .model_base import ModelBase
+
+ModelT = TypeVar('ModelT', bound=ModelBase)
+
+
+class CrudBase(Generic[ModelT]):
+    """Generic CRUD repository for :class:`ModelBase` subclasses.
+
+    The session is always passed explicitly, so callers keep full control
+    of the transaction boundary — CRUD methods only *flush*, never *commit*::
+
+        user_crud = CrudBase(User)
+
+        async with factory.transaction() as session:      # auto commit
+            await user_crud.create(session, User(name="alice"))
+
+        async with factory.read_session() as session:     # read replica
+            users = await user_crud.list(session, is_active=True, limit=10)
+
+    Filters accept both keyword arguments and raw SQLAlchemy expressions::
+
+        await user_crud.list(session, role="admin", expressions=(User.age > 18,))
+        await user_crud.list(session, id=["01A", "02B"])   # sequence -> IN
+        await user_crud.list(session, name=None)           # None -> IS NULL
+    """
+
+    def __init__(self, model: type[ModelT]):
+        """Initialize the CRUD repository.
+
+        Args:
+            model: The SQLAlchemy model class to operate on.
+        """
+        self.model = model
+
+    # ── Read ─────────────────────────────────────────────────────────────
+
+    async def get(self, session: AsyncSession, id_: str) -> ModelT | None:
+        """Fetch a single row by primary key.
+
+        Checks the session identity map first.
+
+        Args:
+            session: The async session to use.
+            id_: The primary key value.
+
+        Returns:
+            The model instance, or ``None`` if not found.
+        """
+        return await session.get(self.model, id_)
+
+    async def get_one(
+            self,
+            session: AsyncSession,
+            *,
+            expressions: Sequence[ColumnElement] | None = None,
+            **filters: Any,
+    ) -> ModelT | None:
+        """Fetch the first row matching the filters.
+
+        Args:
+            session: The async session to use.
+            expressions: Raw SQLAlchemy filter expressions.
+            **filters: Equality keyword filters.
+
+        Returns:
+            The first matching instance, or ``None`` if no row matches.
+        """
+        stmt = select(self.model).where(*self._conditions(expressions, filters)).limit(1)
+        result = await session.execute(stmt)
+        return result.scalars().first()
+
+    async def list(
+            self,
+            session: AsyncSession,
+            *,
+            expressions: Sequence[ColumnElement] | None = None,
+            order_by: Sequence[str | ColumnElement] | None = None,
+            offset: int = 0,
+            limit: int | None = 100,
+            **filters: Any,
+    ) -> Sequence[ModelT]:
+        """Fetch rows matching the filters with optional ordering / pagination.
+
+        Args:
+            session: Async session to run the query on (use a read session).
+            expressions: Raw SQLAlchemy filter expressions.
+            order_by: Column expressions or field names; prefix a name with
+                ``-`` for descending order (e.g. ``'-created_at'``).
+            offset: Number of rows to skip.
+            limit: Max rows to return (``None`` = no limit).
+            **filters: Equality keyword filters.
+
+        Returns:
+            A list of matching model instances.
+        """
+        stmt = select(self.model).where(*self._conditions(expressions, filters))
+
+        if order_by:
+            stmt = stmt.order_by(*(self._resolve_order(spec) for spec in order_by))
+
+        if offset:
+            stmt = stmt.offset(offset)
+
+        if limit is not None:
+            stmt = stmt.limit(limit)
+
+        result = await session.execute(stmt)
+        return result.scalars().all()
+
+    async def count(
+            self,
+            session: AsyncSession,
+            *,
+            expressions: Sequence[ColumnElement] | None = None,
+            **filters: Any,
+    ) -> int:
+        """Count rows matching the filters.
+
+        Args:
+            session: The async session to use.
+            expressions: Raw SQLAlchemy filter expressions.
+            **filters: Equality keyword filters.
+
+        Returns:
+            The number of matching rows.
+        """
+        stmt = select(func.count()).select_from(self.model).where(*self._conditions(expressions, filters))
+        result = await session.execute(stmt)
+        return result.scalar_one()
+
+    async def paginate(
+            self,
+            session: AsyncSession,
+            *,
+            page: int = 1,
+            size: int = 10,
+            expressions: Sequence[ColumnElement] | None = None,
+            order_by: Sequence[str | ColumnElement] | None = None,
+            **filters: Any,
+    ) -> dict[str, Any]:
+        """分页查询，自动执行 count + list 并返回分页结果字典。
+
+        内部复用 ``_conditions`` 保证 count 与 list 使用完全相同的过滤条件，
+        避免调用方手动写两遍 filter::
+
+            result = await crud.paginate(
+                session, page=1, size=10,
+                is_active=True, order_by=['-created_at'],
+            )
+            # result = {'items': [...], 'page': 1, 'size': 10, 'total': 100, 'pages': 10}
+
+            # 可配合 PageData 使用
+            from fastapi_augment.schemas import PageData
+            page_data = PageData.build(result['items'], page=result['page'],
+                                         size=result['size'], total=result['total'])
+
+        Args:
+            session: The async session to use (typically a read session).
+            page: 当前页码（从 1 开始）
+            size: 每页数量
+            expressions: Raw SQLAlchemy filter expressions.
+            order_by: Column expressions or field names; prefix ``-`` for descending.
+            **filters: Equality keyword filters.
+
+        Returns:
+            包含 items / page / size / total / pages 的字典。
+        """
+
+        conditions = self._conditions(expressions, filters)
+
+        # count
+        count_stmt = select(func.count()).select_from(self.model).where(*conditions)
+        total = (await session.execute(count_stmt)).scalar_one()
+
+        # list
+        offset = (page - 1) * size
+        list_stmt = select(self.model).where(*conditions)
+
+        if order_by:
+            list_stmt = list_stmt.order_by(*(self._resolve_order(spec) for spec in order_by))
+
+        list_stmt = list_stmt.offset(offset).limit(size)
+        items = (await session.execute(list_stmt)).scalars().all()
+
+        pages = ceil(total / size) if size > 0 else 0
+        return {
+            'items': items,
+            'page': page,
+            'size': size,
+            'total': total,
+            'pages': pages,
+        }
+
+    async def exists(
+            self,
+            session: AsyncSession,
+            *,
+            expressions: Sequence[ColumnElement] | None = None,
+            **filters: Any,
+    ) -> bool:
+        """Check whether at least one row matches the filters.
+
+        Args:
+            session: The async session to use.
+            expressions: Raw SQLAlchemy filter expressions.
+            **filters: Equality keyword filters.
+
+        Returns:
+            ``True`` if at least one matching row exists, ``False`` otherwise.
+        """
+        stmt = select(exists().where(*self._conditions(expressions, filters)))
+        result = await session.execute(stmt)
+        return bool(result.scalar())
+
+    # ── Update ───────────────────────────────────────────────────────────
+
+    async def update(self, session: AsyncSession, obj: ModelT, **values: Any) -> ModelT:
+        """Update an ORM instance in place.
+
+        Flushes; never commits.
+
+        Args:
+            session: The async session to use.
+            obj: The model instance to update.
+            **values: Field names and their new values.
+
+        Returns:
+            The updated model instance.
+
+        Raises:
+            AttributeError: If a key does not correspond to a mapped attribute.
+        """
+        for key in values:
+            self._attr(key)
+
+        for key, value in values.items():
+            setattr(obj, key, value)
+
+        await session.flush()
+        return obj
+
+    async def update_by_id(self, session: AsyncSession, id_: str, **values: Any) -> int:
+        """Update a row by primary key with a single UPDATE statement.
+
+        Args:
+            session: The async session to use.
+            id_: The primary key value.
+            **values: Field names and their new values.
+
+        Returns:
+            The number of affected rows (0 = not found or nothing to update).
+        """
+        if not values:
+            return 0
+
+        stmt = update(self.model).where(self.model.id == id_).values(**values)
+        result = cast(CursorResult[Any], await session.execute(stmt))
+        return result.rowcount or 0
+
+    # ── Delete ───────────────────────────────────────────────────────────
+
+    async def delete_by_id(self, session: AsyncSession, id_: str) -> bool:
+        """Delete a row by primary key.
+
+        Args:
+            session: The async session to use.
+            id_: The primary key value.
+
+        Returns:
+            ``True`` if a row was deleted, ``False`` otherwise.
+        """
+        stmt = delete(self.model).where(self.model.id == id_)
+        result = cast(CursorResult[Any], await session.execute(stmt))
+        return bool(result.rowcount)
+
+    async def delete_where(
+            self,
+            session: AsyncSession,
+            *,
+            expressions: Sequence[ColumnElement] | None = None,
+            **filters: Any,
+    ) -> int:
+        """Delete all rows matching the filters.
+
+        Args:
+            session: The async session to use.
+            expressions: Raw SQLAlchemy filter expressions.
+            **filters: Equality keyword filters.
+
+        Returns:
+            The number of deleted rows.
+        """
+        stmt = delete(self.model).where(*self._conditions(expressions, filters))
+        result = cast(CursorResult[Any], await session.execute(stmt))
+        return result.rowcount or 0
+
+    # ── Helpers ──────────────────────────────────────────────────────────
+
+    def _attr(self, name: str) -> InstrumentedAttribute[Any]:
+        """Resolve a field name to a mapped attribute.
+
+        Args:
+            name: The field name to resolve.
+
+        Returns:
+            The corresponding :class:`~sqlalchemy.orm.InstrumentedAttribute`.
+
+        Raises:
+            AttributeError: If the name does not correspond to a mapped attribute.
+        """
+        attr = getattr(self.model, name, None)
+
+        if not isinstance(attr, InstrumentedAttribute):
+            raise AttributeError(f'{self.model.__name__} has no mapped attribute {name!r}')
+
+        return attr
+
+    def _conditions(
+            self,
+            expressions: Sequence[ColumnElement] | None,
+            filters: Mapping[str, Any],
+    ) -> Sequence[ColumnElement]:
+        """Combine raw expressions and keyword filters into WHERE conditions.
+
+        Args:
+            expressions: Raw SQLAlchemy filter expressions.
+            filters: Equality keyword filters. Sequences (list, set, tuple,
+                frozenset) are converted to ``IN`` clauses; ``None`` values
+                become ``IS NULL`` checks.
+
+        Returns:
+            A list of column expressions suitable for ``.where()``.
+        """
+        conditions: list[ColumnElement] = list(expressions or [])
+
+        for key, value in filters.items():
+            column = self._attr(key)
+
+            if isinstance(value, (list, set, tuple, frozenset)):
+                conditions.append(column.in_(value))
+            else:
+                conditions.append(column == value)
+
+        return conditions
+
+    def _resolve_order(self, spec: str | ColumnElement) -> ColumnElement:
+        """Convert an order-by spec into a column expression.
+
+        Args:
+            spec: A column expression, or a field name. Prefix with ``-``
+                for descending order (e.g. ``'-created_at'``).
+
+        Returns:
+            A column expression with the appropriate asc/desc direction.
+        """
+        if isinstance(spec, str):
+            descending = spec.startswith('-')
+            column = self._attr(spec.lstrip('+-'))
+            return column.desc() if descending else column.asc()
+
+        return spec
+
+    # ── Static Methods ───────────────────────────────────────────────────
+
+    @staticmethod
+    async def create(session: AsyncSession, obj: ModelT) -> ModelT:
+        """Persist a new object.
+
+        Flushes so PKs / defaults are populated; never commits.
+
+        Args:
+            session: The async session to use.
+            obj: The model instance to persist.
+
+        Returns:
+            The same instance with populated defaults.
+        """
+        session.add(obj)
+        await session.flush()
+        return obj
+
+    @staticmethod
+    async def create_many(session: AsyncSession, objs: Sequence[ModelT]) -> Sequence[ModelT]:
+        """Persist multiple objects in one batch.
+
+        Flushes; never commits.
+
+        Args:
+            session: The async session to use.
+            objs: The model instances to persist.
+
+        Returns:
+            A list of the persisted instances.
+        """
+        session.add_all(objs)
+        await session.flush()
+        return objs
+
+    @staticmethod
+    async def delete(session: AsyncSession, obj: ModelT) -> None:
+        """Delete an ORM instance.
+
+        Flushes; never commits.
+
+        Args:
+            session: The async session to use.
+            obj: The model instance to delete.
+        """
+        await session.delete(obj)
+        await session.flush()

fastapi_augment/db/sqlalchemy/engine.py

@@ -0,0 +1,238 @@
+"""
+@Author         : hangu
+@CreateDate     : 2026/8/31
+@Description    : SQLAlchemy async engine factory, supporting single / master-replica / cluster topologies.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from itertools import cycle
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+
+
+# ── Configuration Models ─────────────────────────────────────────────────────
+
+
+@dataclass
+class NodeConfig:
+    """Single database node configuration.
+
+    Args:
+        url: Async SQLAlchemy connection URL, e.g. ``postgresql+asyncpg://user:pass@host/db``.
+        pool_size: Connection pool size (0 = unlimited).
+        max_overflow: Max connections allowed beyond *pool_size*.
+        pool_timeout: Seconds to wait for a connection from the pool.
+        pool_recycle: Recycle connections after N seconds (-1 = disable).
+        pool_pre_ping: Emit a test statement on checkout to verify liveness.
+        echo: Log all SQL statements.
+        connect_args: Extra arguments passed to the DBAPI ``connect()``.
+    """
+
+    url: str
+    pool_size: int = 5
+    max_overflow: int = 10
+    pool_timeout: float = 30.0
+    pool_recycle: int = 3600
+    pool_pre_ping: bool = True
+    echo: bool = False
+    connect_args: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not self.url:
+            raise ValueError('url must not be empty')
+        if self.pool_size < 0:
+            raise ValueError(f'pool_size must be >= 0, got {self.pool_size}')
+        if self.max_overflow < 0:
+            raise ValueError(f'max_overflow must be >= 0, got {self.max_overflow}')
+        if self.pool_timeout < 0:
+            raise ValueError(f'pool_timeout must be >= 0, got {self.pool_timeout}')
+        if self.pool_recycle < -1:
+            raise ValueError(f'pool_recycle must be >= -1, got {self.pool_recycle}')
+
+
+@dataclass
+class ClusterTopology:
+    """
+    Database cluster topology.
+
+    Supports three deployment patterns::
+
+        Single        : only ``primary`` is set
+        Master-Replica: ``primary`` + ``replicas``
+        Cluster       : ``primary`` + ``replicas`` + ``readonly``
+
+    Read routing priority: replicas -> readonly -> primary (fallback).
+    """
+
+    primary: NodeConfig
+    replicas: list[NodeConfig] = field(default_factory=list)
+    readonly: list[NodeConfig] = field(default_factory=list)
+
+    @property
+    def is_single(self) -> bool:
+        """Indicates whether the topology consists of a single primary node (no replicas or readonly nodes).
+
+        Returns:
+            True if the topology is single, False otherwise.
+        """
+        return not self.replicas and not self.readonly
+
+    @property
+    def is_master_replica(self) -> bool:
+        """Indicates whether the topology consists of a single primary node and one or more replicas (no readonly nodes).
+
+        Returns:
+            True if the topology is master-replica, False otherwise.
+        """
+        return bool(self.replicas) and not self.readonly
+
+    @property
+    def is_cluster(self) -> bool:
+        """Indicates whether the topology consists of a single primary node, one or more replicas, and one or more readonly nodes.
+
+        Returns:
+            True if the topology is a cluster, False otherwise.
+        """
+        return bool(self.replicas) or bool(self.readonly)
+
+    def get_all_read_sources(self) -> list[NodeConfig]:
+        """Return all nodes eligible for read traffic (replicas first, then readonly).
+
+        Returns:
+            A list of node configurations.
+        """
+        return [*self.replicas, *self.readonly]
+
+
+# ── Engine Manager ───────────────────────────────────────────────────────────
+
+
+class EngineManager:
+    """
+    Manages the full lifecycle of async SQLAlchemy engines.
+
+    Usage::
+
+        topology = ClusterTopology(
+            primary=NodeConfig(url='postgresql+asyncpg://...'),
+            replicas=[NodeConfig(url='postgresql+asyncpg://replica-1/...')],
+        )
+        manager = EngineManager(topology)
+        manager.start()
+
+        write_engine = manager.write_engine
+        read_engine  = manager.next_read_engine()  # round-robin
+
+        await manager.dispose()
+    """
+
+    def __init__(self, topology: ClusterTopology) -> None:
+        self._topology = topology
+        self._engines: dict[str, AsyncEngine] = {}
+        self._write_key: str = 'primary'
+        self._read_keys: list[str] = []
+        self._read_cycle: cycle[str] | None = None
+
+    # ── Lifecycle ────────────────────────────────────────────────────────
+
+    def start(self) -> EngineManager:
+        """Create all engines based on the topology. Returns *self* for chaining.
+
+        Returns:
+            The engine manager instance itself, for method chaining.
+        """
+        self._engines[self._write_key] = self._create_engine(self._topology.primary)
+
+        read_sources = self._topology.get_all_read_sources()
+        for idx, node in enumerate(read_sources):
+            key = f'read_{idx}'
+            self._engines[key] = self._create_engine(node)
+            self._read_keys.append(key)
+
+        if self._read_keys:
+            self._read_cycle = cycle(self._read_keys)
+
+        return self
+
+    async def dispose(self) -> None:
+        """Dispose all engine connection pools gracefully."""
+        for engine in self._engines.values():
+            await engine.dispose()
+        self._engines.clear()
+        self._read_keys.clear()
+        self._read_cycle = None
+
+    # ── Engine Access ────────────────────────────────────────────────────
+
+    @property
+    def write_engine(self) -> AsyncEngine:
+        """The primary engine used for all write operations.
+
+        Returns:
+            The primary async SQLAlchemy engine.
+        """
+        return self._engines[self._write_key]
+
+    def next_read_engine(self) -> AsyncEngine:
+        """Return the next read engine via round-robin; falls back to the write engine.
+
+        Returns:
+            An async SQLAlchemy engine.
+        """
+        if self._read_cycle is not None:
+            return self._engines[next(self._read_cycle)]
+        return self.write_engine
+
+    def get_engine(self, name: str) -> AsyncEngine:
+        """Get a specific engine by its key (e.g. ``'primary'``, ``'read_0'``).
+
+        Args:
+            name: Engine key.
+
+        Returns:
+            An async SQLAlchemy engine.
+        """
+        return self._engines[name]
+
+    @property
+    def engines(self) -> dict[str, AsyncEngine]:
+        """All engines managed by the engine manager.
+
+        Returns:
+            A dictionary of engine key -> async SQLAlchemy engine.
+        """
+        return dict(self._engines)
+
+    @property
+    def topology(self) -> ClusterTopology:
+        """The database topology configuration.
+
+        Returns:
+            A ClusterTopology instance.
+        """
+        return self._topology
+
+    # ── Internal ─────────────────────────────────────────────────────────
+
+    @staticmethod
+    def _create_engine(node: NodeConfig) -> AsyncEngine:
+        """Create an async SQLAlchemy engine based on the given node configuration.
+
+        Args:
+            node: Database node configuration.
+
+        Returns:
+            An async SQLAlchemy engine.
+        """
+        return create_async_engine(
+            node.url,
+            pool_size=node.pool_size,
+            max_overflow=node.max_overflow,
+            pool_timeout=node.pool_timeout,
+            pool_recycle=node.pool_recycle,
+            pool_pre_ping=node.pool_pre_ping,
+            echo=node.echo,
+            connect_args=node.connect_args,
+        )