Работающий импорт с изменением БД

This commit is contained in:
Volobuev Andrey
2025-04-04 09:19:37 +03:00
parent 992073d2f5
commit 6ffc432b42
8 changed files with 491 additions and 266 deletions

View File

@@ -7,11 +7,13 @@ from pydantic import BaseModel, Field
from .utils.fileio import *
from .exceptions import *
from .models import SupersetConfig
from .utils.logger import SupersetLogger
class SupersetClient:
def __init__(self, config: SupersetConfig):
self.config = config
self.logger = config.logger or SupersetLogger(console=False)
self.session = requests.Session()
self._setup_session()
self._authenticate()
@@ -25,6 +27,7 @@ class SupersetClient:
if not self.config.verify_ssl:
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
self.logger.debug(f"Проверка сертификатов SSL отключена")
self.session.mount('https://', adapter)
self.session.verify = self.config.verify_ssl
@@ -45,20 +48,25 @@ class SupersetClient:
)
response.raise_for_status()
self.access_token = response.json()["access_token"]
self.logger.info(f"Токен Bearer {self.access_token} получен c {login_url}")
# Затем получаем CSRF токен с использованием access_token
csrf_url = f"{self.config.base_url}/security/csrf_token/"
response = self.session.get(
csrf_url,
headers={"Authorization": f"Bearer {self.access_token}"},
verify=self.config.verify_ssl
)
response.raise_for_status()
self.csrf_token = response.json()["result"]
self.logger.info(f"Токен CSRF {self.csrf_token} получен c {csrf_url}")
except HTTPError as e:
if e.response.status_code == 401:
error_msg = "Invalid credentials" if "login" in e.request.url else "CSRF token fetch failed"
raise AuthenticationError(f"{error_msg}. Check auth configuration") from e
error_msg = f"Неверные данные для аутенфикации для {login_url}" if "login" in e.request.url else f"Не удалось получить CSRF токен с {csrf_url}"
self.logger.error(f"Ошибка получения: {error_msg}")
raise AuthenticationError(f"{error_msg}. Проверь данные аутенфикации") from e
raise
@@ -78,7 +86,7 @@ class SupersetClient:
:dashboard_id_or_slug - id или короткая ссылка
"""
url = f"{self.config.base_url}/dashboard/{dashboard_id_or_slug}"
self.logger.debug(f"Получаем информацию по дашборду с /{url}...")
try:
response = self.session.get(
url,
@@ -86,9 +94,11 @@ class SupersetClient:
timeout=self.config.timeout
)
response.raise_for_status()
self.logger.info(f"ОК - Получили информацию по дашборду с {url}")
return response.json()["result"]
except requests.exceptions.RequestException as e:
raise SupersetAPIError(f"Failed to get dashboard: {str(e)}") from e
self.logger.error(f"Ошибка при получении информации о дашборде: {str(e)}", exc_info=True)
raise SupersetAPIError(f"Ошибка при получении информации о дашборде: {str(e)}") from e
def get_dashboards(self, query: Optional[Dict] = None) -> Tuple[int, List[Dict]]:
@@ -102,11 +112,12 @@ class SupersetClient:
Tuple[int, List[Dict]]: Кортеж, содержащий общее количество дашбордов и список всех дашбордов.
"""
url = f"{self.config.base_url}/dashboard/"
self.logger.debug(f"Получаем информацию по дашбордам с {url}...")
modified_query: Dict = {}
all_results: List[Dict] = []
total_count: int = 0
current_page: int = 0
try:
total_count = self.session.get(
url,
@@ -114,7 +125,9 @@ class SupersetClient:
headers=self.headers,
timeout=self.config.timeout
).json()['count']
self.logger.info(f"ОК - Получили кол-во дашбордов ({total_count}) с {url}")
except requests.exceptions.RequestException as e:
self.logger.error(f"Ошибка при получении кол-ва дашбордов: {str(e)}", exc_info=True)
raise SupersetAPIError(f"Ошибка при получении кол-ва дашбордов: {str(e)}") from e
#Инициализация параметров запроса с учетом переданного query
@@ -151,18 +164,19 @@ class SupersetClient:
params={"q": json.dumps(modified_query)} ,
timeout=self.config.timeout
)
response.raise_for_status()
data = response.json()
all_results.extend(data.get("result", []))
current_page += 1
self.logger.info(f"ОК - Получили информацию по дашбордам с {url}")
# Проверка, достигли ли последней страницы
return total_count, all_results
except requests.exceptions.RequestException as e:
raise SupersetAPIError(f"Ошибка при получении дашбордов: {str(e)}") from e
self.logger.error(f"Ошибка при получении информации о дашбордах: {str(e)}", exc_info=True)
raise SupersetAPIError(f"Ошибка при получении информации о дашбордах: {str(e)}") from e
def export_dashboard(self, dashboard_id: int) -> Tuple[bytes, str]:
def export_dashboard(self, dashboard_id: int, logger: Optional[SupersetLogger] = None) -> Tuple[bytes, str]:
"""Экспортирует дашборд из Superset в виде ZIP-архива и возвращает его содержимое с именем файла.
Параметры:
@@ -193,6 +207,8 @@ class SupersetClient:
"""
url = f"{self.config.base_url}/dashboard/export/"
params = {"q": f"[{dashboard_id}]"}
logger = logger or SupersetLogger(name="client", console=False)
self.logger.debug(f"Экспортируем дашборд ID {dashboard_id} c {url}...")
try:
response = self.session.get(
@@ -204,9 +220,11 @@ class SupersetClient:
response.raise_for_status()
filename = get_filename_from_headers(response.headers) or f"dashboard_{dashboard_id}.zip"
self.logger.info(f"Дашборд сохранен в {filename}")
return response.content, filename
except requests.exceptions.RequestException as e:
self.logger.error(f"Ошибка при экспорте: {str(e)}", exc_info=True)
raise SupersetAPIError(f"Export failed: {str(e)}") from e
@@ -241,7 +259,7 @@ class SupersetClient:
- При конфликте имен может потребоваться ручное разрешение через параметры импорта
"""
url = f"{self.config.base_url}/dashboard/import/"
self.logger.debug(f"Импортируем дашборд ID {zip_path} на {url}...")
headers_without_content_type = {k: v for k, v in self.headers.items() if k.lower() != 'content-type'}
zip_name = zip_path.name
@@ -266,8 +284,10 @@ class SupersetClient:
# Обрабатываем ответ
try:
response.raise_for_status()
self.logger.info(f"Дашборд импортирован успешно")
return response.json()
except requests.exceptions.HTTPError as e:
self.logger.error(f"Ошибка при импорте: {str(e)}", exc_info=True)
error_detail = f"{e.response.status_code} {e.response.reason}"
if e.response.text:
error_detail += f"\nТело ответа: {e.response.text}"