270 lines
11 KiB
Python
270 lines
11 KiB
Python
# #region TargetSchemaValidation [C:3] [TYPE Module] [SEMANTICS translate, schema, validation, target-table]
|
||
# @BRIEF Проверка схемы целевой таблицы: запрос колонок через Superset SQL Lab,
|
||
# сравнение с ожидаемыми (из build_columns), возврат diff.
|
||
# @LAYER Service
|
||
# @RELATION DEPENDS_ON -> [SupersetSqlLabExecutor]
|
||
# @RELATION DEPENDS_ON -> [schemas.translate.TargetSchemaValidationRequest]
|
||
# @RELATION DEPENDS_ON -> [schemas.translate.TargetSchemaValidationResponse]
|
||
# @PRE Superset окружение доступно, target_database_id валиден.
|
||
# @POST Возвращает актуальные, ожидаемые, отсутствующие и лишние колонки.
|
||
# @SIDE_EFFECT Выполняет SQL-запрос через Superset SQL Lab.
|
||
# @INVARIANT Если таблица не существует — table_exists=False без ошибки.
|
||
# @INVARIANT Если Superset недоступен — error с описанием, all_present=False.
|
||
# @INVARIANT Логика ожидаемых колонок sync с build_columns() в orchestrator_sql_rows.py.
|
||
# #endregion TargetSchemaValidation
|
||
|
||
import logging
|
||
import re
|
||
from typing import Any
|
||
|
||
from ...schemas.translate import (
|
||
TargetSchemaColumnInfo,
|
||
TargetSchemaValidationRequest,
|
||
TargetSchemaValidationResponse,
|
||
)
|
||
from ...core.config_manager import ConfigManager
|
||
from .superset_executor import SupersetSqlLabExecutor
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
_TABLE_NAME_RE = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$')
|
||
|
||
|
||
# #region _build_expected_columns [C:2] [TYPE Function]
|
||
# @BRIEF Собирает список ожидаемых колонок по конфигурации column mapping.
|
||
# Логика идентична build_columns() из orchestrator_sql_rows.py,
|
||
# но не требует объекта TranslationJob.
|
||
def _build_expected_columns(req: TargetSchemaValidationRequest) -> list[TargetSchemaColumnInfo]:
|
||
"""
|
||
Определяет, какие колонки INSERT будет пытаться заполнить,
|
||
на основе конфигурации column mapping.
|
||
"""
|
||
names: list[str] = []
|
||
|
||
# Ключевые колонки (user-defined, для upsert matching)
|
||
if req.target_key_cols:
|
||
names.extend(req.target_key_cols)
|
||
|
||
# Целевая колонка для переведённого текста
|
||
effective_target = req.target_column or req.translation_column
|
||
if effective_target:
|
||
names.append(effective_target)
|
||
|
||
# Колонка с кодом языка перевода
|
||
if req.target_language_column:
|
||
names.append(req.target_language_column)
|
||
|
||
# Колонка с исходным текстом (source reference)
|
||
if req.target_source_column:
|
||
names.append(req.target_source_column)
|
||
|
||
# Колонка с определённым исходным языком
|
||
if req.target_source_language_column:
|
||
names.append(req.target_source_language_column)
|
||
|
||
# Служебные колонки — всегда
|
||
names.append("context")
|
||
names.append("is_original")
|
||
|
||
# Дедупликация с сохранением порядка
|
||
seen: set[str] = set()
|
||
deduped: list[str] = []
|
||
for c in names:
|
||
if c and c not in seen:
|
||
deduped.append(c)
|
||
seen.add(c)
|
||
|
||
return [TargetSchemaColumnInfo(name=n) for n in deduped]
|
||
# #endregion _build_expected_columns
|
||
|
||
|
||
# #region _parse_column_results [C:2] [TYPE Function]
|
||
# @BRIEF Извлекает информацию о колонках из ответа Superset SQL Lab.
|
||
#
|
||
# Superset SQL Lab возвращает разные форматы в зависимости от версии.
|
||
# Эта функция ищет колонки во всех возможных местах ответа.
|
||
def _parse_column_results(result: dict[str, Any]) -> tuple[list[dict[str, Any]], bool]:
|
||
"""
|
||
Парсит результат SQL Lab запроса для извлечения списка колонок таблицы.
|
||
Returns (columns_info, table_exists).
|
||
"""
|
||
if not result:
|
||
return [], False
|
||
|
||
raw = result.get("raw_response", result)
|
||
|
||
# Пробуем извлечь данные из разных форматов ответа Superset
|
||
columns_data = None
|
||
data_rows = None
|
||
|
||
# Формат 1: прямой ответ с колонками (sync mode)
|
||
if "columns" in raw and isinstance(raw["columns"], list):
|
||
columns_data = raw["columns"]
|
||
data_rows = raw.get("data") if isinstance(raw.get("data"), list) else None
|
||
|
||
# Формат 2: вложенный result (async polling)
|
||
elif "result" in raw and isinstance(raw["result"], dict):
|
||
res = raw["result"]
|
||
if "columns" in res:
|
||
columns_data = res["columns"]
|
||
data_rows = res.get("data") if isinstance(res.get("data"), list) else None
|
||
|
||
# Формат 3: ключ results (get_query_results)
|
||
elif "results" in raw and isinstance(raw["results"], dict):
|
||
res = raw["results"]
|
||
if "columns" in res:
|
||
columns_data = res["columns"]
|
||
if "data" in res:
|
||
data_rows = res["data"]
|
||
if isinstance(res.get("result"), dict):
|
||
sub = res["result"]
|
||
if not columns_data and "columns" in sub:
|
||
columns_data = sub["columns"]
|
||
if not data_rows and "data" in sub:
|
||
data_rows = sub["data"]
|
||
|
||
if not columns_data:
|
||
return [], False
|
||
|
||
# Проверяем существование таблицы: есть ли хотя бы одна строка
|
||
table_exists = data_rows is not None and len(data_rows) > 0
|
||
|
||
# Извлекаем имена и типы колонок из метаданных ответа
|
||
columns_info: list[dict[str, Any]] = []
|
||
for col in columns_data:
|
||
col_name = (
|
||
col.get("name")
|
||
or col.get("column_name")
|
||
or col.get("Column", "")
|
||
)
|
||
col_type = (
|
||
col.get("type")
|
||
or col.get("data_type")
|
||
or col.get("Type", "")
|
||
or col.get("col_type", "")
|
||
)
|
||
is_nullable = col.get("is_nullable")
|
||
# PostgreSQL returns 'YES'/'NO' for is_nullable
|
||
if isinstance(is_nullable, str):
|
||
is_nullable = is_nullable.upper() == "YES"
|
||
|
||
if col_name:
|
||
columns_info.append({
|
||
"name": str(col_name),
|
||
"type": str(col_type) if col_type else None,
|
||
"is_nullable": is_nullable if isinstance(is_nullable, bool) else True,
|
||
})
|
||
|
||
return columns_info, table_exists
|
||
# #endregion _parse_column_results
|
||
|
||
|
||
# #region validate_target_table_schema [C:3] [TYPE Function]
|
||
# @BRIEF Основная функция: проверяет схему целевой таблицы через Superset SQL Lab.
|
||
def validate_target_table_schema(
|
||
req: TargetSchemaValidationRequest,
|
||
config_manager: ConfigManager,
|
||
) -> TargetSchemaValidationResponse:
|
||
"""
|
||
Проверяет схему целевой таблицы:
|
||
1. Вычисляет ожидаемые колонки (build_columns логика)
|
||
2. Запрашивает актуальные колонки через Superset SQL Lab
|
||
(SELECT column_name, data_type, is_nullable FROM information_schema.columns ...)
|
||
3. Сравнивает и возвращает diff
|
||
"""
|
||
# 1. Ожидаемые колонки
|
||
expected = _build_expected_columns(req)
|
||
expected_names = {c.name for c in expected}
|
||
|
||
# 2. Валидация имён схемы и таблицы (защита от SQL injection в identifier)
|
||
if req.target_table and not _TABLE_NAME_RE.match(req.target_table):
|
||
return TargetSchemaValidationResponse(
|
||
table_exists=False,
|
||
error=f"Invalid target table name: '{req.target_table}'. Only alphanumeric and underscore allowed.",
|
||
expected_columns=expected,
|
||
actual_columns=[],
|
||
missing_columns=expected,
|
||
extra_columns=[],
|
||
all_present=False,
|
||
)
|
||
if req.target_schema and not _TABLE_NAME_RE.match(req.target_schema):
|
||
return TargetSchemaValidationResponse(
|
||
table_exists=False,
|
||
error=f"Invalid target schema name: '{req.target_schema}'. Only alphanumeric and underscore allowed.",
|
||
expected_columns=expected,
|
||
actual_columns=[],
|
||
missing_columns=expected,
|
||
extra_columns=[],
|
||
all_present=False,
|
||
)
|
||
|
||
# 3. Запрос через Superset SQL Lab
|
||
try:
|
||
executor = SupersetSqlLabExecutor(config_manager, req.environment_id)
|
||
executor.resolve_database_id(target_database_id=req.target_database_id)
|
||
|
||
safe_schema = (req.target_schema or "public").replace("'", "''")
|
||
safe_table = req.target_table.replace("'", "''")
|
||
|
||
sql = (
|
||
f"SELECT column_name, data_type, is_nullable "
|
||
f"FROM information_schema.columns "
|
||
f"WHERE table_schema = '{safe_schema}' "
|
||
f"AND table_name = '{safe_table}' "
|
||
f"ORDER BY ordinal_position"
|
||
)
|
||
|
||
result = executor.execute_and_poll(sql=sql, max_polls=15, poll_interval_seconds=1.0)
|
||
|
||
status = result.get("status", "")
|
||
if status == "failed":
|
||
error_msg = result.get("error_message", "SQL Lab query failed")
|
||
return TargetSchemaValidationResponse(
|
||
table_exists=False,
|
||
error=f"Superset SQL Lab error: {error_msg}",
|
||
expected_columns=expected,
|
||
actual_columns=[],
|
||
missing_columns=expected,
|
||
extra_columns=[],
|
||
all_present=False,
|
||
)
|
||
|
||
raw_columns, table_exists = _parse_column_results(result)
|
||
|
||
actual_cols = []
|
||
for col in raw_columns:
|
||
actual_cols.append(TargetSchemaColumnInfo(
|
||
name=col["name"],
|
||
data_type=col.get("type"),
|
||
is_nullable=col.get("is_nullable", True),
|
||
))
|
||
|
||
actual_names = {c.name for c in actual_cols}
|
||
|
||
# 3. Diff
|
||
missing = [col for col in expected if col.name not in actual_names]
|
||
extra = [col for col in actual_cols if col.name not in expected_names]
|
||
|
||
return TargetSchemaValidationResponse(
|
||
table_exists=table_exists,
|
||
expected_columns=expected,
|
||
actual_columns=actual_cols,
|
||
missing_columns=missing,
|
||
extra_columns=extra,
|
||
all_present=len(missing) == 0,
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.warning("Target schema validation failed: %s", str(e), exc_info=True)
|
||
return TargetSchemaValidationResponse(
|
||
table_exists=False,
|
||
error=f"Failed to validate target schema: {e}",
|
||
expected_columns=expected,
|
||
actual_columns=[],
|
||
missing_columns=expected,
|
||
extra_columns=[],
|
||
all_present=False,
|
||
)
|
||
# #endregion validate_target_table_schema
|
||
# #endregion TargetSchemaValidation
|