atinypgtool 0.0.1__py3-none-any.whl

atinypgtool/__init__.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+from atinypgtool.core import close, close_all, init, with_cursor
+from atinypgtool.sql import SQLHelper
+from atinypgtool.utils import CursorPlaceholder
+
+__all__ = (
+    'init',
+    'close',
+    'close_all',
+    'with_cursor',
+    'CursorPlaceholder',
+    'SQLHelper',
+)

atinypgtool/core.py

@@ -0,0 +1,196 @@
+# -*- coding: UTF-8 -*-
+import asyncio
+import contextlib
+import functools
+import inspect
+import logging
+import typing
+
+import orjson
+from psycopg import AsyncConnection, AsyncCursor
+from psycopg.types.json import set_json_dumps, set_json_loads
+from psycopg_pool import AsyncConnectionPool
+from psycopg_pool.abc import AsyncKwargsParam
+
+from atinypgtool.utils import ConfigureFunc, SequencePlaceholder
+
+set_json_loads(orjson.loads)
+set_json_dumps(orjson.dumps)
+
+_NAMED_POOL_DICT: dict[str, AsyncConnectionPool] = {}
+_POOL_CHECKER_TASK_DICT: dict[str, asyncio.Task[None]] = {}
+
+logger = logging.getLogger(__name__)
+
+
+async def _check_pool_forever(
+    *,
+    name: str,
+    pool: AsyncConnectionPool,
+    interval: int,
+) -> None:
+    while True:
+        await asyncio.sleep(interval)
+        try:
+            await pool.check()
+        except Exception as e:
+            logger.error(
+                'Postgres pool "%s" health check failed',
+                name,
+                exc_info=e,
+            )
+
+
+def _ensure_pool_checker(
+    *,
+    name: str,
+    pool: AsyncConnectionPool,
+    interval: int,
+) -> None:
+    task = _POOL_CHECKER_TASK_DICT.get(name)
+    if task and not task.done():
+        return
+    _POOL_CHECKER_TASK_DICT[name] = asyncio.create_task(
+        _check_pool_forever(name=name, pool=pool, interval=interval),
+        name=f'postgres-pool-checker-{name}',
+    )
+
+
+async def _stop_pool_checker(*, name: str) -> None:
+    task = _POOL_CHECKER_TASK_DICT.pop(name, None)
+    if not task:
+        return
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await task
+
+
+async def _stop_pool_checkers() -> None:
+    for name in tuple(_POOL_CHECKER_TASK_DICT):
+        await _stop_pool_checker(name=name)
+
+
+async def init(
+    *,
+    name: str,
+    dsn: str,
+    pool_size: int = 16,
+    open_timeout: int = 15,
+    health_check_interval: int = 0,
+    configure_funcs: typing.Sequence[ConfigureFunc] = SequencePlaceholder,
+) -> None:
+    if name in _NAMED_POOL_DICT:
+        raise SyntaxError(f'Pool "{name}" already exists')
+    if pool_size < 1:
+        raise SyntaxError('Pool size should be greater than 0')
+    if health_check_interval < 0:
+        raise SyntaxError('Health check interval should not be less than 0')
+    minsize, maxsize = pool_size, pool_size
+    if minsize > 4:
+        minsize = 4
+    kwargs: AsyncKwargsParam = {
+        'autocommit': False,
+        # 新建连接超过 5 秒未完成时快速失败，避免请求长时间卡在重连上。
+        'connect_timeout': 5,
+        # 开启 TCP keepalive 探测，用于发现已经被网络或数据库断开的空闲连接。
+        'keepalives': 1,
+        # TCP 连接空闲 15 秒后发送第一次 keepalive 探测。
+        'keepalives_idle': 15,
+        # 第一次探测后，如果没有响应，每隔 5 秒继续重试探测。
+        'keepalives_interval': 5,
+        # 连续 3 次 keepalive 探测无响应后，将连接视为已断开。
+        'keepalives_count': 3,
+    }
+
+    configure = None
+    if configure_funcs:
+
+        async def _configure(conn: AsyncConnection) -> None:
+            for func in configure_funcs:
+                if inspect.iscoroutinefunction(func):
+                    await func(conn)
+                else:
+                    func(conn)
+
+        configure = _configure
+
+    pool = AsyncConnectionPool(
+        conninfo=dsn,
+        min_size=minsize,
+        max_size=maxsize,
+        open=False,
+        configure=configure,
+        name=name,
+        kwargs=kwargs,
+        check=AsyncConnectionPool.check_connection,
+        num_workers=4,
+        timeout=5,
+        reconnect_timeout=30,
+    )
+    await pool.open(wait=True, timeout=open_timeout)
+    _NAMED_POOL_DICT[name] = pool
+    if health_check_interval > 0:
+        _ensure_pool_checker(
+            name=name,
+            pool=pool,
+            interval=health_check_interval,
+        )
+
+
+async def close(*, name: str) -> None:
+    if name not in _NAMED_POOL_DICT:
+        raise ValueError(f'Pool "{name}" not found')
+    pool = _NAMED_POOL_DICT.pop(name)
+    await _stop_pool_checker(name=name)
+    await pool.close()
+
+
+async def close_all() -> None:
+    await _stop_pool_checkers()
+    for pool in _NAMED_POOL_DICT.values():
+        await pool.close()
+    _NAMED_POOL_DICT.clear()
+
+
+def with_cursor(*, name: str, transaction: bool) -> typing.Callable:
+    def wrapper(func: typing.Callable) -> typing.Callable:
+        argspec = inspect.getfullargspec(func)
+        if all('cursor' not in x for x in (argspec.args, argspec.kwonlyargs)):
+            raise SyntaxError('`cursor` is a required argument')
+
+        @functools.wraps(func)
+        async def wrapped(*args, **kwargs):  # noqa: ANN002, ANN003, ANN202
+            if name not in _NAMED_POOL_DICT:
+                raise SyntaxError(f'Pool "{name}" not found')
+            if 'cursor' in kwargs and kwargs['cursor']:
+                raise SyntaxError('`cursor` is a reserved argument')
+            async with _NAMED_POOL_DICT[name].connection() as conn:  # noqa: SIM117
+                async with conn.cursor() as cursor:
+                    kwargs['cursor'] = cursor
+                    if transaction:
+                        async with conn.transaction():
+                            result = await func(*args, **kwargs)
+                    else:
+                        result = await func(*args, **kwargs)
+            return result
+
+        return wrapped
+
+    return wrapper
+
+
+@contextlib.asynccontextmanager
+async def with_cursor_context(
+    *,
+    name: str,
+    transaction: bool,
+) -> typing.AsyncGenerator[AsyncCursor, None]:
+    if name not in _NAMED_POOL_DICT:
+        raise RuntimeError(f'Pool "{name}" not found')
+    async with _NAMED_POOL_DICT[name].connection() as conn:  # noqa: SIM117
+        async with conn.cursor() as cursor:
+            if transaction:
+                async with conn.transaction():
+                    yield cursor
+            else:
+                yield cursor

