From 7f1937f10b92df67a34f8cf0dc7c681a3b16b5c9 Mon Sep 17 00:00:00 2001 From: busya Date: Fri, 5 Jun 2026 15:14:44 +0300 Subject: [PATCH] 036: wire SupersetClientRegistry into translate + services flow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces direct SupersetClient(env) calls with shared get_superset_client(env) from client_registry. Changed files (9): - client_registry.py: added get_superset_client(), _build_env_id(), accepts Environment models (not just dicts), fixed async_client attr - superset_executor.py: _get_client() now async, uses shared client - preview_executor.py, _run_source.py, service.py, service_datasource.py - health_service.py, resource_service.py, debug.py Impact per environment: - 6 separate httpx.AsyncClient instances → 1 shared client - 6 CSRF cookie fetches → 1 (on first access) - 6 connection pools → 1 shared pool - Shared semaphore for backpressure - Shared cookie jar (fixes CSRF 'tokens do not match' on SQL Lab execute) --- backend/src/core/utils/client_registry.py | 110 +++++++++++++----- backend/src/plugins/debug.py | 6 +- backend/src/plugins/translate/_run_source.py | 4 +- .../src/plugins/translate/preview_executor.py | 4 +- backend/src/plugins/translate/service.py | 4 +- .../plugins/translate/service_datasource.py | 6 +- .../plugins/translate/superset_executor.py | 13 ++- backend/src/services/health_service.py | 5 +- backend/src/services/resource_service.py | 8 +- 9 files changed, 105 insertions(+), 55 deletions(-) diff --git a/backend/src/core/utils/client_registry.py b/backend/src/core/utils/client_registry.py index 45575660..088c0992 100644 --- a/backend/src/core/utils/client_registry.py +++ b/backend/src/core/utils/client_registry.py @@ -4,13 +4,18 @@ # @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. Shutdown without a registry leaves -# dangling connections. +# 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. +# 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 @@ -20,9 +25,9 @@ from .async_network import AsyncAPIClient class _ClientSlot: - """Internal slot holding a client and its shared resources.""" - def __init__(self, client: AsyncAPIClient, semaphore: asyncio.Semaphore, lock: asyncio.Lock): - self.client = client + """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 @@ -32,56 +37,99 @@ _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. -# Client is created lazily on first access. Subsequent calls return the same client. -# @PRE env_config is a valid config dict with base_url and auth. -# @POST Returns AsyncAPIClient instance with shared semaphore and auth lock. +# 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. -# @RELATION CALLS -> [AsyncAPIClient] async def get_client( - env_config: dict[str, Any], + env: Any, connection_pool_size: int = 20, request_timeout: int = 30, ) -> AsyncAPIClient: - # Build stable env_id from config. - env_id = f"{env_config.get('base_url','')}|{env_config.get('auth',{}).get('username','')}" + env_id = _build_env_id(env) # Fast path without lock. slot = _registry.get(env_id) if slot is not None: - return slot.client + 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.client - verify_ssl = env_config.get("verify_ssl", True) - client = AsyncAPIClient( - config=env_config, + return slot.async_client + # Convert Environment model or dict to config dict + if hasattr(env, "model_dump"): + config = env.model_dump() + elif hasattr(env, "dict"): + config = env.dict() + elif isinstance(env, dict): + config = env + else: + config = {} + verify_ssl = config.get("verify_ssl", True) + async_client = AsyncAPIClient( + config=config, verify_ssl=verify_ssl, timeout=request_timeout, ) semaphore = asyncio.Semaphore(connection_pool_size) lock = asyncio.Lock() - _registry[env_id] = _ClientSlot(client=client, semaphore=semaphore, lock=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 client - - + 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_config: dict[str, Any]) -> asyncio.Semaphore: - env_id = f"{env_config.get('base_url','')}|{env_config.get('auth',{}).get('username','')}" +async def get_semaphore(env: Any) -> asyncio.Semaphore: + env_id = _build_env_id(env) slot = _registry.get(env_id) if slot is None: - # Create the client first, which creates the semaphore. - await get_client(env_config) + await get_client(env) slot = _registry[env_id] return slot.semaphore @@ -90,14 +138,14 @@ async def get_semaphore(env_config: dict[str, Any]) -> asyncio.Semaphore: # #region get_auth_lock [C:2] [TYPE Function] -# @BRIEF Return the shared auth lock for the given env_config. +# @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_config: dict[str, Any]) -> asyncio.Lock: - env_id = f"{env_config.get('base_url','')}|{env_config.get('auth',{}).get('username','')}" +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_config) + await get_client(env) slot = _registry[env_id] return slot.lock @@ -115,7 +163,7 @@ async def shutdown() -> None: for env_id in env_ids: slot = _registry[env_id] try: - await slot.client.aclose() + await slot.async_client.aclose() except Exception as exc: logger.explore("SupersetClientRegistry.shutdown", extra={"payload": {"env_id": env_id}, "error": str(exc)}) diff --git a/backend/src/plugins/debug.py b/backend/src/plugins/debug.py index 286425b3..a739c979 100644 --- a/backend/src/plugins/debug.py +++ b/backend/src/plugins/debug.py @@ -8,8 +8,8 @@ from typing import Any from ..core.logger import belief_scope, logger from ..core.plugin_base import PluginBase -from ..core.superset_client import SupersetClient from ..core.task_manager.context import TaskContext +from ..core.utils.client_registry import get_superset_client # #region DebugPlugin [TYPE Class] @@ -166,7 +166,7 @@ class DebugPlugin(PluginBase): log.error(f"Environment '{name}' not found.") raise ValueError(f"Environment '{name}' not found.") - client = SupersetClient(env_config) + client = await get_superset_client(env_config) await client.authenticate() count, dbs = await client.get_databases() log.debug(f"Found {count} databases in {name}") @@ -202,7 +202,7 @@ class DebugPlugin(PluginBase): log.error(f"Environment '{env_name}' not found.") raise ValueError(f"Environment '{env_name}' not found.") - client = SupersetClient(env_config) + client = await get_superset_client(env_config) await client.authenticate() dataset_response = await client.get_dataset(dataset_id) diff --git a/backend/src/plugins/translate/_run_source.py b/backend/src/plugins/translate/_run_source.py index ebcb0a53..c6a532fc 100644 --- a/backend/src/plugins/translate/_run_source.py +++ b/backend/src/plugins/translate/_run_source.py @@ -43,8 +43,8 @@ async def fetch_source_rows(db: Session, config_manager: ConfigManager, job_id: env_config = environments[0] if env_config: - from ...core.superset_client import SupersetClient - client = SupersetClient(env_config) + from ...core.utils.client_registry import get_superset_client + client = await get_superset_client(env_config) dataset_detail = await client.get_dataset_detail(int(job.source_datasource_id)) query_context = client.build_dataset_preview_query_context( dataset_id=int(job.source_datasource_id), dataset_record=dataset_detail, diff --git a/backend/src/plugins/translate/preview_executor.py b/backend/src/plugins/translate/preview_executor.py index 7faf2d80..5f978d95 100644 --- a/backend/src/plugins/translate/preview_executor.py +++ b/backend/src/plugins/translate/preview_executor.py @@ -16,7 +16,7 @@ from sqlalchemy.orm import Session from ...core.config_manager import ConfigManager from ...core.logger import belief_scope, logger -from ...core.superset_client import SupersetClient +from ...core.utils.client_registry import get_superset_client from ...models.translate import TranslationJob from ._llm_async_http import call_openai_compatible from .preview_resolve_provider import resolve_provider_model as _resolve_provider_model @@ -54,7 +54,7 @@ class PreviewExecutor: else: raise ValueError("No Superset environments configured") - client = SupersetClient(env_config) + client = await get_superset_client(env_config) dataset_detail = await client.get_dataset_detail(int(job.source_datasource_id)) query_context = client.build_dataset_preview_query_context( dataset_id=int(job.source_datasource_id), dataset_record=dataset_detail, diff --git a/backend/src/plugins/translate/service.py b/backend/src/plugins/translate/service.py index d578adb8..d3178444 100644 --- a/backend/src/plugins/translate/service.py +++ b/backend/src/plugins/translate/service.py @@ -235,11 +235,11 @@ class TranslateJobService: # region fetch_available_datasources [TYPE Function] # @PURPOSE: List available Superset datasets for translation job creation. async def fetch_available_datasources(self, env_id: str, search: str | None = None) -> list: - from ...core.superset_client import SupersetClient + from ...core.utils.client_registry import get_superset_client env = self.config_manager.get_environment(env_id) if not env: raise ValueError(f"Environment '{env_id}' not found") - client = SupersetClient(env) + client = await get_superset_client(env) _, datasets = await client.get_datasets() result = [] for ds in datasets: diff --git a/backend/src/plugins/translate/service_datasource.py b/backend/src/plugins/translate/service_datasource.py index a2762ec0..8ce7cd12 100644 --- a/backend/src/plugins/translate/service_datasource.py +++ b/backend/src/plugins/translate/service_datasource.py @@ -10,7 +10,7 @@ from typing import Any from ...core.config_manager import ConfigManager from ...core.logger import logger -from ...core.superset_client import SupersetClient +from ...core.utils.client_registry import get_superset_client from ...schemas.translate import DatasourceColumnResponse, DatasourceColumnsResponse # Supported database dialects for translation @@ -68,7 +68,7 @@ async def fetch_datasource_metadata( if not env_config: raise ValueError(f"Superset environment '{env_id}' not found in configuration") - client = SupersetClient(env_config) + client = await get_superset_client(env_config) dataset_detail = await client.get_dataset_detail(dataset_id) raw_columns = dataset_detail.get("columns", []) @@ -127,7 +127,7 @@ async def get_datasource_columns( if not env_config: raise ValueError(f"Superset environment '{env_id}' not found") - client = SupersetClient(env_config) + client = await get_superset_client(env_config) dataset_detail = await client.get_dataset_detail(datasource_id) database_info = dataset_detail.get("database", {}) diff --git a/backend/src/plugins/translate/superset_executor.py b/backend/src/plugins/translate/superset_executor.py index 183f64ac..28e0e1d7 100644 --- a/backend/src/plugins/translate/superset_executor.py +++ b/backend/src/plugins/translate/superset_executor.py @@ -17,6 +17,7 @@ import uuid from ...core.config_manager import ConfigManager from ...core.logger import belief_scope, logger from ...core.superset_client import SupersetClient +from ...core.utils.client_registry import get_superset_client # #region SupersetSqlLabExecutor [C:4] [TYPE Class] @@ -39,7 +40,7 @@ class SupersetSqlLabExecutor: # @PRE env_id must correspond to a valid environment config. # @POST Returns authenticated SupersetClient. # @SIDE_EFFECT Authenticates against Superset API. - def _get_client(self) -> SupersetClient: + async def _get_client(self) -> SupersetClient: with belief_scope("SupersetSqlLabExecutor._get_client"): if self._client is not None: return self._client @@ -49,7 +50,7 @@ class SupersetSqlLabExecutor: if not env_config: raise ValueError(f"Superset environment '{self.env_id}' not found") - self._client = SupersetClient(env_config) + self._client = await get_superset_client(env_config) return self._client # endregion _get_client @@ -64,7 +65,7 @@ class SupersetSqlLabExecutor: target_database_id: str | None = None, ) -> int: with belief_scope("SupersetSqlLabExecutor.resolve_database_id"): - client = self._get_client() + client = await self._get_client() # If a specific target_database_id is provided, use it directly if target_database_id: @@ -153,7 +154,7 @@ class SupersetSqlLabExecutor: run_async: bool = True, ) -> dict[str, Any]: with belief_scope("SupersetSqlLabExecutor.execute_sql"): - client = self._get_client() + client = await self._get_client() db_id = database_id or self._database_id if not db_id: db_id = await self.resolve_database_id() @@ -265,7 +266,7 @@ class SupersetSqlLabExecutor: poll_interval_seconds: float = 2.0, ) -> dict[str, Any]: with belief_scope("SupersetSqlLabExecutor.poll_execution_status"): - client = self._get_client() + client = await self._get_client() logger.reason("Polling SQL execution status", { "query_id": query_id, @@ -408,7 +409,7 @@ class SupersetSqlLabExecutor: # @SIDE_EFFECT Makes HTTP GET to Superset. async def get_query_results(self, query_id: str) -> dict[str, Any]: with belief_scope("SupersetSqlLabExecutor.get_query_results"): - client = self._get_client() + client = await self._get_client() try: response = await client.network.request( method="GET", diff --git a/backend/src/services/health_service.py b/backend/src/services/health_service.py index a019af32..0cd0cf68 100644 --- a/backend/src/services/health_service.py +++ b/backend/src/services/health_service.py @@ -16,8 +16,8 @@ from sqlalchemy import func from sqlalchemy.orm import Session from ..core.logger import logger -from ..core.superset_client import SupersetClient from ..core.task_manager import TaskManager +from ..core.utils.client_registry import get_superset_client from ..core.task_manager.cleanup import TaskCleanupService from ..models.llm import ValidationRecord from ..schemas.health import DashboardHealthItem, HealthSummaryResponse @@ -121,7 +121,8 @@ class HealthService: ) dashboard_meta_map = cached_meta_data[1] else: - dashboards = await SupersetClient(env).get_dashboards_summary() + client = await get_superset_client(env) + dashboards = await client.get_dashboards_summary() dashboard_meta_map = { str(item.get("id")): { "slug": item.get("slug"), diff --git a/backend/src/services/resource_service.py b/backend/src/services/resource_service.py index 9cb723ec..cd3a3205 100644 --- a/backend/src/services/resource_service.py +++ b/backend/src/services/resource_service.py @@ -14,8 +14,8 @@ from datetime import UTC, datetime from typing import Any from ..core.logger import belief_scope, logger -from ..core.superset_client import SupersetClient from ..core.task_manager.models import Task +from ..core.utils.client_registry import get_superset_client from ..services.git_service import GitService @@ -53,7 +53,7 @@ class ResourceService: require_slug: bool = False, ) -> list[dict[str, Any]]: with belief_scope("get_dashboards_with_status", f"env={env.id}"): - client = SupersetClient(env) + client = await get_superset_client(env) dashboards = await client.get_dashboards_summary(require_slug=require_slug) # Enhance each dashboard with Git status and task status @@ -110,7 +110,7 @@ class ResourceService: "get_dashboards_page_with_status", f"env={env.id}, page={page}, page_size={page_size}, search={search}", ): - client = SupersetClient(env) + client = await get_superset_client(env) total, dashboards_page = await client.get_dashboards_summary_page( page=page, page_size=page_size, @@ -302,7 +302,7 @@ class ResourceService: tasks: list[Task] | None = None ) -> list[dict[str, Any]]: with belief_scope("get_datasets_with_status", f"env={env.id}"): - client = SupersetClient(env) + client = await get_superset_client(env) datasets = await client.get_datasets_summary() # Enhance each dataset with task status and linked dashboard count