toms-fast 0.2.1__py3-none-any.whl

tomskit/sqlalchemy/config.py

@@ -0,0 +1,125 @@
+from typing import Any
+from urllib.parse import quote_plus
+from pydantic import Field, NonNegativeInt, PositiveInt, computed_field
+from pydantic_settings import BaseSettings
+
+
+class DatabaseConfig(BaseSettings):
+    """
+    Configuration settings for the database
+    """
+
+    DB_HOST: str = Field(
+        description="db host",
+        default="localhost",
+    )
+
+    DB_PORT: PositiveInt = Field(
+        description="db port",
+        default=5432,
+    )
+
+    DB_USERNAME: str = Field(
+        description="db username",
+        default="",
+    )
+
+    DB_PASSWORD: str = Field(
+        description="db password",
+        default="",
+    )
+
+    DB_DATABASE: str = Field(
+        description="db database",
+        default="tomskitdb",
+    )
+
+    DB_CHARSET: str = Field(
+        description="db charset",
+        default="",
+    )
+
+    DB_EXTRAS: str = Field(
+        description="db extras options. Example: keepalives_idle=60&keepalives=1",
+        default="",
+    )
+
+    SQLALCHEMY_DATABASE_URI_SCHEME: str = Field(
+        description="db uri scheme",
+        default="mysql+aiomysql",
+    )
+
+    SQLALCHEMY_DATABASE_SYNC_URI_SCHEME: str = Field(
+        description="db uri scheme",
+        default="mysql+pymysql",
+    )
+
+    @computed_field # type: ignore
+    @property
+    def SQLALCHEMY_DATABASE_URI(self) -> str:
+        db_extras = (
+            f"{self.DB_EXTRAS}&client_encoding={self.DB_CHARSET}" if self.DB_CHARSET else self.DB_EXTRAS
+        ).strip("&")
+        db_extras = f"?{db_extras}" if db_extras else ""
+        return (
+            f"{self.SQLALCHEMY_DATABASE_URI_SCHEME}://"
+            f"{quote_plus(self.DB_USERNAME)}:{quote_plus(self.DB_PASSWORD)}@{self.DB_HOST}:{self.DB_PORT}/{self.DB_DATABASE}"
+            f"{db_extras}"
+        )
+    
+    @computed_field # type: ignore
+    @property
+    def SQLALCHEMY_DATABASE_SYNC_URI(self) -> str:
+        db_extras = (
+            f"{self.DB_EXTRAS}&client_encoding={self.DB_CHARSET}" if self.DB_CHARSET else self.DB_EXTRAS
+        ).strip("&")
+        db_extras = f"?{db_extras}" if db_extras else ""
+        return (
+            f"{self.SQLALCHEMY_DATABASE_SYNC_URI_SCHEME}://"
+            f"{quote_plus(self.DB_USERNAME)}:{quote_plus(self.DB_PASSWORD)}@{self.DB_HOST}:{self.DB_PORT}/{self.DB_DATABASE}"
+            f"{db_extras}"
+        )
+
+
+    SQLALCHEMY_POOL_SIZE: NonNegativeInt = Field(
+        description="pool size of SqlAlchemy",
+        default=300,
+    )
+
+    SQLALCHEMY_MAX_OVERFLOW: NonNegativeInt = Field(
+        description="max overflows for SqlAlchemy",
+        default=10,
+    )
+
+    SQLALCHEMY_POOL_RECYCLE: NonNegativeInt = Field(
+        description="SqlAlchemy pool recycle",
+        default=3600,
+    )
+
+    SQLALCHEMY_POOL_PRE_PING: bool = Field(
+        description="whether to enable pool pre-ping in SqlAlchemy",
+        default=False,
+    )
+
+    SQLALCHEMY_ECHO: bool = Field(
+        description="whether to enable SqlAlchemy echo",
+        default=False,
+    )
+
+    SQLALCHEMY_POOL_ECHO: bool = Field(
+        description="whether to enable pool echo in SqlAlchemy",
+        default=False,
+    )
+
+    @computed_field # type: ignore
+    @property
+    def SQLALCHEMY_ENGINE_OPTIONS(self) -> dict[str, Any]:
+        return {
+            "pool_size": self.SQLALCHEMY_POOL_SIZE,
+            "max_overflow": self.SQLALCHEMY_MAX_OVERFLOW,
+            "pool_recycle": self.SQLALCHEMY_POOL_RECYCLE,
+            "pool_pre_ping": self.SQLALCHEMY_POOL_PRE_PING,
+            "echo": self.SQLALCHEMY_ECHO,
+            "echo_pool": self.SQLALCHEMY_POOL_ECHO,
+            # "connect_args": {"options": "-c timezone=UTC"},
+        }