atinypgtool/sql.py

@@ -0,0 +1,105 @@
+# -*- coding: utf-8 -*-
+import dataclasses
+import inspect
+import typing
+
+from psycopg.sql import SQL
+
+from atinypgtool.utils import SequencePlaceholder
+
+
+class SQLHelper:
+    def __init__(self, base_model_cls: type | None = None) -> None:
+        if base_model_cls is not None:
+            try:
+                assert isinstance(base_model_cls, type)
+                assert dataclasses.is_dataclass(base_model_cls)
+            except AssertionError as e:
+                raise SyntaxError(
+                    f'Base model "{base_model_cls}" '
+                    f'should be a dataclass or None',
+                ) from e
+        self._base_model_cls = base_model_cls
+
+    def _gen_fields(self, *, dataclass: type) -> tuple[list[str], list[str]]:
+        default_fields = []
+        if self._base_model_cls is not None and issubclass(
+            dataclass,
+            self._base_model_cls,  # noqa: type: ignore
+        ):
+            default_fields = [
+                i.strip('_')
+                for i in inspect.get_annotations(self._base_model_cls)
+            ]
+        custom_fields = [
+            i.strip('_') for i in inspect.get_annotations(dataclass)
+        ]
+        return default_fields, custom_fields
+
+    def gen_select_base(
+        self,
+        *,
+        dataclass: type,
+        table_name: str,
+    ) -> str:
+        default_fields, custom_fields = self._gen_fields(dataclass=dataclass)
+        return (
+            f'SELECT {", ".join(default_fields + custom_fields)} '
+            f'FROM {table_name}'
+        )
+
+    def gen_get(self, *, dataclass: type, table_name: str) -> SQL:
+        default_fields, custom_fields = self._gen_fields(dataclass=dataclass)
+        return self.make_sql(
+            sql_str=(
+                f'SELECT {", ".join(default_fields + custom_fields)} '
+                f'FROM {table_name} WHERE id=%s'
+            ),
+        )
+
+    def gen_insert(
+        self,
+        *,
+        dataclass: type,
+        table_name: str,
+        limited_fields: typing.Sequence[str] = SequencePlaceholder,
+    ) -> SQL:
+        _, custom_fields = self._gen_fields(dataclass=dataclass)
+        if limited_fields and (
+            not set(limited_fields).issubset(set(custom_fields))
+        ):
+            unexcepted = set(limited_fields) - set(custom_fields)
+            raise SyntaxError(f'Field names not allowed: {unexcepted}')
+        fields = limited_fields or custom_fields
+        return self.make_sql(
+            sql_str=(
+                f'INSERT INTO {table_name} ({", ".join(fields)}) '
+                f'VALUES ({", ".join(["%s"] * len(fields))})'
+            ),
+        )
+
+    def gen_insert_returning(
+        self,
+        *,
+        dataclass: type,
+        table_name: str,
+        limited_fields: typing.Sequence[str] = SequencePlaceholder,
+    ) -> SQL:
+        default_fields, custom_fields = self._gen_fields(dataclass=dataclass)
+        if limited_fields and (
+            not set(limited_fields).issubset(set(custom_fields))
+        ):
+            unexcepted = set(limited_fields) - set(custom_fields)
+            raise SyntaxError(f'Field names not allowed: {unexcepted}')
+        fields = limited_fields or custom_fields
+        return self.make_sql(
+            sql_str=(
+                f'INSERT INTO {table_name} ({", ".join(fields)}) '
+                f'VALUES ({", ".join(["%s"] * len(fields))}) '
+                f'RETURNING {", ".join(default_fields + custom_fields)}'
+            ),
+        )
+
+    @staticmethod
+    def make_sql(*, sql_str: str) -> SQL:
+        return SQL(obj=sql_str.strip())  # type: ignore

