atquery 2.11.15__py3-none-any.whl

atquery/__init__.py

@@ -0,0 +1,47 @@
+from .api_v2 import (
+    NULL,
+    ATConnection,
+    ATDataSource,
+    ATQuery,
+    ATResult,
+    ATRow,
+    ATRowCountError,
+    ATSQLError,
+    ATString,
+    ATWarning,
+    Resource,
+    create_async_datasource,
+    create_datasource,
+    escape4like,
+    execute,
+    execute_script,
+    export_csv,
+    insert_csv,
+    load,
+    sql,
+    upsert_csv,
+)
+
+__all__ = [
+    "NULL",
+    "ATConnection",
+    "ATDataSource",
+    "ATQuery",
+    "ATResult",
+    "ATRow",
+    "ATRowCountError",
+    "ATSQLError",
+    "ATString",
+    "ATWarning",
+    "Resource",
+    "create_async_datasource",
+    "create_datasource",
+    "escape4like",
+    "execute",
+    "execute_script",
+    "insert_csv",
+    "load",
+    "sql",
+    "upsert_csv",
+    "export_csv",
+]

atquery/aio/__init__.py

@@ -0,0 +1,101 @@
+"""atquery.aio - async 専用 DataSource / Query の公開 API
+
+Usage:
+    import atquery.aio
+
+    async_ds = atquery.aio.create_datasource("mysql://user:pass@host/db")
+    async_ds = atquery.aio.create_datasource("mysql://user:pass@host/db", pool_size=10)
+
+    async with async_ds.connect() as conn:
+        users = await at.aquery("SELECT__users.sql").getall(conn, {"status": "active"})
+"""
+import re
+from importlib import import_module
+from typing import Any
+from urllib.parse import parse_qsl, urlparse
+
+from atquery.aio.engine import (
+    ATAsyncConnection,
+    ATAsyncConnectionPool,
+    ATAsyncDataSource,
+    ATAsyncResult,
+)
+from atquery.aio.query import ATAsyncQuery, ATAsyncTypedQueryProxy
+
+_ASYNC_DATASOURCE_TYPES: dict[str, tuple[str, str]] = {
+    "mysql": ("atquery.dialect.mysql.aio", "ATMySQLAsyncDataSource"),
+    "mysql+mysqlconnector": ("atquery.dialect.mysql.aio", "ATMySQLAsyncDataSource"),
+    "mysql+mysqlconnector#poolable": ("atquery.dialect.mysql.aio", "ATMySQLAsyncPoolableDataSource"),
+    "mysql#poolable": ("atquery.dialect.mysql.aio", "ATMySQLAsyncPoolableDataSource"),
+}
+
+_USER = re.compile(r"^.+:\/\/([^:]+:[^@]+@)")
+
+
+def create_datasource(uri: str, pool_size: int = 0) -> ATAsyncDataSource:
+    """async 専用 DataSource を生成する。
+
+    Args:
+        uri: 接続URI。例: "mysql://user:pass@host/db"
+        pool_size: コネクションプールサイズ。0の場合はプールなし。
+
+    Returns:
+        ATAsyncDataSource: async専用DataSource
+    """
+    username = None
+    password = None
+    if match := _USER.match(uri):
+        user = match.group(1)
+        (username, password) = user.split(":", 1)
+        password = password[:-1]  # '@' を除く
+        uri = uri.replace(user, "", 1)
+
+    sr = urlparse(uri)
+    scheme = sr.scheme
+    params: dict[str, Any] = dict(parse_qsl(sr.query))
+
+    if pool_size:
+        scheme += "#poolable"
+        params["pool_size"] = pool_size
+
+    info = _ASYNC_DATASOURCE_TYPES.get(scheme)
+    if info is None:
+        raise ValueError(f'scheme "{scheme}" は async DataSource でサポートされていません。')
+
+    cls = getattr(import_module(info[0]), info[1])
+    return cls(
+        user=username or sr.username,
+        password=password or sr.password,
+        host=sr.hostname,
+        port=sr.port,
+        database=sr.path[1:],
+        **params,
+    )
+
+
+def load(py: str, filepath: str, use_cache: bool = True) -> ATAsyncQuery:
+    """SQLファイルを読み込んで ATAsyncQuery のインスタンスを生成する。
+
+    Args:
+        py: ファイル検索の起点モジュール（通常は __file__）
+        filepath: SQLファイルの相対パス
+        use_cache: キャッシュを使用するか
+
+    Returns:
+        ATAsyncQuery: インスタンス
+    """
+    from atquery.factory import Resource
+
+    return Resource.aload(py, filepath, use_cache=use_cache)
+
+
+__all__ = [
+    "ATAsyncConnection",
+    "ATAsyncConnectionPool",
+    "ATAsyncDataSource",
+    "ATAsyncResult",
+    "ATAsyncQuery",
+    "ATAsyncTypedQueryProxy",
+    "create_datasource",
+    "load",
+]