tomskit/sqlalchemy/database.py

@@ -0,0 +1,125 @@
+"""
+数据库扩展模块
+"""
+
+from contextvars import ContextVar
+from typing import Any, Optional
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.sql import Select
+
+from tomskit.sqlalchemy import Pagination, SelectPagination, SQLAlchemy
+
+class DatabaseSession(SQLAlchemy):
+    database_session_ctx: ContextVar[Optional[AsyncSession]] = ContextVar('database_session', default=None)
+
+    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:
+        return await SelectPagination(
+            select=select,
+            session=self.session,
+            page=page,
+            per_page=per_page,
+            max_per_page=max_per_page,
+            error_out=error_out,
+            count=count,
+        )  # type: ignore
+    
+    @property
+    def session(self):
+        return self.database_session_ctx.get()
+    
+    def create_session(self)-> AsyncSession:
+        """
+        创建一个新的会话并手动将其设置为ContextVar。
+        """
+        session = self._SessionLocal()  # type: ignore
+        self.database_session_ctx.set(session)
+        return session
+    
+    async def close_session(self, session):
+        """
+        Close the session and reset the context variable manually.
+        """
+        await session.aclose()
+        self.database_session_ctx.set(None)
+
+    def initialize_session_pool(self, db_url: str, engine_options: Optional[dict[str, Any]] = None):
+        """
+        Initialize the database with the given database URL.
+        Create the AsyncEngine and SessionLocal for database operations.
+        """
+        # Create the asynchronous engine
+        default_options = {
+            "pool_size": 10,        # 连接池大小 (Connection pool size)
+            "max_overflow": 20,     # 允许的额外连接数 (Extra connections allowed)
+            "pool_timeout": 30,     # 获取连接的超时时间 (Timeout for acquiring a connection)
+            "pool_recycle": 1800,   # 空闲后回收连接的时间 (Recycle connections after being idle)
+            "echo": False,          # 调试时打印SQL查询 (Echo SQL queries for debugging)
+            "echo_pool": False      # 调试时打印连接池信息 (Echo pool information for debugging)
+        }
+        
+        engine_options = engine_options.copy() if engine_options else {}
+        for key, default_val in default_options.items():
+            if key not in engine_options:
+                engine_options[key] = default_val
+
+        self._engine = create_async_engine(db_url, **engine_options)
+
+        # Create the session factory for AsyncSession
+        self._SessionLocal = async_sessionmaker(
+            bind=self._engine, 
+            class_=AsyncSession,
+            expire_on_commit=False
+        )
+
+    # Start of Selection
+    def get_session_pool_info(self) -> dict:
+        """
+        获取当前数据库连接池的详细信息
+        pool_size: 连接池大小
+        pool_checkedin: 已检查入的连接数
+        pool_checkedout: 已检查出的连接数
+        pool_overflow: 溢出的连接数
+        """
+        if self._engine is None or self._engine.pool is None:
+            return {"error": "数据库引擎未初始化"}
+        
+        return {
+            "pool_size": self._engine.pool.size(),  # type: ignore
+            "pool_checkedin": self._engine.pool.checkedin(), # type: ignore
+            "pool_checkedout": self._engine.pool.checkedout(), # type: ignore
+            "pool_overflow": self._engine.pool.overflow(), # type: ignore
+        }
+
+
+    async def close_session_pool(self):
+        if self._engine is not None:
+            await self._engine.dispose()
+
+    def create_celery_session(self, config: dict[str, Any]):
+        self.initialize_session_pool(
+            config["SQLALCHEMY_DATABASE_URI"],
+            config["SQLALCHEMY_ENGINE_OPTIONS"]
+        )
+        session = self.create_session()
+        return session
+
+    async def close_celery_session(self, session):
+        """
+        关闭会话并手动重置上下文变量。
+        """
+        await session.aclose()
+        self.database_session_ctx.set(None)
+        if self._engine is not None:
+            await self._engine.dispose()
+
+
+db = DatabaseSession()

