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

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