476 lines
24 KiB
Python
476 lines
24 KiB
Python
# [MODULE_PATH] superset_tool.client
|
||
# [FILE] client.py
|
||
# [SEMANTICS] superset, api, client, logging, error-handling, slug-support
|
||
|
||
# --------------------------------------------------------------
|
||
# [IMPORTS]
|
||
# --------------------------------------------------------------
|
||
import json
|
||
import zipfile
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||
|
||
from requests import Response
|
||
|
||
from superset_tool.models import SupersetConfig
|
||
from superset_tool.exceptions import ExportError, InvalidZipFormatError
|
||
from superset_tool.utils.fileio import get_filename_from_headers
|
||
from superset_tool.utils.logger import SupersetLogger
|
||
from superset_tool.utils.network import APIClient
|
||
# [END_IMPORTS]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Service('SupersetClient')]
|
||
# [RELATION: Service('SupersetClient')] -> [DEPENDS_ON] -> [PythonModule('superset_tool.utils.network')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Класс‑обёртка над Superset REST‑API.
|
||
:preconditions:
|
||
- ``config`` – валидный объект :class:`SupersetConfig`.
|
||
- Доступен рабочий HTTP‑клиент :class:`APIClient`.
|
||
:postconditions:
|
||
- Объект готов к выполнению запросов (GET, POST, DELETE и т.д.).
|
||
:raises:
|
||
- :class:`TypeError` при передаче неверного типа конфигурации.
|
||
"""
|
||
class SupersetClient:
|
||
"""
|
||
:ivar SupersetLogger logger: Логгер, используемый в клиенте.
|
||
:ivar SupersetConfig config: Текущая конфигурация подключения.
|
||
:ivar APIClient network: Объект‑обёртка над ``requests``.
|
||
:ivar bool delete_before_reimport: Флаг, указывающий,
|
||
что при ошибке импорта дашборд следует удалить и повторить импорт.
|
||
"""
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('__init__')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Инициализировать клиент и передать ему логгер.
|
||
:preconditions: ``config`` – экземпляр :class:`SupersetConfig`.
|
||
:postconditions: Атрибуты ``logger``, ``config`` и ``network`` созданы,
|
||
``delete_before_reimport`` установлен в ``False``.
|
||
"""
|
||
def __init__(self, config: SupersetConfig, logger: Optional[SupersetLogger] = None):
|
||
self.logger = logger or SupersetLogger(name="SupersetClient")
|
||
self.logger.info("[INFO][SupersetClient.__init__] Initializing SupersetClient.")
|
||
self._validate_config(config)
|
||
self.config = config
|
||
self.network = APIClient(
|
||
config=config.dict(),
|
||
verify_ssl=config.verify_ssl,
|
||
timeout=config.timeout,
|
||
logger=self.logger,
|
||
)
|
||
self.delete_before_reimport: bool = False
|
||
self.logger.info("[INFO][SupersetClient.__init__] SupersetClient initialized.")
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_validate_config')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Проверить, что передан объект :class:`SupersetConfig`.
|
||
:preconditions: ``config`` – произвольный объект.
|
||
:postconditions: При несовпадении типов возбуждается :class:`TypeError`.
|
||
"""
|
||
def _validate_config(self, config: SupersetConfig) -> None:
|
||
self.logger.debug("[DEBUG][_validate_config][ENTER] Validating SupersetConfig.")
|
||
if not isinstance(config, SupersetConfig):
|
||
self.logger.error("[ERROR][_validate_config][FAILURE] Invalid config type.")
|
||
raise TypeError("Конфигурация должна быть экземпляром SupersetConfig")
|
||
self.logger.debug("[DEBUG][_validate_config][SUCCESS] Config is valid.")
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Property('headers')]
|
||
# --------------------------------------------------------------
|
||
@property
|
||
def headers(self) -> dict:
|
||
"""Базовые HTTP‑заголовки, используемые клиентом."""
|
||
return self.network.headers
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('get_dashboards')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Получить список дашбордов с поддержкой пагинации.
|
||
:preconditions: None.
|
||
:postconditions: Возвращается кортеж ``(total_count, list_of_dashboards)``.
|
||
"""
|
||
def get_dashboards(self, query: Optional[Dict] = None) -> Tuple[int, List[Dict]]:
|
||
self.logger.info("[INFO][get_dashboards][ENTER] Fetching dashboards.")
|
||
validated_query = self._validate_query_params(query)
|
||
total_count = self._fetch_total_object_count(endpoint="/dashboard/")
|
||
paginated_data = self._fetch_all_pages(
|
||
endpoint="/dashboard/",
|
||
pagination_options={
|
||
"base_query": validated_query,
|
||
"total_count": total_count,
|
||
"results_field": "result",
|
||
},
|
||
)
|
||
self.logger.info("[INFO][get_dashboards][SUCCESS] Got dashboards.")
|
||
return total_count, paginated_data
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('export_dashboard')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Скачать дашборд в виде ZIP‑архива.
|
||
:preconditions: ``dashboard_id`` – существующий идентификатор.
|
||
:postconditions: Возвращается бинарное содержимое и имя файла.
|
||
"""
|
||
def export_dashboard(self, dashboard_id: int) -> Tuple[bytes, str]:
|
||
self.logger.info("[INFO][export_dashboard][ENTER] Exporting dashboard %s.", dashboard_id)
|
||
response = self.network.request(
|
||
method="GET",
|
||
endpoint="/dashboard/export/",
|
||
params={"q": json.dumps([dashboard_id])},
|
||
stream=True,
|
||
raw_response=True,
|
||
)
|
||
self._validate_export_response(response, dashboard_id)
|
||
filename = self._resolve_export_filename(response, dashboard_id)
|
||
self.logger.info("[INFO][export_dashboard][SUCCESS] Exported dashboard %s.", dashboard_id)
|
||
return response.content, filename
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('import_dashboard')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Импортировать дашборд из ZIP‑файла. При неуспешном импорте,
|
||
если ``delete_before_reimport`` = True, сначала удаляется
|
||
дашборд по ID, затем импорт повторяется.
|
||
:preconditions:
|
||
- ``file_name`` – путь к существующему ZIP‑архиву (str|Path).
|
||
- ``dash_id`` – (опционально) ID дашборда, который следует удалить.
|
||
:postconditions: Возвращается словарь‑ответ API при успехе.
|
||
"""
|
||
def import_dashboard(
|
||
self,
|
||
file_name: Union[str, Path],
|
||
dash_id: Optional[int] = None,
|
||
dash_slug: Optional[str] = None, # сохраняем для возможного логирования
|
||
) -> Dict:
|
||
# -----------------------------------------------------------------
|
||
# 1️⃣ Приводим путь к строке (API‑клиент ожидает str)
|
||
# -----------------------------------------------------------------
|
||
file_path: str = str(file_name) # <--- гарантируем тип str
|
||
self._validate_import_file(file_path)
|
||
|
||
try:
|
||
import_response = self._do_import(file_path)
|
||
self.logger.info("[INFO][import_dashboard] Imported %s.", file_path)
|
||
return import_response
|
||
|
||
except Exception as exc:
|
||
# -----------------------------------------------------------------
|
||
# 2️⃣ Логируем первую неудачу, пытаемся удалить и повторить,
|
||
# только если включён флаг ``delete_before_reimport``.
|
||
# -----------------------------------------------------------------
|
||
self.logger.error(
|
||
"[ERROR][import_dashboard] First import attempt failed: %s",
|
||
exc,
|
||
exc_info=True,
|
||
)
|
||
if not self.delete_before_reimport:
|
||
raise
|
||
|
||
# -----------------------------------------------------------------
|
||
# 3️⃣ Выбираем, как искать дашборд для удаления.
|
||
# При наличии ``dash_id`` – удаляем его.
|
||
# Иначе, если известен ``dash_slug`` – переводим его в ID ниже.
|
||
# -----------------------------------------------------------------
|
||
target_id: Optional[int] = dash_id
|
||
if target_id is None and dash_slug is not None:
|
||
# Попытка динамического определения ID через slug.
|
||
# Мы делаем отдельный запрос к /dashboard/ (поисковый фильтр).
|
||
self.logger.debug("[DEBUG][import_dashboard] Resolving ID by slug '%s'.", dash_slug)
|
||
try:
|
||
_, candidates = self.get_dashboards(
|
||
query={"filters": [{"col": "slug", "op": "eq", "value": dash_slug}]}
|
||
)
|
||
if candidates:
|
||
target_id = candidates[0]["id"]
|
||
self.logger.debug("[DEBUG][import_dashboard] Resolved slug → ID %s.", target_id)
|
||
except Exception as e:
|
||
self.logger.warning(
|
||
"[WARN][import_dashboard] Could not resolve slug '%s' to ID: %s",
|
||
dash_slug,
|
||
e,
|
||
)
|
||
|
||
# Если всё‑равно нет ID – считаем невозможным корректно удалить.
|
||
if target_id is None:
|
||
self.logger.error("[ERROR][import_dashboard] No ID available for delete‑retry.")
|
||
raise
|
||
|
||
# -----------------------------------------------------------------
|
||
# 4️⃣ Удаляем найденный дашборд (по ID)
|
||
# -----------------------------------------------------------------
|
||
try:
|
||
self.delete_dashboard(target_id)
|
||
self.logger.info("[INFO][import_dashboard] Deleted dashboard ID %s, retrying import.", target_id)
|
||
except Exception as del_exc:
|
||
self.logger.error("[ERROR][import_dashboard] Delete failed: %s", del_exc, exc_info=True)
|
||
raise
|
||
|
||
# -----------------------------------------------------------------
|
||
# 5️⃣ Повторный импорт (тот же файл)
|
||
# -----------------------------------------------------------------
|
||
try:
|
||
import_response = self._do_import(file_path)
|
||
self.logger.info("[INFO][import_dashboard] Re‑import succeeded.")
|
||
return import_response
|
||
except Exception as rec_exc:
|
||
self.logger.error(
|
||
"[ERROR][import_dashboard] Re‑import after delete failed: %s",
|
||
rec_exc,
|
||
exc_info=True,
|
||
)
|
||
raise
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_do_import')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Выполнить один запрос на импорт без обработки исключений.
|
||
:preconditions: ``file_name`` уже проверен и существует.
|
||
:postconditions: Возвращается словарь‑ответ API.
|
||
"""
|
||
def _do_import(self, file_name: Union[str, Path]) -> Dict:
|
||
return self.network.upload_file(
|
||
endpoint="/dashboard/import/",
|
||
file_info={
|
||
"file_obj": Path(file_name),
|
||
"file_name": Path(file_name).name,
|
||
"form_field": "formData",
|
||
},
|
||
extra_data={"overwrite": "true"},
|
||
timeout=self.config.timeout * 2,
|
||
)
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('delete_dashboard')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Удалить дашборд **по ID или slug**.
|
||
:preconditions:
|
||
- ``dashboard_id`` – int ID **или** str slug дашборда.
|
||
:postconditions: На уровне API считается, что ресурс удалён
|
||
(HTTP 200/204). Логируется результат операции.
|
||
"""
|
||
def delete_dashboard(self, dashboard_id: Union[int, str]) -> None:
|
||
# ``dashboard_id`` может быть целым числом или строковым slug.
|
||
self.logger.info("[INFO][delete_dashboard][ENTER] Deleting dashboard %s.", dashboard_id)
|
||
response = self.network.request(
|
||
method="DELETE",
|
||
endpoint=f"/dashboard/{dashboard_id}",
|
||
)
|
||
# Superset обычно возвращает 200/204. Если есть поле ``result`` – проверяем.
|
||
if response.get("result", True) is not False:
|
||
self.logger.info("[INFO][delete_dashboard] Dashboard %s deleted.", dashboard_id)
|
||
else:
|
||
self.logger.warning("[WARN][delete_dashboard] Unexpected response while deleting %s.", dashboard_id)
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_extract_dashboard_id_from_zip')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Попытаться извлечь **ID** дашборда из ``metadata.yaml`` внутри ZIP‑архива.
|
||
:preconditions: ``file_name`` – путь к корректному ZIP‑файлу.
|
||
:postconditions: Возвращается ``int`` ID или ``None``.
|
||
"""
|
||
def _extract_dashboard_id_from_zip(self, file_name: Union[str, Path]) -> Optional[int]:
|
||
try:
|
||
import yaml
|
||
with zipfile.ZipFile(file_name, "r") as zf:
|
||
for name in zf.namelist():
|
||
if name.endswith("metadata.yaml"):
|
||
with zf.open(name) as meta_file:
|
||
meta = yaml.safe_load(meta_file.read())
|
||
dash_id = meta.get("dashboard_uuid") or meta.get("dashboard_id")
|
||
if dash_id is not None:
|
||
return int(dash_id)
|
||
except Exception as exc:
|
||
self.logger.error("[ERROR][_extract_dashboard_id_from_zip] %s", exc, exc_info=True)
|
||
return None
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_extract_dashboard_slug_from_zip')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Попытаться извлечь **slug** дашборда из ``metadata.yaml`` внутри ZIP‑архива.
|
||
:preconditions: ``file_name`` – путь к корректному ZIP‑файлу.
|
||
:postconditions: Возвращается строка‑slug или ``None``.
|
||
"""
|
||
def _extract_dashboard_slug_from_zip(self, file_name: Union[str, Path]) -> Optional[str]:
|
||
try:
|
||
import yaml
|
||
with zipfile.ZipFile(file_name, "r") as zf:
|
||
for name in zf.namelist():
|
||
if name.endswith("metadata.yaml"):
|
||
with zf.open(name) as meta_file:
|
||
meta = yaml.safe_load(meta_file.read())
|
||
slug = meta.get("slug")
|
||
if slug:
|
||
return str(slug)
|
||
except Exception as exc:
|
||
self.logger.error("[ERROR][_extract_dashboard_slug_from_zip] %s", exc, exc_info=True)
|
||
return None
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_validate_export_response')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Проверить, что ответ от ``/dashboard/export/`` – ZIP‑архив с данными.
|
||
:preconditions: ``response`` – объект :class:`requests.Response`.
|
||
:postconditions: При несоответствии возбуждается :class:`ExportError`.
|
||
"""
|
||
def _validate_export_response(self, response: Response, dashboard_id: int) -> None:
|
||
self.logger.debug("[DEBUG][_validate_export_response][ENTER] Validating response for %s.", dashboard_id)
|
||
content_type = response.headers.get("Content-Type", "")
|
||
if "application/zip" not in content_type:
|
||
self.logger.error("[ERROR][_validate_export_response][FAILURE] Invalid content type: %s", content_type)
|
||
raise ExportError(f"Получен не ZIP‑архив (Content-Type: {content_type})")
|
||
if not response.content:
|
||
self.logger.error("[ERROR][_validate_export_response][FAILURE] Empty response content.")
|
||
raise ExportError("Получены пустые данные при экспорте")
|
||
self.logger.debug("[DEBUG][_validate_export_response][SUCCESS] Response validated.")
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_resolve_export_filename')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Определить имя файла, полученного из заголовков ответа.
|
||
:preconditions: ``response.headers`` содержит (возможно) ``Content‑Disposition``.
|
||
:postconditions: Возвращается строка‑имя файла.
|
||
"""
|
||
def _resolve_export_filename(self, response: Response, dashboard_id: int) -> str:
|
||
self.logger.debug("[DEBUG][_resolve_export_filename][ENTER] Resolving filename.")
|
||
filename = get_filename_from_headers(response.headers)
|
||
if not filename:
|
||
from datetime import datetime
|
||
timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
|
||
filename = f"dashboard_export_{dashboard_id}_{timestamp}.zip"
|
||
self.logger.warning("[WARN][_resolve_export_filename] Generated filename: %s", filename)
|
||
self.logger.debug("[DEBUG][_resolve_export_filename][SUCCESS] Filename: %s", filename)
|
||
return filename
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_validate_query_params')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Сформировать корректный набор параметров запроса.
|
||
:preconditions: ``query`` – любой словарь или ``None``.
|
||
:postconditions: Возвращается словарь с обязательными полями.
|
||
"""
|
||
def _validate_query_params(self, query: Optional[Dict]) -> Dict:
|
||
base_query = {
|
||
"columns": ["slug", "id", "changed_on_utc", "dashboard_title", "published"],
|
||
"page": 0,
|
||
"page_size": 1000,
|
||
}
|
||
validated = {**base_query, **(query or {})}
|
||
self.logger.debug("[DEBUG][_validate_query_params] %s", validated)
|
||
return validated
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_fetch_total_object_count')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Получить общее количество объектов по указанному endpoint.
|
||
:preconditions: ``endpoint`` – строка, начинающаяся с «/».
|
||
:postconditions: Возвращается целое число.
|
||
"""
|
||
def _fetch_total_object_count(self, endpoint: str) -> int:
|
||
query_params_for_count = {"page": 0, "page_size": 1}
|
||
count = self.network.fetch_paginated_count(
|
||
endpoint=endpoint,
|
||
query_params=query_params_for_count,
|
||
count_field="count",
|
||
)
|
||
self.logger.debug("[DEBUG][_fetch_total_object_count] %s → %s", endpoint, count)
|
||
return count
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_fetch_all_pages')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Обойти все страницы пагинированного API.
|
||
:preconditions: ``pagination_options`` – словарь, сформированный
|
||
в ``_validate_query_params`` и ``_fetch_total_object_count``.
|
||
:postconditions: Возвращается список всех объектов.
|
||
"""
|
||
def _fetch_all_pages(self, endpoint: str, pagination_options: Dict) -> List[Dict]:
|
||
all_data = self.network.fetch_paginated_data(
|
||
endpoint=endpoint,
|
||
pagination_options=pagination_options,
|
||
)
|
||
self.logger.debug("[DEBUG][_fetch_all_pages] Fetched %s items from %s.", len(all_data), endpoint)
|
||
return all_data
|
||
# [END_ENTITY]
|
||
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('_validate_import_file')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Проверить, что файл существует, является ZIP‑архивом и
|
||
содержит ``metadata.yaml``.
|
||
:preconditions: ``zip_path`` – путь к файлу.
|
||
:postconditions: При невалидном файле возбуждается :class:`InvalidZipFormatError`.
|
||
"""
|
||
def _validate_import_file(self, zip_path: Union[str, Path]) -> None:
|
||
path = Path(zip_path)
|
||
if not path.exists():
|
||
self.logger.error("[ERROR][_validate_import_file] File not found: %s", zip_path)
|
||
raise FileNotFoundError(f"Файл {zip_path} не существует")
|
||
if not zipfile.is_zipfile(path):
|
||
self.logger.error("[ERROR][_validate_import_file] Not a zip file: %s", zip_path)
|
||
raise InvalidZipFormatError(f"Файл {zip_path} не является ZIP‑архивом")
|
||
with zipfile.ZipFile(path, "r") as zf:
|
||
if not any(n.endswith("metadata.yaml") for n in zf.namelist()):
|
||
self.logger.error("[ERROR][_validate_import_file] No metadata.yaml in %s", zip_path)
|
||
raise InvalidZipFormatError(f"Архив {zip_path} не содержит 'metadata.yaml'")
|
||
self.logger.debug("[DEBUG][_validate_import_file] File %s validated.", zip_path)
|
||
# [END_ENTITY]
|
||
# --------------------------------------------------------------
|
||
# [ENTITY: Method('get_datasets')]
|
||
# --------------------------------------------------------------
|
||
"""
|
||
:purpose: Получить список датасетов с поддержкой пагинации.
|
||
:preconditions: None.
|
||
:postconditions: Возвращается кортеж ``(total_count, list_of_datasets)``.
|
||
"""
|
||
|
||
def get_datasets(self, query: Optional[Dict] = None) -> Tuple[int, List[Dict]]:
|
||
self.logger.info("[INFO][get_datasets][ENTER] Fetching datasets.")
|
||
validated_query = self._validate_query_params(query)
|
||
total_count = self._fetch_total_object_count(endpoint="/dataset/")
|
||
paginated_data = self._fetch_all_pages(
|
||
endpoint="/dataset/",
|
||
pagination_options={
|
||
"base_query": validated_query,
|
||
"total_count": total_count,
|
||
"results_field": "result",
|
||
},
|
||
)
|
||
self.logger.info("[INFO][get_datasets][SUCCESS] Got datasets.")
|
||
return total_count, paginated_data
|
||
# [END_ENTITY]
|
||
|
||
|
||
# [END_FILE client.py] |