tomskit/sqlalchemy/pagination.py

@@ -0,0 +1,359 @@
+from __future__ import annotations
+
+import typing as t
+from math import ceil
+
+import sqlalchemy as sa
+import sqlalchemy.orm as sa_orm
+
+from tomskit.server import raise_api_error
+
+
+class Pagination:
+    """Apply an offset and limit to the query based on the current page and number of
+    items per page.
+
+    Don't create pagination objects manually. They are created by
+    :meth:`.SQLAlchemy.paginate` and :meth:`.Query.paginate`.
+
+    This is a base class, a subclass must implement :meth:`_query_items` and
+    :meth:`_query_count`. Those methods will use arguments passed as ``kwargs`` to
+    perform the queries.
+
+    :param page: The current page, used to calculate the offset. Defaults to the
+        ``page`` query arg during a request, or 1 otherwise.
+    :param per_page: The maximum number of items on a page, used to calculate the
+        offset and limit. Defaults to the ``per_page`` query arg during a request,
+        or 20 otherwise.
+    :param max_per_page: The maximum allowed value for ``per_page``, to limit a
+        user-provided value. Use ``None`` for no limit. Defaults to 100.
+    :param error_out: Abort with a ``404 Not Found`` error if no items are returned
+        and ``page`` is not 1, or if ``page`` or ``per_page`` is less than 1, or if
+        either are not ints.
+    :param count: Calculate the total number of values by issuing an extra count
+        query. For very complex queries this may be inaccurate or slow, so it can be
+        disabled and set manually if necessary.
+    :param kwargs: Information about the query to paginate. Different subclasses will
+        require different arguments.
+
+    .. versionchanged:: 3.0
+        Iterating over a pagination object iterates over its items.
+
+    .. versionchanged:: 3.0
+        Creating instances manually is not a public API.
+    """
+
+    def __init__(
+        self,
+        page: int | None = None,
+        per_page: int | None = None,
+        max_per_page: int | None = 100,
+        error_out: bool = True,
+        count: bool = True,
+        **kwargs: t.Any,
+    ) -> None:
+        self._query_args = kwargs
+        page, per_page = self._prepare_page_args(
+            page=page,
+            per_page=per_page,
+            max_per_page=max_per_page,
+            error_out=error_out,
+        )
+
+        self.page: int = page
+        """The current page."""
+
+        self.per_page: int = per_page
+        """The maximum number of items on a page."""
+
+        self.max_per_page: int | None = max_per_page
+        """The maximum allowed value for ``per_page``."""
+
+        self.error_out = error_out
+        self.count = count
+        # items = self._query_items()
+
+
+        # if not items and page != 1 and error_out:
+        #     raise_api_error(404)
+
+        # self.items: list[t.Any] = items
+        """The items on the current page. Iterating over the pagination object is
+        equivalent to iterating over the items.
+        """
+
+        # if count:
+        #     total = self._query_count()
+        # else:
+        #     total = None
+
+        # self.total: int | None = total
+        """The total number of items across all pages."""
+
+    def __await__(self):
+        return self.initialize().__await__()
+    
+    async def initialize(self):
+        items = await self._query_items()
+        if not items and self.page != 1 and self.error_out:
+            raise_api_error(404, message="Page not found")
+        self.items: list[t.Any] = items # type: ignore
+
+        if self.count:
+            total = await self._query_count()
+        else:
+            total = None
+
+        self.total: int | None = total # type: ignore
+
+        return self
+
+    @staticmethod
+    def _prepare_page_args(
+        *,
+        page: int | None = None,
+        per_page: int | None = None,
+        max_per_page: int | None = None,
+        error_out: bool = True,
+    ) -> tuple[int, int]:
+        
+        if page is None:
+            page = 1
+
+        if per_page is None:
+            per_page = 20
+
+        if max_per_page is not None:
+            per_page = min(per_page, max_per_page)
+
+        if page < 1:
+            if error_out:
+                raise_api_error(404, message=f"Invalid page number: {page}")
+            else:
+                page = 1
+
+        if per_page < 1:
+            if error_out:
+                raise_api_error(404, message=f"Invalid items per page value: {per_page}")
+            else:
+                per_page = 20
+
+        return page, per_page
+
+    @property
+    def _query_offset(self) -> int:
+        """The index of the first item to query, passed to ``offset()``.
+
+        :meta private:
+
+        .. versionadded:: 3.0
+        """
+        return (self.page - 1) * self.per_page
+
+    async def _query_items(self) -> list[t.Any]:
+        """Execute the query to get the items on the current page.
+
+        Uses init arguments stored in :attr:`_query_args`.
+
+        :meta private:
+
+        .. versionadded:: 3.0
+        """
+        raise NotImplementedError
+
+    async def _query_count(self) -> int:
+        """Execute the query to get the total number of items.
+
+        Uses init arguments stored in :attr:`_query_args`.
+
+        :meta private:
+
+        .. versionadded:: 3.0
+        """
+        raise NotImplementedError
+
+    @property
+    def first(self) -> int:
+        """The number of the first item on the page, starting from 1, or 0 if there are
+        no items.
+
+        .. versionadded:: 3.0
+        """
+        if len(self.items) == 0:
+            return 0
+
+        return (self.page - 1) * self.per_page + 1
+
+    @property
+    def last(self) -> int:
+        """The number of the last item on the page, starting from 1, inclusive, or 0 if
+        there are no items.
+
+        .. versionadded:: 3.0
+        """
+        first = self.first
+        return max(first, first + len(self.items) - 1)
+
+    @property
+    def pages(self) -> int:
+        """The total number of pages."""
+        if self.total == 0 or self.total is None:
+            return 0
+
+        return ceil(self.total / self.per_page)
+
+    @property
+    def has_prev(self) -> bool:
+        """``True`` if this is not the first page."""
+        return self.page > 1
+
+    @property
+    def prev_num(self) -> int | None:
+        """The previous page number, or ``None`` if this is the first page."""
+        if not self.has_prev:
+            return None
+
+        return self.page - 1
+
+    async def prev(self, *, error_out: bool = False) -> Pagination:
+        """Query the :class:`Pagination` object for the previous page.
+
+        :param error_out: Abort with a ``404 Not Found`` error if no items are returned
+            and ``page`` is not 1, or if ``page`` or ``per_page`` is less than 1, or if
+            either are not ints.
+        """
+        p = await type(self)(
+            page=self.page - 1,
+            per_page=self.per_page,
+            error_out=error_out,
+            count=False,
+            **self._query_args,
+        )
+        p.total = self.total
+        return p
+
+    @property
+    def has_next(self) -> bool:
+        """``True`` if this is not the last page."""
+        return self.page < self.pages
+
+    @property
+    def next_num(self) -> int | None:
+        """The next page number, or ``None`` if this is the last page."""
+        if not self.has_next:
+            return None
+
+        return self.page + 1
+
+    async def next(self, *, error_out: bool = False) -> Pagination:
+        """Query the :class:`Pagination` object for the next page.
+
+        :param error_out: Abort with a ``404 Not Found`` error if no items are returned
+            and ``page`` is not 1, or if ``page`` or ``per_page`` is less than 1, or if
+            either are not ints.
+        """
+        p = await type(self)(
+            page=self.page + 1,
+            per_page=self.per_page,
+            max_per_page=self.max_per_page,
+            error_out=error_out,
+            count=False,
+            **self._query_args,
+        )
+        p.total = self.total
+        return p
+
+    def iter_pages(
+        self,
+        *,
+        left_edge: int = 2,
+        left_current: int = 2,
+        right_current: int = 4,
+        right_edge: int = 2,
+    ) -> t.Iterator[int | None]:
+        """Yield page numbers for a pagination widget. Skipped pages between the edges
+        and middle are represented by a ``None``.
+
+        For example, if there are 20 pages and the current page is 7, the following
+        values are yielded.
+
+        .. code-block:: python
+
+            1, 2, None, 5, 6, 7, 8, 9, 10, 11, None, 19, 20
+
+        :param left_edge: How many pages to show from the first page.
+        :param left_current: How many pages to show left of the current page.
+        :param right_current: How many pages to show right of the current page.
+        :param right_edge: How many pages to show from the last page.
+
+        .. versionchanged:: 3.0
+            Improved efficiency of calculating what to yield.
+
+        .. versionchanged:: 3.0
+            ``right_current`` boundary is inclusive.
+
+        .. versionchanged:: 3.0
+            All parameters are keyword-only.
+        """
+        pages_end = self.pages + 1
+
+        if pages_end == 1:
+            return
+
+        left_end = min(1 + left_edge, pages_end)
+        yield from range(1, left_end)
+
+        if left_end == pages_end:
+            return
+
+        mid_start = max(left_end, self.page - left_current)
+        mid_end = min(self.page + right_current + 1, pages_end)
+
+        if mid_start - left_end > 0:
+            yield None
+
+        yield from range(mid_start, mid_end)
+
+        if mid_end == pages_end:
+            return
+
+        right_start = max(mid_end, pages_end - right_edge)
+
+        if right_start - mid_end > 0:
+            yield None
+
+        yield from range(right_start, pages_end)
+
+    def __iter__(self) -> t.Iterator[t.Any]:
+        yield from self.items
+
+
+class SelectPagination(Pagination):
+    """Returned by :meth:`.SQLAlchemy.paginate`. Takes ``select`` and ``session``
+    arguments in addition to the :class:`Pagination` arguments.
+
+    .. versionadded:: 3.0
+    """
+
+    def __await__(self):
+        return self.initialize().__await__()
+    
+    async def initialize(self):
+        await super().initialize()
+        return self
+
+    async def _query_items(self) -> list[t.Any]:
+        select = self._query_args["select"]
+        select = select.limit(self.per_page).offset(self._query_offset)
+        session = self._query_args["session"]
+        result =  await session.execute(select)
+        return list(result.unique().scalars())
+
+    async def _query_count(self) -> int:
+        select = self._query_args["select"]
+        sub = select.options(sa_orm.lazyload("*")).order_by(None).subquery()
+        session = self._query_args["session"]
+        result = await session.execute(sa.select(sa.func.count()).select_from(sub))
+        out = result.scalar()
+        return out  # type: ignore[no-any-return]
+

tomskit/sqlalchemy/property.py

@@ -0,0 +1,19 @@
+import functools
+
+class cached_async_property:
+    def __init__(self, func):
+        self.func = func
+        self.attr_name = f"_cached_{func.__name__}"
+        functools.update_wrapper(self, func)
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            return self
+
+        async def wrapper():
+            if not hasattr(instance, self.attr_name):
+                value = await self.func(instance)
+                setattr(instance, self.attr_name, value)
+            return getattr(instance, self.attr_name)
+
+        return wrapper()