atinypgtool/utils.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+import dataclasses
+import typing
+
+from psycopg import AsyncConnection, AsyncCursor
+from psycopg.pq import PGconn
+
+
+class _SequencePlaceholder[T](typing.Sequence):
+    @typing.overload
+    def __getitem__(self, index: int, /) -> T: ...
+
+    @typing.overload
+    def __getitem__(self, index: slice, /) -> typing.Sequence[T]: ...
+
+    def __getitem__(self, _, /):
+        raise SyntaxError('SequencePlaceholder is empty')
+
+    def __len__(self) -> int:
+        return 0
+
+    def __bool__(self) -> bool:
+        return False
+
+
+SequencePlaceholder = _SequencePlaceholder[typing.Any]()
+
+
+type ConfigureFunc = typing.Callable[
+    [AsyncConnection],
+    typing.Awaitable | None,
+]
+
+
+@typing.runtime_checkable
+class DataclassInstance(typing.Protocol):
+    __dataclass_fields__: typing.ClassVar[
+        dict[str, dataclasses.Field[typing.Any]]
+    ]
+
+
+class _PGConnPlaceholder(PGconn):
+    pass
+
+
+class _CursorPlaceholder(AsyncCursor):
+    def __init__(self) -> None:
+        super().__init__(
+            connection=AsyncConnection(pgconn=_PGConnPlaceholder()),
+        )
+
+    def __bool__(self) -> bool:
+        return False
+
+
+CursorPlaceholder = _CursorPlaceholder()

atinypgtool-0.0.1.dist-info/METADATA

@@ -0,0 +1,13 @@
+Metadata-Version: 2.4
+Name: atinypgtool
+Version: 0.0.1
+Summary: AsyncIO-based Tiny PostgreSQL Tool
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: orjson
+Requires-Dist: psycopg[binary,pool]
+Requires-Dist: psycopg_pool
+Dynamic: license-file
+
+# ez-psycopg

atinypgtool-0.0.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+atinypgtool/__init__.py,sha256=QfPBszp-bUE0ukn72992P9sQPJvt9AghPtFUYNJpnAs,293
+atinypgtool/core.py,sha256=S0Jn2toXaGggSFw1hssmtsbiVPOipTAAih3uK_UJ3xA,5998
+atinypgtool/sql.py,sha256=VW4DtGA_c_zqsag6WC8M_K0kNqnAccHz2K3ksTK6CL0,3586
+atinypgtool/utils.py,sha256=tri9votbV-rnUH_VT8ompDGRS4c7pJ8pLdT-6AHZEh4,1186
+atinypgtool-0.0.1.dist-info/licenses/LICENSE,sha256=tc4Ew4FFA8XrGfDzY3obl6RgMc-ORvOftbv-799gUYk,1072
+atinypgtool-0.0.1.dist-info/METADATA,sha256=UqLVMKGkQd2eVZOv9yhqhCr41tCbeD0rnt_E4Hqhrws,307
+atinypgtool-0.0.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+atinypgtool-0.0.1.dist-info/top_level.txt,sha256=RhTyPWKIJHqrmc1S-kHDio2Kc9LFFiN0XTAD9KSicAM,12
+atinypgtool-0.0.1.dist-info/RECORD,,

atinypgtool-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

atinypgtool-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 M@gicAgCl(DEAD)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

atinypgtool-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+atinypgtool