ss-tools/backend/src/core/utils/client_registry.py

# #region SupersetClientRegistryModule [C:4] [TYPE Module] [SEMANTICS async, superset, client, registry, semaphore]
# @BRIEF Singleton registry of long-lived async Superset clients per environment (env_id).
#   Stores shared httpx.AsyncClient + shared asyncio.Semaphore + asyncio.Lock for auth refresh.
# @LAYER Infrastructure
# @RELATION DEPENDS_ON -> [async_network.AsyncAPIClient]
# @RELATION DEPENDS_ON -> [config_models.Environment]
# @RELATION DEPENDS_ON -> [SupersetClient]
# @INVARIANT One client and one semaphore per env_id for the entire application lifetime.
# @INVARIANT All clients are closed on shutdown (calling shutdown()).
# @INVARIANT All consumers share a single SupersetClient per environment, sharing
#           connection pool, cookie jar (CSRF session), and semaphore.
# @RATIONALE Global per-env semaphore is impossible without a singleton registry.
#   Request-scoped clients lose connection pooling, CSRF cookies, and semaphore.
#   Shutdown without a registry leaves dangling connections.
# @REJECTED Per-request clients — lose connection pooling, semaphore becomes per-client.
#   Inline creation in each use-case — connection leaks, no shared CSRF cookie.
#   AsyncAPIClient-only registry (discarded by callers) — _detail_routes.py pattern
#   shows callers create SupersetClient anyway, defeating the purpose.

import asyncio
from typing import Any

from ..logger import logger
from .async_network import AsyncAPIClient


class _ClientSlot:
    """Internal slot holding a shared client and its resources."""
    def __init__(self, async_client: AsyncAPIClient, semaphore: asyncio.Semaphore, lock: asyncio.Lock):
        self.async_client = async_client
        self.semaphore = semaphore
        self.lock = lock


# Module-level registry — singleton dict.
_registry: dict[str, _ClientSlot] = {}
_create_lock = asyncio.Lock()


# #region _build_env_id [C:1] [TYPE Function]
# @BRIEF Build stable env_id from an Environment-like object.
def _build_env_id(env: Any) -> str:
    """Build stable identifier from environment config."""
    if hasattr(env, "id") and env.id:
        return f"{env.id}|{getattr(env, 'base_url', '')}"
    if isinstance(env, dict):
        return f"{env.get('base_url','')}|{env.get('username','')}"
    return str(id(env))


# #region get_client [C:3] [TYPE Function]
# @BRIEF Get or create an AsyncAPIClient for the given environment config.
#   Internal — consumers should use get_superset_client.
# @PRE  env is a valid Environment model or dict with base_url and auth.
# @POST Returns shared AsyncAPIClient instance.
# @SIDE_EFFECT  On first call: creates httpx.AsyncClient, semaphore, and auth lock.
async def get_client(
    env: Any,
    connection_pool_size: int = 20,
    request_timeout: int = 30,
) -> AsyncAPIClient:
    env_id = _build_env_id(env)
    # Fast path without lock.
    slot = _registry.get(env_id)
    if slot is not None:
        return slot.async_client
    # Slow path with lock — create once.
    async with _create_lock:
        slot = _registry.get(env_id)
        if slot is not None:
            return slot.async_client
        # Convert Environment model to config dict with the format expected by AsyncAPIClient
        if hasattr(env, "url") and hasattr(env, "username"):
            # Environment model — build config like SupersetClientBase.__init__ does
            config = {
                "base_url": env.url,
                "auth": {
                    "username": env.username,
                    "password": env.password,
                    "provider": "db",
                    "refresh": "true",
                },
            }
            verify_ssl = getattr(env, "verify_ssl", True)
            timeout = getattr(env, "timeout", request_timeout)
        elif hasattr(env, "model_dump"):
            config = env.model_dump()
            verify_ssl = config.get("verify_ssl", True)
            timeout = request_timeout
        elif isinstance(env, dict):
            config = env
            verify_ssl = config.get("verify_ssl", True)
            timeout = request_timeout
        else:
            config = {}
            verify_ssl = True
            timeout = request_timeout
        async_client = AsyncAPIClient(
            config=config,
            verify_ssl=verify_ssl,
            timeout=timeout,
        )
        semaphore = asyncio.Semaphore(connection_pool_size)
        lock = asyncio.Lock()
        _registry[env_id] = _ClientSlot(
            async_client=async_client, semaphore=semaphore, lock=lock
        )
        logger.reason("SupersetClientRegistry.get_client",
                      extra={"payload": {"env_id": env_id, "pool_size": connection_pool_size}})
        return async_client
# #endregion get_client


# #region get_superset_client [C:3] [TYPE Function]
# @BRIEF Get or create a shared SupersetClient for the given environment.
#   Returns a lightweight SupersetClient wrapping the shared AsyncAPIClient.
#   All callers across the project should use this instead of SupersetClient(env).
# @PRE  env is a valid Environment model.
# @POST Returns SupersetClient with shared httpx.AsyncClient, CSRF cookies, semaphore.
# @SIDE_EFFECT  On first call per env: creates httpx.AsyncClient, auth, CSRF cookie.
# @RELATION CALLS -> [SupersetClient]
# @RELATION DEPENDS_ON -> [get_client]
async def get_superset_client(
    env: Any,
    connection_pool_size: int = 20,
    request_timeout: int = 30,
) -> Any:
    """Get or create SupersetClient with shared AsyncAPIClient (connection pool + CSRF cookies)."""
    # Import lazily to avoid circular imports
    from ..superset_client import SupersetClient

    # Get or create the shared AsyncAPIClient (httpx + cookies + semaphore)
    shared_async = await get_client(env, connection_pool_size, request_timeout)

    # Wrap in a lightweight SupersetClient with the shared client
    return SupersetClient(env, client=shared_async)
# #endregion get_superset_client


# #region get_semaphore [C:2] [TYPE Function]
# @BRIEF Return the shared semaphore for the given env_config.
# @PRE  Client must have been created via get_client first.
# @POST Returns asyncio.Semaphore instance.
async def get_semaphore(env: Any) -> asyncio.Semaphore:
    env_id = _build_env_id(env)
    slot = _registry.get(env_id)
    if slot is None:
        await get_client(env)
        slot = _registry[env_id]
    return slot.semaphore


# #endregion get_semaphore


# #region get_auth_lock [C:2] [TYPE Function]
# @BRIEF Return the shared auth lock for the given environment.
# @PRE  Client must have been created via get_client first.
# @POST Returns asyncio.Lock instance for serializing auth refresh.
async def get_auth_lock(env: Any) -> asyncio.Lock:
    env_id = _build_env_id(env)
    slot = _registry.get(env_id)
    if slot is None:
        await get_client(env)
        slot = _registry[env_id]
    return slot.lock


# #endregion get_auth_lock


# #region shutdown [C:3] [TYPE Function]
# @BRIEF Close all async clients and clear the registry.
# @PRE  Application is shutting down.
# @POST All httpx.AsyncClient instances are closed. Registry cleared.
# @SIDE_EFFECT  Closes network connections for all envs.
async def shutdown() -> None:
    env_ids = list(_registry.keys())
    for env_id in env_ids:
        slot = _registry[env_id]
        try:
            await slot.async_client.aclose()
        except Exception as exc:
            logger.explore("SupersetClientRegistry.shutdown",
                           extra={"payload": {"env_id": env_id}, "error": str(exc)})
    _registry.clear()
    logger.reason("SupersetClientRegistry.shutdown",
                  extra={"payload": {"closed": len(env_ids)}})


# #endregion shutdown
# #endregion SupersetClientRegistryModule