pz-cli 5.0.0__py3-none-any.whl

pz_foundation/db/session.py

@@ -0,0 +1,369 @@
+from __future__ import annotations
+
+import asyncio
+import functools
+from contextlib import asynccontextmanager
+from contextvars import ContextVar, Token
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable, Literal, Optional, TypeVar, cast
+
+from .config import build_foundation_tortoise_config
+from .connection_manager import _is_pool_exhausted_error, _is_transient_db_error
+from .tortoise_manager import tortoise_manager
+
+try:
+    from tortoise import connections
+except Exception:  # pragma: no cover
+    connections = None  # type: ignore[assignment]
+
+try:
+    from tortoise.transactions import in_transaction
+except Exception:  # pragma: no cover
+    in_transaction = None  # type: ignore[assignment]
+
+
+@dataclass(frozen=True)
+class DBRequestContext:
+    preference: Literal["read_only", "write"] = "write"
+    force_write: bool = False
+    prefer_read: bool = False
+    reason: Optional[str] = None
+
+    def should_use_read(self) -> bool:
+        return not self.force_write and (self.preference == "read_only" or self.prefer_read)
+
+
+_DEFAULT_DB_CONTEXT = DBRequestContext()
+_db_request_context: ContextVar[DBRequestContext] = ContextVar(
+    "pz_db_request_context",
+    default=_DEFAULT_DB_CONTEXT,
+)
+
+
+def _normalize_preference(preference: str) -> Literal["read_only", "write"]:
+    if isinstance(preference, str) and preference.lower() in {"read", "read_only"}:
+        return "read_only"
+    return "write"
+
+
+def get_request_context() -> DBRequestContext:
+    return _db_request_context.get()
+
+
+def push_request_context(
+    preference: str,
+    *,
+    force_write: bool = False,
+    prefer_read: bool = False,
+    reason: Optional[str] = None,
+) -> tuple[Token, DBRequestContext]:
+    context = DBRequestContext(
+        preference=_normalize_preference(preference),
+        force_write=force_write,
+        prefer_read=prefer_read,
+        reason=reason,
+    )
+    token = _db_request_context.set(context)
+    return token, context
+
+
+def reset_request_context(token: Token) -> None:
+    _db_request_context.reset(token)
+
+
+T = TypeVar("T")
+AsyncFunc = Callable[..., Any]
+TransactionFunc = Callable[[Any], Awaitable[T]]
+
+
+async def _ensure_connection_ready(alias: str) -> None:
+    await tortoise_manager.ensure_initialized(
+        config_builder=build_foundation_tortoise_config,
+        cluster_factory=lambda: None,
+    )
+    tortoise_manager.get_connection(alias)
+
+
+async def _reconnect(alias: str) -> None:
+    await tortoise_manager.ensure_initialized(
+        config_builder=build_foundation_tortoise_config,
+        cluster_factory=lambda: None,
+    )
+    await tortoise_manager.reconnect(alias)
+
+
+def _build_transaction_context(alias: str):
+    if in_transaction is not None:
+        return in_transaction(alias)
+
+    conn = tortoise_manager.get_connection(alias)
+    tx_factory = getattr(conn, "in_transaction", None)
+    if tx_factory is None:
+        @asynccontextmanager
+        async def _passthrough():
+            yield {"alias": alias, "mode": "transaction"}
+
+        return _passthrough()
+    return tx_factory()
+
+
+@asynccontextmanager
+async def _transaction(alias: str):
+    max_retries = 3
+    retry_delay = 0.5
+    tx_cm = None
+    tx = None
+
+    for attempt in range(max_retries):
+        await _ensure_connection_ready(alias)
+        tx_cm = _build_transaction_context(alias)
+        try:
+            tx = await tx_cm.__aenter__()
+            break
+        except Exception as exc:  # noqa: BLE001
+            if _is_pool_exhausted_error(exc) or not _is_transient_db_error(exc) or attempt == max_retries - 1:
+                raise
+            await _reconnect(alias)
+            await asyncio.sleep(retry_delay * (2**attempt))
+    else:  # pragma: no cover
+        raise RuntimeError(f"failed to open transaction for connection '{alias}'")
+
+    try:
+        yield tx
+    except BaseException as exc:
+        suppress = await tx_cm.__aexit__(type(exc), exc, exc.__traceback__)
+        if not suppress:
+            raise
+    else:
+        await tx_cm.__aexit__(None, None, None)
+
+
+async def run_in_transaction(
+    func: TransactionFunc[T],
+    *,
+    connection_name: str = "default",
+) -> T:
+    async with _transaction(connection_name) as conn:
+        return await func(conn)
+
+
+class ReadWriteDBSession:
+    def __init__(self) -> None:
+        self._has_read_connection = False
+        self._check_lock = asyncio.Lock()
+
+    @staticmethod
+    def _is_connection_closed(conn: Any) -> bool:
+        pool = getattr(conn, "_pool", None)
+        return bool(getattr(pool, "closed", False))
+
+    async def _check_read_connection(self) -> None:
+        if not self._has_read_connection:
+            async with self._check_lock:
+                if not self._has_read_connection:
+                    try:
+                        tortoise_manager.get_connection("read_only")
+                        self._has_read_connection = True
+                    except Exception:
+                        self._has_read_connection = False
+
+    async def _ensure_read_connection_available(self) -> bool:
+        return await self._ensure_connection_available("read_only")
+
+    async def _ensure_connection_available(self, alias: str) -> bool:
+        if alias == "read_only":
+            await self._check_read_connection()
+            if not self._has_read_connection:
+                return False
+        try:
+            await _ensure_connection_ready(alias)
+            conn = tortoise_manager.get_connection(alias)
+        except Exception:
+            if alias == "read_only":
+                self._has_read_connection = False
+            return False
+        if not self._is_connection_closed(conn):
+            return True
+        try:
+            await tortoise_manager.reconnect(alias)
+            conn = tortoise_manager.get_connection(alias)
+        except Exception:
+            if alias == "read_only":
+                self._has_read_connection = False
+            return False
+        if self._is_connection_closed(conn):
+            if alias == "read_only":
+                self._has_read_connection = False
+            return False
+        if alias == "read_only":
+            self._has_read_connection = True
+        return True
+
+    async def _ensure_write_connection_available(self) -> None:
+        await self._ensure_connection_available("default")
+
+    @asynccontextmanager
+    async def read_transaction(self):
+        context = get_request_context()
+        if not context.force_write and await self._ensure_read_connection_available():
+            try:
+                async with self._use_connection("read_only"):
+                    async with _transaction("read_only") as tx:
+                        yield tx
+                    return
+            except Exception:
+                pass
+        await self._ensure_write_connection_available()
+        async with self._use_connection("default"):
+            async with _transaction("default") as tx:
+                yield tx
+
+    @asynccontextmanager
+    async def write_transaction(self):
+        await self._ensure_write_connection_available()
+        async with self._use_connection("default"):
+            async with _transaction("default") as tx:
+                yield tx
+
+    async def execute_read(self, query: str, params: Optional[list] = None):
+        context = get_request_context()
+        if not context.force_write and await self._ensure_read_connection_available():
+            try:
+                conn = tortoise_manager.get_connection("read_only")
+                executor = getattr(conn, "execute_query", None)
+                async with self._use_connection("read_only"):
+                    if executor is None:
+                        return {"alias": "read_only", "query": query, "params": params or []}
+                    return await executor(query, params or [])
+            except Exception:
+                pass
+        await self._ensure_write_connection_available()
+        conn = tortoise_manager.get_connection("default")
+        executor = getattr(conn, "execute_query", None)
+        async with self._use_connection("default"):
+            if executor is None:
+                return {"alias": "default", "query": query, "params": params or []}
+            return await executor(query, params or [])
+
+    async def execute_write(self, query: str, params: Optional[list] = None):
+        await self._ensure_write_connection_available()
+        conn = tortoise_manager.get_connection("default")
+        executor = getattr(conn, "execute_query", None)
+        async with self._use_connection("default"):
+            if executor is None:
+                return {"alias": "default", "query": query, "params": params or []}
+            return await executor(query, params or [])
+
+    def using_read_db(self, func: AsyncFunc) -> AsyncFunc:
+        @functools.wraps(func)
+        async def wrapper(*args, **kwargs):
+            base_context = get_request_context()
+            token, context = push_request_context(
+                "read_only",
+                force_write=base_context.force_write,
+                prefer_read=True,
+                reason=f"@using_read_db:{func.__qualname__}",
+            )
+            try:
+                if context.should_use_read() and await self._ensure_read_connection_available():
+                    async with self._use_connection("read_only"):
+                        return await func(*args, **kwargs)
+                async with self._use_connection("default"):
+                    return await func(*args, **kwargs)
+            finally:
+                reset_request_context(token)
+
+        return cast(AsyncFunc, wrapper)
+
+    def using_write_db(self, func: AsyncFunc) -> AsyncFunc:
+        @functools.wraps(func)
+        async def wrapper(*args, **kwargs):
+            token, _ = push_request_context(
+                "write",
+                force_write=True,
+                prefer_read=False,
+                reason=f"@using_write_db:{func.__qualname__}",
+            )
+            try:
+                await self._ensure_write_connection_available()
+                async with self._use_connection("default"):
+                    return await func(*args, **kwargs)
+            finally:
+                reset_request_context(token)
+
+        return cast(AsyncFunc, wrapper)
+
+    @asynccontextmanager
+    async def _use_connection(self, alias: str):
+        if connections is None:
+            yield
+            return
+        try:
+            conn = tortoise_manager.get_connection(alias)
+        except KeyError:
+            yield
+            return
+        try:
+            token = connections.set(alias, conn)
+        except Exception:
+            token = None
+        try:
+            yield
+        finally:
+            if token is not None:
+                try:
+                    connections.reset(token)
+                except Exception:
+                    pass
+
+    @asynccontextmanager
+    async def request_scope(
+        self,
+        preference: str,
+        *,
+        force_write: bool = False,
+        prefer_read: bool = False,
+        reason: Optional[str] = None,
+    ):
+        token, context = push_request_context(
+            preference,
+            force_write=force_write,
+            prefer_read=prefer_read,
+            reason=reason,
+        )
+        try:
+            use_read = False
+            if not context.force_write:
+                use_read = context.should_use_read() and await self._ensure_read_connection_available()
+
+            if use_read:
+                async with self._use_connection("read_only"):
+                    yield "read_only"
+            else:
+                await self._ensure_write_connection_available()
+                async with self._use_connection("default"):
+                    yield "write"
+        finally:
+            reset_request_context(token)
+
+
+db_session = ReadWriteDBSession()
+read_transaction = db_session.read_transaction
+write_transaction = db_session.write_transaction
+using_read_db = db_session.using_read_db
+using_write_db = db_session.using_write_db
+
+
+__all__ = [
+    "DBRequestContext",
+    "ReadWriteDBSession",
+    "db_session",
+    "get_request_context",
+    "push_request_context",
+    "read_transaction",
+    "reset_request_context",
+    "run_in_transaction",
+    "using_read_db",
+    "using_write_db",
+    "write_transaction",
+]

pz_foundation/db/sqlite.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def build_sqlite_connections(settings) -> dict[str, Any]:
+    return {
+        "default": {
+            "engine": "tortoise.backends.sqlite",
+            "credentials": {"file_path": getattr(settings, "path", "") or ":memory:"},
+        }
+    }
+
+
+__all__ = ["build_sqlite_connections"]