Source code for sqlspec.observability._dispatcher

"""Lifecycle dispatcher for drivers and registry hooks."""

import inspect
from collections.abc import Callable, Iterable
from typing import Any, Literal

from sqlspec.utils.logging import get_logger

__all__ = ("LifecycleContext", "LifecycleDispatcher", "LifecycleHook")


logger = get_logger("sqlspec.observability.lifecycle")

LifecycleContext = dict[str, Any]
LifecycleHook = Callable[[LifecycleContext], Any]

LifecycleEvent = Literal[
    "on_pool_create",
    "on_pool_destroying",
    "on_pool_destroy",
    "on_connection_create",
    "on_connection_destroy",
    "on_session_start",
    "on_session_end",
    "on_query_start",
    "on_query_complete",
    "on_error",
]
EVENT_ATTRS: tuple[LifecycleEvent, ...] = (
    "on_pool_create",
    "on_pool_destroying",
    "on_pool_destroy",
    "on_connection_create",
    "on_connection_destroy",
    "on_session_start",
    "on_session_end",
    "on_query_start",
    "on_query_complete",
    "on_error",
)




class LifecycleDispatcher:
    """Dispatches lifecycle hooks with guard flags and diagnostics counters."""

    __slots__ = (
        "_counters",
        "_hooks",
        "_is_enabled",
        "has_connection_create",
        "has_connection_destroy",
        "has_error",
        "has_pool_create",
        "has_pool_destroy",
        "has_pool_destroying",
        "has_query_complete",
        "has_query_start",
        "has_session_end",
        "has_session_start",
    )



    def __init__(self, hooks: "dict[str, Iterable[LifecycleHook]] | None" = None) -> None:
        self.has_pool_create = False
        self.has_pool_destroying = False
        self.has_pool_destroy = False
        self.has_connection_create = False
        self.has_connection_destroy = False
        self.has_session_start = False
        self.has_session_end = False
        self.has_query_start = False
        self.has_query_complete = False
        self.has_error = False

        normalized: dict[LifecycleEvent, list[LifecycleHook]] = {}
        for event_name in EVENT_ATTRS:
            callables = hooks.get(event_name) if hooks else None
            normalized[event_name] = list(callables) if callables else []
            if normalized[event_name]:
                self._enable_event_guard(event_name)
        self._hooks: dict[LifecycleEvent, list[LifecycleHook]] = normalized
        self._counters: dict[LifecycleEvent, int] = dict.fromkeys(EVENT_ATTRS, 0)
        self._is_enabled = any(self._hooks.values())

@property def is_enabled(self) -> bool: """Return True when at least one hook is registered.""" return self._is_enabled

def emit_pool_create_sync(self, context: LifecycleContext) -> None: """Fire pool creation hooks synchronously.""" self._emit_sync("on_pool_create", context)