busya/ss-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

036: wire SupersetClientRegistry into translate + services flow

Browse Source 
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)
This commit is contained in:
busya
2026-06-05 15:14:44 +03:00
parent 4cef6af041
commit 7f1937f10b
9 changed files with 105 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

110
backend/src/core/utils/client_registry.py
View File

@@ -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)})

6
backend/src/plugins/debug.py
View File

@@ -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)

4
backend/src/plugins/translate/_run_source.py
View File

@@ -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,

4
backend/src/plugins/translate/preview_executor.py
View File

@@ -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,

4
backend/src/plugins/translate/service.py
View File

@@ -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:

6
backend/src/plugins/translate/service_datasource.py
View File

@@ -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", {})

13
backend/src/plugins/translate/superset_executor.py
View File

@@ -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",

5
backend/src/services/health_service.py
View File

@@ -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"),

8
backend/src/services/resource_service.py
View File

@@ -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