pz-cli 5.0.0__py3-none-any.whl

pz_foundation/db/connection_probe.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import inspect
+from typing import Any
+
+from .lifecycle_status import DBConnectionLifecycleStatus, connection_lifecycle_status
+
+
+@dataclass(frozen=True, slots=True)
+class DBConnectionReadinessProbeResult:
+    alias: str
+    lifecycle: DBConnectionLifecycleStatus
+    ready: bool
+    query_executed: bool
+    transaction_safe: bool
+    status: str
+    error: BaseException | None = None
+    pool_status: dict[str, Any] = field(default_factory=dict)
+
+
+async def probe_connection_readiness(
+    tortoise_manager: Any,
+    *,
+    alias: str = "default",
+    query: str = "SELECT 1",
+) -> DBConnectionReadinessProbeResult:
+    lifecycle = connection_lifecycle_status(tortoise_manager)
+    if lifecycle.alias_status(alias) != "ok":
+        return DBConnectionReadinessProbeResult(
+            alias=alias,
+            lifecycle=lifecycle,
+            ready=False,
+            query_executed=False,
+            transaction_safe=True,
+            status=lifecycle.alias_status(alias),
+        )
+
+    try:
+        connection = tortoise_manager.get_connection(alias)
+    except Exception as exc:  # noqa: BLE001
+        return DBConnectionReadinessProbeResult(
+            alias=alias,
+            lifecycle=lifecycle,
+            ready=False,
+            query_executed=False,
+            transaction_safe=True,
+            status="connection_lookup_failed",
+            error=exc,
+        )
+
+    pool_status = describe_connection_pool(connection)
+    if not _connection_allows_non_transactional_probe(connection):
+        return DBConnectionReadinessProbeResult(
+            alias=alias,
+            lifecycle=lifecycle,
+            ready=False,
+            query_executed=False,
+            transaction_safe=False,
+            status="transactional_probe_not_allowed",
+            error=RuntimeError("database readiness probe requires an autocommit-safe connection"),
+            pool_status=pool_status,
+        )
+
+    try:
+        executor = getattr(connection, "execute_query")
+        result = executor(query)
+        if inspect.isawaitable(result):
+            await result
+        return DBConnectionReadinessProbeResult(
+            alias=alias,
+            lifecycle=lifecycle,
+            ready=True,
+            query_executed=True,
+            transaction_safe=True,
+            status="ok",
+            pool_status=describe_connection_pool(connection),
+        )
+    except Exception as exc:  # noqa: BLE001
+        return DBConnectionReadinessProbeResult(
+            alias=alias,
+            lifecycle=lifecycle,
+            ready=False,
+            query_executed=True,
+            transaction_safe=True,
+            status="query_failed",
+            error=exc,
+            pool_status=describe_connection_pool(connection),
+        )
+
+
+def _connection_allows_non_transactional_probe(connection: Any) -> bool:
+    template = getattr(connection, "_template", None)
+    if not isinstance(template, dict):
+        return True
+
+    kwargs = template.get("kwargs")
+    if isinstance(kwargs, dict) and "autocommit" in kwargs:
+        return bool(kwargs.get("autocommit"))
+
+    if "autocommit" in template:
+        return bool(template.get("autocommit"))
+
+    return True
+
+
+def describe_connection_pool(connection: Any) -> dict[str, Any]:
+    pool = getattr(connection, "_pool", None)
+    if pool is None:
+        return {"present": False}
+
+    payload: dict[str, Any] = {
+        "present": True,
+        "closed": bool(getattr(pool, "closed", False)),
+    }
+    for source_name, output_name in (
+        ("_min_size", "min_size"),
+        ("_max_size", "max_size"),
+        ("_nconns", "connections"),
+        ("_growing", "growing"),
+        ("max_waiting", "max_waiting"),
+    ):
+        value = getattr(pool, source_name, None)
+        if _is_scalar(value):
+            payload[output_name] = value
+
+    ready_queue = getattr(pool, "_pool", None)
+    waiting_queue = getattr(pool, "_waiting", None)
+    try:
+        payload["available"] = len(ready_queue) if ready_queue is not None else None
+    except Exception:
+        pass
+    try:
+        payload["waiting"] = len(waiting_queue) if waiting_queue is not None else None
+    except Exception:
+        pass
+    worker_tasks = tuple(getattr(pool, "_workers", ()) or ())
+    scheduler_task = getattr(pool, "_sched_runner", None)
+    maintenance_tasks = tuple(task for task in (*worker_tasks, scheduler_task) if task is not None)
+    if maintenance_tasks:
+        current_loop_id = _current_loop_id()
+        payload["workers"] = len(worker_tasks)
+        payload["workers_done"] = sum(1 for task in worker_tasks if _task_done(task))
+        payload["scheduler_done"] = _task_done(scheduler_task)
+        if current_loop_id is not None:
+            payload["worker_loop_matches"] = all(
+                _task_loop_id(task) in (None, current_loop_id) for task in maintenance_tasks
+            )
+
+    get_stats = getattr(pool, "get_stats", None)
+    if callable(get_stats):
+        try:
+            stats = get_stats()
+        except Exception:
+            stats = None
+        if isinstance(stats, dict):
+            sanitized = {
+                str(key): value
+                for key, value in stats.items()
+                if _is_scalar(value)
+            }
+            if sanitized:
+                payload["stats"] = sanitized
+    return payload
+
+
+def _is_scalar(value: Any) -> bool:
+    return value is None or isinstance(value, (bool, int, float, str))
+
+
+def _current_loop_id() -> int | None:
+    try:
+        import asyncio
+
+        return id(asyncio.get_running_loop())
+    except Exception:
+        return None
+
+
+def _task_done(task: Any) -> bool:
+    done = getattr(task, "done", None)
+    if not callable(done):
+        return False
+    try:
+        return bool(done())
+    except Exception:
+        return False
+
+
+def _task_loop_id(task: Any) -> int | None:
+    get_loop = getattr(task, "get_loop", None)
+    if not callable(get_loop):
+        return None
+    try:
+        return id(get_loop())
+    except Exception:
+        return None
+
+
+__all__ = [
+    "DBConnectionReadinessProbeResult",
+    "describe_connection_pool",
+    "probe_connection_readiness",
+]