atquery/aio/engine.py

@@ -0,0 +1,431 @@
+"""atquery async engine - DataSource / Connection / Result / ConnectionPool
+
+ATAsyncDataSource, ATAsyncConnection, ATAsyncResult, ATAsyncConnectionPool を提供する。
+engine.py に asyncio 依存を持ち込まないため、このモジュールに分離している。
+"""
+import asyncio
+import collections
+import inspect
+import logging
+from typing import Any, Callable, Generic, Sequence, TypeVar
+
+from atquery.engine import (
+    ATConnectionPoolTimeoutError,
+    ATRow,
+    ATSQLError,
+    _reject_multi_statement_sql,
+)
+from atquery.type import NULL
+
+_pool_log = logging.getLogger("atquery.pool")
+
+ConnType = TypeVar("ConnType")
+
+
+class _AsyncConnectContext:
+    """async with / await の両方に対応するコンテキスト。
+
+    asyncpg と同じパターン:
+        async with ds.connect() as conn: ...
+        conn = await ds.connect()
+    """
+
+    def __init__(self, coro) -> None:
+        self._coro = coro
+        self._conn: "ATAsyncConnection | None" = None
+
+    def __await__(self):
+        return self._coro.__await__()
+
+    async def __aenter__(self) -> "ATAsyncConnection":
+        self._conn = await self._coro
+        return self._conn  # type: ignore[return-value]  # _conn は __aenter__ 後必ず非 None
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self._conn is not None:
+            if self._conn.transaction:
+                if exc_type is None:
+                    await self._conn.commit()
+                else:
+                    await self._conn.rollback()
+            await self._conn.close()
+
+
+class ATAsyncResult:
+    """非同期結果セット。非同期カーソルのラッパー。"""
+
+    def __init__(
+        self,
+        cursor: Any,
+        column_names: Sequence[str],
+        column_types: Sequence[str | None] | None = None,
+    ) -> None:
+        self.cursor = cursor
+        self.column_names = tuple(column_names)
+        self.column_types: tuple[str | None, ...] = (
+            tuple(column_types) if column_types else (None,) * len(self.column_names)
+        )
+
+    def __aiter__(self) -> "ATAsyncResult":
+        return self
+
+    async def __anext__(self) -> ATRow:
+        row = await self.cursor.fetchone()
+        if row is None:
+            raise StopAsyncIteration
+        return ATRow(self.column_names, row)
+
+    async def __aenter__(self) -> "ATAsyncResult":
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        await self.close()
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self.column_names})[{self.rowcount}]"
+
+    @property
+    def rowcount(self) -> int:
+        return self.cursor.rowcount
+
+    async def fetchone(self) -> ATRow | None:
+        row = await self.cursor.fetchone()
+        if row is None:
+            return None
+        return ATRow(self.column_names, row)
+
+    async def fetchall(self) -> list[ATRow]:
+        rows = await self.cursor.fetchall()
+        return [ATRow(self.column_names, row) for row in rows]
+
+    async def fetchmany(self, size: int) -> list[ATRow]:
+        rows = await self.cursor.fetchmany(size)
+        return [ATRow(self.column_names, row) for row in rows]
+
+    async def scalar(self) -> Any:
+        row = await self.fetchone()
+        return row[0] if row else None
+
+    async def close(self) -> None:
+        await self.cursor.close()
+
+
+class ATAsyncConnection(Generic[ConnType]):
+    """非同期接続の基底クラス。"""
+
+    connection: Any = None
+    transaction: bool = False
+
+    def __init__(self, connection: ConnType) -> None:
+        self.connection = connection
+
+    @property
+    def dialect(self) -> str:
+        raise NotImplementedError()
+
+    async def commit(self) -> None:
+        await self.connection.commit()
+
+    async def rollback(self) -> None:
+        await self.connection.rollback()
+
+    async def close(self) -> None:
+        await self.connection.close()
+
+    async def execute(
+        self,
+        sql: str,
+        params: dict[str, Any] | None = None,
+        buffered: bool = False,
+    ) -> ATAsyncResult:
+        if not sql:
+            raise ValueError(f'"sql" is required: ${sql}')
+
+        _reject_multi_statement_sql(sql, params)
+
+        params = {k: (None if v is NULL else v) for k, v in params.items()} if params else {}
+
+        try:
+            return await self.__aexecute__(self.connection, sql, params, buffered)
+        except ATSQLError as ex:
+            raise ex
+        except Exception as ex:
+            raise ATSQLError(str(ex), ex, sql=sql, sql_file=None, params=params)
+
+    async def __aexecute__(
+        self, conn: ConnType, sql: str, params: dict, buffered: bool
+    ) -> ATAsyncResult:
+        raise NotImplementedError()
+
+    async def execute_script(self, script: str) -> None:
+        raise NotImplementedError()
+
+    async def is_connected(self) -> bool:
+        try:
+            result = await self.execute("SELECT 1")
+            await result.fetchone()
+            return True
+        except Exception:
+            return False
+
+    def _mark_closed_sync(self) -> None:
+        """event loop なしで同期的にコネクションをクローズ済みにマークする。
+
+        別ループからの dispose() 等で abandon する際に呼び出す。
+        SSL transport の ``__del__`` が閉じたループにアクセスしてエラーを起こすのを防ぐ。
+        デフォルト実装は何もしない。DBMS 固有の処理はサブクラスでオーバーライドする。
+        """
+        pass
+
+
+class ATAsyncDataSource:
+    """async 専用 DataSource の基底クラス。"""
+
+    def __init__(
+        self,
+        user: str | None = None,
+        password: str | None = None,
+        host: str | None = None,
+        port: int | None = None,
+        database: str | None = None,
+        **kwargs,
+    ) -> None:
+        self.user = user
+        self.password = password
+        self.host = host
+        self.port = port
+        self.database = database
+        self.dbconfig = kwargs
+        self.warning_handler: Callable | None = None
+        self.on_connect: Callable | None = None
+
+    def __str__(self) -> str:
+        port = f":{self.port}" if self.port else ""
+        return f"{self.__class__.__name__}({self.user}@{self.host}{port}/{self.database})"
+
+    async def __aenter__(self) -> "ATAsyncDataSource":
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        await self.dispose()
+
+    def connect(self) -> _AsyncConnectContext:
+        return _AsyncConnectContext(self._do_connect())
+
+    def begin(self) -> _AsyncConnectContext:
+        async def _begin() -> ATAsyncConnection:
+            conn = await self._do_connect()
+            conn.transaction = True
+            return conn
+
+        return _AsyncConnectContext(_begin())
+
+    async def _do_connect(self) -> "ATAsyncConnection":
+        conn = await self.__aconnect__()
+        if self.on_connect is not None:
+            if not inspect.iscoroutinefunction(self.on_connect):
+                await conn.close()
+                raise TypeError("on_connect must be an async function (async def) in ATAsyncDataSource")
+            try:
+                await self.on_connect(conn)
+            except Exception:
+                await conn.close()
+                raise
+        return conn
+
+    async def __aconnect__(self) -> "ATAsyncConnection":
+        raise NotImplementedError()
+
+    async def dispose(self) -> None:
+        pass
+
+
+class ATAsyncConnectionPool:
+    """asyncio ベースの非同期コネクションプール。
+
+    sync の ATConnectionPool と同じ Condition + LIFO deque 方式を使用する。
+    asyncio.Condition により、async/await で自然なブロッキングを実現する。
+
+    - ``max_connections``: idle + active の総数上限
+    - ``max_idle_connections``: アイドルキャッシュの上限（LIFO で管理）
+    - ``acquire_timeout``: ``acquire()`` の最大待機秒数。``None`` で無限待機
+    - ``pool_size``: エイリアス。``max_connections = max_idle_connections = pool_size`` に設定する
+    - ``name``: ロガー識別子。複数 DataSource を区別するために使用する（例: ``zeus/my_db``）
+    """
+
+    def __init__(
+        self,
+        max_connections: int = 10,
+        max_idle_connections: int = 5,
+        acquire_timeout: float | None = 30.0,
+        *,
+        pool_size: int | None = None,
+        name: str | None = None,
+    ) -> None:
+        if pool_size is not None:
+            max_connections = pool_size
+            max_idle_connections = pool_size
+        self.max_connections = max_connections
+        self.max_idle_connections = max_idle_connections
+        self.acquire_timeout = acquire_timeout
+        self._idles: collections.deque[ATAsyncConnection] = collections.deque()
+        self._active = 0  # idle + checked_out の総数
+        self._cond: asyncio.Condition | None = None  # 遅延初期化
+        self._app_loop: asyncio.AbstractEventLoop | None = None  # 接続を作ったループ（初回 acquire 時に記録）
+        self._log = logging.getLogger(f"atquery.pool.{name}") if name else _pool_log
+
+    def _get_condition(self) -> asyncio.Condition:
+        if self._cond is None:
+            self._cond = asyncio.Condition()
+        return self._cond
+
+    async def acquire(self) -> "ATAsyncConnection | None":
+        """コネクションを取得する。
+
+        Returns:
+            ATAsyncConnection | None: アイドルキャッシュから取得したコネクション、
+                                       または None（呼び出し元が新規作成する）
+
+        Raises:
+            ATConnectionPoolTimeoutError: timeout 経過した場合
+        """
+        cond = self._get_condition()
+        loop = asyncio.get_event_loop()
+        if self._app_loop is None:
+            self._app_loop = loop  # 初回 acquire でアプリのループを記録
+        deadline = (
+            None if self.acquire_timeout is None else loop.time() + self.acquire_timeout
+        )
+
+        async with cond:
+            while True:
+                # 1. idle を優先して取得（LIFO）
+                while self._idles:
+                    conn = self._idles.pop()
+                    if await conn.is_connected():
+                        self._log.debug("[acquire] reuse idle  active=%d idles=%d", self._active, len(self._idles))
+                        return conn
+                    # 切断済み: 除去してカウント調整
+                    try:
+                        await conn.connection.close()
+                    except Exception:
+                        pass  # 既に破棄済みでも除去は継続する
+                    self._active -= 1
+                    self._log.info("[acquire] drop stale  active=%d idles=%d", self._active, len(self._idles))
+
+                # 2. 新規作成できるならスロット確保して返す
+                if self._active < self.max_connections:
+                    self._active += 1
+                    self._log.debug("[acquire] new slot    active=%d idles=%d", self._active, len(self._idles))
+                    return None
+
+                # 3. 待機
+                self._log.info("[acquire] wait        active=%d idles=%d max=%d", self._active, len(self._idles), self.max_connections)
+                if deadline is not None:
+                    remaining = deadline - loop.time()
+                    if remaining <= 0:
+                        self._log.warning(
+                            "[acquire] TIMEOUT     active=%d idles=%d max=%d tasks=%s",
+                            self._active, len(self._idles), self.max_connections,
+                            [t.get_name() for t in asyncio.all_tasks()],
+                        )
+                        raise ATConnectionPoolTimeoutError(
+                            f"Connection pool timed out after {self.acquire_timeout}s "
+                            f"(max_connections={self.max_connections})"
+                        )
+                    try:
+                        await asyncio.wait_for(cond.wait(), remaining)
+                    except asyncio.TimeoutError:
+                        self._log.warning(
+                            "[acquire] TIMEOUT     active=%d idles=%d max=%d tasks=%s",
+                            self._active, len(self._idles), self.max_connections,
+                            [t.get_name() for t in asyncio.all_tasks()],
+                        )
+                        raise ATConnectionPoolTimeoutError(
+                            f"Connection pool timed out after {self.acquire_timeout}s "
+                            f"(max_connections={self.max_connections})"
+                        )
+                else:
+                    await cond.wait()
+
+    async def adiscard_slot(self) -> None:
+        """acquire() が None を返した後、接続作成に失敗した場合にスロットを返却する。
+
+        acquire() は新規接続スロット確保のために _active をインクリメントして None を返す。
+        呼び出し元 (DataSource.__aconnect__) で実際の DB 接続が失敗した場合は、
+        このメソッドを呼んでスロットを解放しないと _active が過剰カウントされ続け、
+        やがてプールが永久にブロックされる。
+        """
+        cond = self._get_condition()
+        async with cond:
+            self._active -= 1
+            self._log.info("[discard_slot]        active=%d idles=%d", self._active, len(self._idles))
+            cond.notify()
+
+    async def adiscard(self, conn: "ATAsyncConnection") -> None:
+        """コネクションを破棄する。
+
+        リセット失敗など、プールへの返却が不可能な場合に使用する。
+        """
+        cond = self._get_condition()
+        async with cond:
+            self._active -= 1
+            self._log.info("[discard]             active=%d idles=%d", self._active, len(self._idles))
+            cond.notify()
+        try:
+            await conn.connection.close()
+        except Exception:
+            pass  # ループが閉じている・既に切断済みの場合は無視
+
+    async def arelease(self, conn: "ATAsyncConnection") -> None:
+        """コネクションをアイドルキャッシュに返却する。
+
+        アイドルキャッシュが満杯の場合はコネクションをクローズする。
+        いずれの場合もウェイターに通知する。
+        """
+        close_conn = False
+        cond = self._get_condition()
+        async with cond:
+            if len(self._idles) < self.max_idle_connections:
+                self._idles.append(conn)
+                self._log.debug("[release] -> idle     active=%d idles=%d", self._active, len(self._idles))
+                cond.notify()
+            else:
+                self._active -= 1
+                self._log.debug("[release] -> close    active=%d idles=%d", self._active, len(self._idles))
+                cond.notify()
+                close_conn = True
+        if close_conn:
+            await conn.connection.close()
+
+    async def dispose(self) -> None:
+        """全アイドルコネクションをクローズする。
+
+        呼び出し元のイベントループが接続を作ったループと同一の場合（lifespan shutdown 等）は
+        COM_QUIT を送ってグレースフルクローズする。
+        別ループから呼ばれた場合（atexit + asyncio.run() 等）は接続を abandon し、
+        TCP ソケットの回収は OS に任せる。
+        """
+        current_loop = asyncio.get_running_loop()
+        same_loop = (self._app_loop is None or self._app_loop is current_loop)
+
+        conns_to_close: list[ATAsyncConnection] = []
+        cond = self._get_condition()
+        async with cond:
+            while self._idles:
+                conn = self._idles.pop()
+                self._active -= 1
+                if same_loop:
+                    conns_to_close.append(conn)
+                else:
+                    # 別ループ作成の接続 → abandon
+                    # SSL transport の __del__ がクローズ済みループを参照してエラーを起こさないよう
+                    # 同期的にクローズ済みとしてマークする（TCP は OS/MySQL wait_timeout に任せる）
+                    conn._mark_closed_sync()
+            cond.notify_all()
+
+        for conn in conns_to_close:
+            try:
+                await conn.connection.close()
+            except Exception:
+                pass