busya
ee9123bcf2
Full async conversion for all sync callers of async SupersetClient methods: orchestrator_sql.py: generate_and_insert_sql + _resolve_dialect → async orchestrator_run_completion.py: complete_success → async (calls generate_and_insert_sql) orchestrator_exec.py: execute_run → async (awaits complete_success) orchestrator_runner.py: execute_run → async (delegates to engine) orchestrator.py: execute_run + _generate_and_insert_sql → async executor.py: _insert_batch_to_target → async (awaits batch insert) _batch_proc.py: insert_batch_to_target → async _batch_insert.py: insert_batch_to_target + _resolve_insert_backend + _execute_insert_sql → async dataset_mapper.py: get_sqllab_mappings + run_mapping → async + await get_dataset, update_dataset, execute_and_poll mapper.py: await on run_mapping + resolve_database_id calls _run_routes.py: threading.Thread → asyncio.create_task (_background_execute async)
239 lines
13 KiB
Python
239 lines
13 KiB
Python
# #region DatasetMapperModule [TYPE Module] [SEMANTICS dataset, superset, dataset-mapper, sqllab]
|
||
#
|
||
# @BRIEF Этот модуль отвечает за обновление метаданных (verbose_name) в датасетах Superset,
|
||
# извлекая их из Superset SQL Lab (любая БД, подключённая к Superset) или XLSX-файлов.
|
||
# @LAYER Domain
|
||
# @RELATION DEPENDS_ON -> backend.core.superset_client
|
||
# @RELATION DEPENDS_ON -> [EXT:Library:pandas]
|
||
# @PUBLIC_API DatasetMapper
|
||
from typing import Any
|
||
|
||
import pandas as pd # type: ignore
|
||
|
||
from ..logger import belief_scope, logger as app_logger
|
||
|
||
|
||
# #region DatasetMapper [TYPE Class]
|
||
# @BRIEF Класс для маппинга и обновления verbose_name в датасетах Superset.
|
||
class DatasetMapper:
|
||
# #region __init__ [TYPE Function]
|
||
# @PURPOSE: Initializes the mapper.
|
||
# @POST Объект DatasetMapper инициализирован.
|
||
def __init__(self):
|
||
pass
|
||
# #endregion __init__
|
||
|
||
# #region get_sqllab_mappings [TYPE Function]
|
||
# @PURPOSE: Извлекает маппинги column_name -> verbose_name через SQL Lab Superset.
|
||
# @PRE sqllab_executor должен быть инициализирован с database_id.
|
||
# @PRE dataset_id должен существовать в Superset.
|
||
# @POST Возвращается словарь column_name -> verbose_name из результатов SQL-запроса.
|
||
# Если sql_query не указан, строится запрос к information_schema.columns.
|
||
# @PARAM client (SupersetClient) - Авторизованный клиент Superset.
|
||
# @PARAM dataset_id (int) - ID датасета (для получения table_name/schema).
|
||
# @PARAM sqllab_executor (SupersetSqlLabExecutor) - Исполнитель SQL Lab.
|
||
# @PARAM database_id (int) - ID базы данных в Superset.
|
||
# @PARAM sql_query (Optional[str]) - Произвольный SQL-запрос (должен вернуть column_name + verbose_name).
|
||
# @RETURN Dict[str, str] - Словарь column_name -> verbose_name.
|
||
async def get_sqllab_mappings(
|
||
self,
|
||
client: Any,
|
||
dataset_id: int,
|
||
sqllab_executor: Any,
|
||
database_id: int,
|
||
sql_query: str | None = None,
|
||
) -> dict[str, str]:
|
||
with belief_scope("Fetch mappings via SQL Lab"):
|
||
# Получаем датасет из Superset, чтобы узнать table_name и schema
|
||
dataset_response = client.get_dataset(dataset_id)
|
||
dataset_data = dataset_response.get("result", dataset_response)
|
||
table_name = dataset_data.get("table_name")
|
||
table_schema = dataset_data.get("schema") or "public"
|
||
database_name = dataset_data.get("database", {}).get("database_name", "")
|
||
|
||
app_logger.info(
|
||
"[get_sqllab_mappings] Dataset %d: %s.%s (db: %s)",
|
||
dataset_id, table_schema, table_name, database_name
|
||
)
|
||
|
||
# Если SQL-запрос не указан, строим дефолтный для information_schema
|
||
if not sql_query:
|
||
sql_query = (
|
||
"SELECT column_name, pg_catalog.col_description(\n"
|
||
" (quote_ident(table_schema) || '.' || quote_ident(table_name))::regclass::oid,\n"
|
||
" ordinal_position\n"
|
||
") AS verbose_name\n"
|
||
f"FROM information_schema.columns\n"
|
||
f"WHERE table_schema = '{table_schema}' AND table_name = '{table_name}';\n"
|
||
)
|
||
app_logger.info("[get_sqllab_mappings] Using default information_schema query for %s.%s", table_schema, table_name)
|
||
else:
|
||
app_logger.info("[get_sqllab_mappings] Using custom SQL query")
|
||
|
||
# Выполняем запрос через SQL Lab
|
||
app_logger.info("[get_sqllab_mappings] Executing SQL Lab query on database %d...", database_id)
|
||
result = await sqllab_executor.execute_and_poll(sql_query, database_id)
|
||
rows = result.get("results") or result.get("data") or result.get("result", [])
|
||
app_logger.info("[get_sqllab_mappings] SQL Lab returned %d rows", len(rows))
|
||
|
||
mappings: dict[str, str] = {}
|
||
for row in rows:
|
||
col_name = row.get("column_name")
|
||
verbose_name = row.get("verbose_name")
|
||
if col_name and verbose_name:
|
||
mappings[col_name] = str(verbose_name)
|
||
|
||
app_logger.info("[get_sqllab_mappings] Extracted %d mappings", len(mappings))
|
||
return mappings
|
||
# #endregion get_sqllab_mappings
|
||
|
||
# #region load_excel_mappings [TYPE Function]
|
||
# @PURPOSE: Загружает маппинги column_name -> verbose_name из XLSX файла.
|
||
# @PRE file_path должен указывать на существующий XLSX файл.
|
||
# @POST Возвращается словарь с маппингами из файла.
|
||
# @THROW: Exception - При ошибках чтения файла или парсинга.
|
||
# @PARAM file_path (str) - Путь к XLSX файлу.
|
||
# @RETURN Dict[str, str] - Словарь с маппингами.
|
||
def load_excel_mappings(self, file_path: str) -> dict[str, str]:
|
||
with belief_scope("Load mappings from Excel"):
|
||
app_logger.info("[load_excel_mappings][Enter] Loading mappings from %s.", file_path)
|
||
try:
|
||
df = pd.read_excel(file_path)
|
||
mappings = df.set_index('column_name')['verbose_name'].to_dict()
|
||
app_logger.info("[load_excel_mappings][Success] Loaded %d mappings.", len(mappings))
|
||
return mappings
|
||
except Exception as e:
|
||
app_logger.error("[load_excel_mappings][Failure] %s", e, exc_info=True)
|
||
raise
|
||
# #endregion load_excel_mappings
|
||
|
||
# #region run_mapping [TYPE Function]
|
||
# @PURPOSE: Основная функция для выполнения маппинга и обновления verbose_name датасета в Superset.
|
||
# @PRE superset_client должен быть авторизован.
|
||
# @PRE dataset_id должен быть существующим ID в Superset.
|
||
# @POST Если найдены изменения, датасет в Superset обновлен через API.
|
||
# @RELATION CALLS -> [EXT:method:DatasetMapper.get_sqllab_mappings]
|
||
# @RELATION CALLS -> [EXT:method:DatasetMapper.load_excel_mappings]
|
||
# @RELATION CALLS -> [EXT:method:SupersetClient.get_dataset]
|
||
# @RELATION CALLS -> [EXT:method:SupersetClient.update_dataset]
|
||
# @PARAM superset_client (Any) - Клиент Superset.
|
||
# @PARAM dataset_id (int) - ID датасета для обновления.
|
||
# @PARAM source (str) - Источник данных ('sqllab' или 'excel').
|
||
# @PARAM sqllab_executor (Optional[Any]) - Исполнитель SQL Lab (для sqllab source).
|
||
# @PARAM database_id (Optional[int]) - ID базы данных Superset (для sqllab source).
|
||
# @PARAM sql_query (Optional[str]) - Произвольный SQL-запрос (для sqllab source).
|
||
# @PARAM excel_path (Optional[str]) - Путь к XLSX файлу.
|
||
async def run_mapping(
|
||
self,
|
||
superset_client: Any,
|
||
dataset_id: int,
|
||
source: str,
|
||
sqllab_executor: Any = None,
|
||
database_id: int | None = None,
|
||
sql_query: str | None = None,
|
||
excel_path: str | None = None,
|
||
):
|
||
with belief_scope(f"Run dataset mapping for ID {dataset_id}"):
|
||
app_logger.info("[run_mapping][Enter] Starting dataset mapping for ID %d from source '%s'.", dataset_id, source)
|
||
mappings: dict[str, str] = {}
|
||
|
||
try:
|
||
if source == "sqllab":
|
||
assert sqllab_executor and database_id, "sqllab_executor and database_id are required for sqllab source."
|
||
mappings.update(await self.get_sqllab_mappings(superset_client, dataset_id, sqllab_executor, database_id, sql_query))
|
||
elif source == "excel":
|
||
assert excel_path, "excel_path is required."
|
||
mappings.update(self.load_excel_mappings(excel_path))
|
||
else:
|
||
app_logger.error("[run_mapping][Failure] Invalid source: %s.", source)
|
||
return
|
||
|
||
dataset_response = await superset_client.get_dataset(dataset_id)
|
||
dataset_data = dataset_response['result']
|
||
|
||
original_columns = dataset_data.get('columns', [])
|
||
updated_columns = []
|
||
changes_made = False
|
||
for column in original_columns:
|
||
col_name = column.get('column_name')
|
||
|
||
new_column = {
|
||
"column_name": col_name,
|
||
"id": column.get("id"),
|
||
"advanced_data_type": column.get("advanced_data_type"),
|
||
"description": column.get("description"),
|
||
"expression": column.get("expression"),
|
||
"extra": column.get("extra"),
|
||
"filterable": column.get("filterable"),
|
||
"groupby": column.get("groupby"),
|
||
"is_active": column.get("is_active"),
|
||
"is_dttm": column.get("is_dttm"),
|
||
"python_date_format": column.get("python_date_format"),
|
||
"type": column.get("type"),
|
||
"uuid": column.get("uuid"),
|
||
"verbose_name": column.get("verbose_name"),
|
||
}
|
||
|
||
new_column = {k: v for k, v in new_column.items() if v is not None}
|
||
if col_name in mappings:
|
||
mapping_value = mappings[col_name]
|
||
if isinstance(mapping_value, str) and new_column.get('verbose_name') != mapping_value:
|
||
new_column['verbose_name'] = mapping_value
|
||
changes_made = True
|
||
|
||
updated_columns.append(new_column)
|
||
|
||
updated_metrics = []
|
||
for metric in dataset_data.get("metrics", []):
|
||
new_metric = {
|
||
"id": metric.get("id"),
|
||
"metric_name": metric.get("metric_name"),
|
||
"expression": metric.get("expression"),
|
||
"verbose_name": metric.get("verbose_name"),
|
||
"description": metric.get("description"),
|
||
"d3format": metric.get("d3format"),
|
||
"currency": metric.get("currency"),
|
||
"extra": metric.get("extra"),
|
||
"warning_text": metric.get("warning_text"),
|
||
"metric_type": metric.get("metric_type"),
|
||
"uuid": metric.get("uuid"),
|
||
}
|
||
updated_metrics.append({k: v for k, v in new_metric.items() if v is not None})
|
||
|
||
if changes_made:
|
||
payload_for_update = {
|
||
"database_id": dataset_data.get("database", {}).get("id"),
|
||
"table_name": dataset_data.get("table_name"),
|
||
"schema": dataset_data.get("schema"),
|
||
"columns": updated_columns,
|
||
"owners": [owner["id"] for owner in dataset_data.get("owners", [])],
|
||
"metrics": updated_metrics,
|
||
"extra": dataset_data.get("extra"),
|
||
"description": dataset_data.get("description"),
|
||
"sql": dataset_data.get("sql"),
|
||
"cache_timeout": dataset_data.get("cache_timeout"),
|
||
"catalog": dataset_data.get("catalog"),
|
||
"default_endpoint": dataset_data.get("default_endpoint"),
|
||
"external_url": dataset_data.get("external_url"),
|
||
"fetch_values_predicate": dataset_data.get("fetch_values_predicate"),
|
||
"filter_select_enabled": dataset_data.get("filter_select_enabled"),
|
||
"is_managed_externally": dataset_data.get("is_managed_externally"),
|
||
"is_sqllab_view": dataset_data.get("is_sqllab_view"),
|
||
"main_dttm_col": dataset_data.get("main_dttm_col"),
|
||
"normalize_columns": dataset_data.get("normalize_columns"),
|
||
"offset": dataset_data.get("offset"),
|
||
"template_params": dataset_data.get("template_params"),
|
||
}
|
||
|
||
payload_for_update = {k: v for k, v in payload_for_update.items() if v is not None}
|
||
await superset_client.update_dataset(dataset_id, payload_for_update)
|
||
app_logger.info("[run_mapping][Success] Dataset %d columns' verbose_name updated.", dataset_id)
|
||
else:
|
||
app_logger.info("[run_mapping][State] No changes in columns' verbose_name, skipping update.")
|
||
except (AssertionError, FileNotFoundError, Exception) as e:
|
||
app_logger.error("[run_mapping][Failure] %s", e, exc_info=True)
|
||
return
|
||
# #endregion run_mapping
|
||
# #endregion DatasetMapper
|
||
# #endregion DatasetMapperModule
|