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