2

2025-06-27 15:20:29 +03:00
parent 04fa28f086
commit c0a6ca7769
5 changed files with 389 additions and 252 deletions
--- a/superset_tool/client.py
+++ b/superset_tool/client.py
@@ -10,27 +10,29 @@

 # [IMPORTS] Стандартная библиотека
 import json
-from typing import Optional, Dict, Tuple, List, Any, Literal, Union,BinaryIO
+from typing import Optional, Dict, Tuple, List, Any, Literal, Union
+import datetime
 from pathlib import Path

 # [IMPORTS] Сторонние библиотеки
 import requests
 import urllib3
-from pydantic import BaseModel, Field
-from requests.exceptions import HTTPError
+import zipfile

 # [IMPORTS] Локальные модули
-from .models import SupersetConfig
-from .exceptions import (
+from superset_tool.models import SupersetConfig
+from superset_tool.exceptions import (
    AuthenticationError,
    SupersetAPIError,
    DashboardNotFoundError,
    NetworkError,
    PermissionDeniedError,
-    ExportError
+    ExportError,
+    InvalidZipFormatError
 )
-from .utils.fileio import get_filename_from_headers
-from .utils.logger import SupersetLogger
+from superset_tool.utils.fileio import get_filename_from_headers
+from superset_tool.utils.logger import SupersetLogger
+from superset_tool.utils.network import APIClient

 # [CONSTANTS] Логирование
 HTTP_METHODS = Literal['GET', 'POST', 'PUT', 'DELETE']
@@ -80,10 +82,14 @@ class SupersetClient:
        self._validate_config(config)
        self.config = config
        self.logger = config.logger or SupersetLogger(name="client")
-        self.session = self._setup_session()
+        self.network = APIClient(
+            base_url=config.base_url,
+            auth=config.auth,
+            verify_ssl=config.verify_ssl
+        )
+        self.tokens = self.network.authenticate()
        
        try:
-            self._authenticate()
            self.logger.info(
                "[COHERENCE_CHECK_PASSED] Клиент успешно инициализирован",
                extra={"base_url": config.base_url}
@@ -136,96 +142,6 @@ class SupersetClient:
            )
            raise ValueError("base_url должен начинаться с http:// или https://")
    
-    # [CHUNK] Настройка сессии и адаптеров
-    def _setup_session(self) -> requests.Session:
-        """[INTERNAL] Конфигурация HTTP-сессии
-        @coherence_check: SSL verification должен соответствовать config
-        """
-        session = requests.Session()
-        adapter = requests.adapters.HTTPAdapter(
-            max_retries=3,
-            pool_connections=10,
-            pool_maxsize=100
-        )
-        
-        session.mount('https://', adapter)
-        session.verify = self.config.verify_ssl
-        
-        if not self.config.verify_ssl:
-            urllib3.disable_warnings()
-            self.logger.debug(
-                "[CONFIG] SSL verification отключен",
-                extra={"base_url": self.config.base_url}
-            )
-            
-        return session
-
-    # [CHUNK] Процесс аутентификации
-    def _authenticate(self):
-        """[AUTH-FLOW] Получение токенов
-        @semantic_steps:
-            1. Получение access_token
-            2. Получение CSRF токена
-        @error_handling:
-            - AuthenticationError при проблемах credentials
-            - NetworkError при проблемах связи
-        """
-        try:
-            # [STEP 1] Получение bearer token
-            login_url = f"{self.config.base_url}/security/login"
-            response = self.session.post(
-                login_url,
-                json={
-                    "username": self.config.auth["username"],
-                    "password": self.config.auth["password"],
-                    "provider": self.config.auth["provider"],
-                    "refresh": self.config.auth["refresh"]
-                },
-                timeout=self.config.timeout
-            )
-            
-            if response.status_code == 401:
-                raise AuthenticationError(
-                    "Invalid credentials",
-                    context={
-                        "endpoint": login_url,
-                        "username": self.config.auth["username"],
-                        "status_code": response.status_code
-                    }
-                )
-                
-            response.raise_for_status()
-            self.access_token = response.json()["access_token"]
-
-            # [STEP 2] Получение CSRF token
-            csrf_url = f"{self.config.base_url}/security/csrf_token/"
-            response = self.session.get(
-                csrf_url,
-                headers={"Authorization": f"Bearer {self.access_token}"},
-                timeout=self.config.timeout
-            )
-            response.raise_for_status()
-            self.csrf_token = response.json()["result"]
-            
-            self.logger.info(
-                "[AUTH_SUCCESS] Токены успешно получены",
-                extra={
-                    "access_token": f"{self.access_token[:5]}...",
-                    "csrf_token": f"{self.csrf_token[:5]}..."
-                }
-            )
-            
-        except requests.exceptions.RequestException as e:
-            error_context = {
-                "method": e.request.method,
-                "url": e.request.url,
-                "status_code": getattr(e.response, 'status_code', None)
-            }
-            
-            if isinstance(e, (requests.Timeout, requests.ConnectionError)):
-                raise NetworkError("Connection failed", context=error_context) from e
-            raise SupersetAPIError("Auth flow failed", context=error_context) from e
-
    @property
    def headers(self) -> dict:
        """[INTERFACE] Базовые заголовки для API-вызовов
@@ -233,8 +149,8 @@ class SupersetClient:
        @post: Всегда возвращает актуальные токены
        """
        return {
-            "Authorization": f"Bearer {self.access_token}",
-            "X-CSRFToken": self.csrf_token,
+            "Authorization": f"Bearer {self.tokens['access_token']}",
+            "X-CSRFToken": self.tokens["csrf_token"],
            "Referer": self.config.base_url,
            "Content-Type": "application/json"
        }
@@ -244,62 +160,33 @@ class SupersetClient:
        """[CONTRACT] Получение метаданных дашборда
        @pre:
            - dashboard_id_or_slug должен существовать
-            - Токены должны быть валидны
+            - Клиент должен быть аутентифицирован (tokens актуальны)
        @post:
-            - Возвращает полные метаданные
+            - Возвращает dict с метаданными дашборда
            - В случае 404 вызывает DashboardNotFoundError
+        @semantic_layers:
+            1. Взаимодействие с API через APIClient
+            2. Обработка специфичных для Superset ошибок
        """
-        url = f"{self.config.base_url}/dashboard/{dashboard_id_or_slug}"
-        
        try:
-            response = self.session.get(
-                url,
-                headers=self.headers,
-                timeout=self.config.timeout
+            response = self.network.request(
+                method="GET",
+                endpoint=f"/dashboard/{dashboard_id_or_slug}",
+                headers=self.headers  # Автоматически включает токены
            )
-            
-            if response.status_code == 404:
+            return response.json()["result"]
+        
+        except requests.HTTPError as e:
+            if e.response.status_code == 404:
                raise DashboardNotFoundError(
                    dashboard_id_or_slug,
-                    context={"url": url}
+                    context={"url": f"{self.config.base_url}/dashboard/{dashboard_id_or_slug}"}
                )
-                
-            response.raise_for_status()
-            return response.json()["result"]
-            
-        except requests.exceptions.RequestException as e:
-            self._handle_api_error("get_dashboard", e, url)
-
-    def export_dashboard(self, dashboard_id: int) -> tuple[bytes, str]:
-        """[CONTRACT] Экспорт дашборда в ZIP
-        @error_handling:
-            - DashboardNotFoundError если дашборд не существует
-            - ExportError при проблемах экспорта
-        """
-        url = f"{self.config.base_url}/dashboard/export/"
+            raise SupersetAPIError(
+                f"API Error: {str(e)}", 
+                status_code=e.response.status_code
+            ) from e
        
-        try:
-            response = self.session.get(
-                url,
-                headers=self.headers,
-                params={"q": f"[{dashboard_id}]"},
-                timeout=self.config.timeout
-            )
-            
-            if response.status_code == 404:
-                raise DashboardNotFoundError(dashboard_id)
-                
-            response.raise_for_status()
-            
-            filename = (
-                get_filename_from_headers(response.headers) 
-                or f"dashboard_{dashboard_id}.zip"
-            )
-            return response.content, filename
-            
-        except requests.exceptions.RequestException as e:
-            self._handle_api_error("export_dashboard", e, url)
-
    # [ERROR-HANDLER] Централизованная обработка ошибок
    def _handle_api_error(self, method_name: str, error: Exception, url: str) -> None:
        """[UNIFIED-ERROR] Обработка API-ошибок
@@ -372,13 +259,12 @@ class SupersetClient:
            - Ответ должен иметь status_code 200
            - Content-Type: application/zip
        """
-        response = self.session.get(
-            url,
-            headers=self.headers,
-            params={"q": f"[{dashboard_id}]"},
-            timeout=self.config.timeout,
-            stream=True
-        )
+        response = self.network.request(
+                method="GET",
+                endpoint="/dashboard/export/",
+                params={"q": f"[{dashboard_id}]"},
+                raw_response=True  # Для получения бинарного содержимого
+            )
        response.raise_for_status()
        return response

@@ -483,8 +369,8 @@ class SupersetClient:
        
        try:
            # Инициализация пагинации
-            total_count = self._fetch_total_count(url)
-            paginated_data = self._fetch_all_pages(url, validated_query, total_count)
+            total_count = self._fetch_total_count()
+            paginated_data = self._fetch_all_pages(validated_query, total_count)
            
            self.logger.info(
                "[API_SUCCESS] Дашборды получены",
@@ -496,8 +382,8 @@ class SupersetClient:
            error_ctx = {"method": "get_dashboards", "query": validated_query}
            self._handle_api_error("Пагинация дашбордов", e, error_ctx)

-    # [SECTION] Импорт/экспорт
-    def import_dashboard(self, zip_path: Union[str, Path]) -> Dict:
+    # [SECTION] Импорт
+    def import_dashboard(self, file_name: Union[str, Path]) -> Dict:
        """[CONTRACT] Импорт дашборда из архива
        @pre:
            - Файл должен существовать и быть валидным ZIP
@@ -506,21 +392,36 @@ class SupersetClient:
            - Возвращает метаданные импортированного дашборда
            - При конфликтах выполняет overwrite
        """
-        self._validate_import_file(zip_path)
+        self._validate_import_file(file_name)
+        self.logger.debug(
+            "[IMPORT_START] Инициирован импорт дашборда",
+            extra={"file": file_name}
+            )
        
        try:
-            with open(zip_path, 'rb') as f:
-                return self._execute_import(
-                    file_obj=f,
-                    file_name=Path(zip_path).name
-                )
+            return self.network.upload_file(
+                endpoint="/dashboard/import/",
+                file_obj=file_name,
+                file_name=file_name,
+                form_field="formData",
+                extra_data={'overwrite': 'true'},
+                timeout=self.config.timeout * 2
+            )
+        
+        except PermissionDeniedError as e:
+            self.logger.error(
+                "[IMPORT_AUTH_FAILED] Недостаточно прав для импорта",
+                exc_info=True
+            )
+            raise
+
        except Exception as e:
            self.logger.error(
-                "[IMPORT_FAILED] Критическая ошибка импорта",
+                "[IMPORT_FAILED] Ошибка импорта дашборда",
                exc_info=True,
-                extra={"file": str(zip_path)}
+                extra={"file": file_name}
            )
-            raise DashboardImportError(f"Import failed: {str(e)}") from e
+            raise SupersetAPIError(f"Ошибка импорта: {str(e)}") from e

    # [SECTION] Приватные методы-помощники
    def _validate_query_params(self, query: Optional[Dict]) -> Dict:
@@ -532,34 +433,62 @@ class SupersetClient:
        }
        return {**base_query, **(query or {})}

-    def _fetch_total_count(self, url: str) -> int:
-        """[HELPER] Получение общего количества дашбордов"""
-        count_response = self.session.get(
-            f"{url}?q={json.dumps({'columns': ['id'], 'page': 0, 'page_size': 1})}",
-            headers=self.headers,
-            timeout=self.config.timeout
-        )
-        count_response.raise_for_status()
-        return count_response.json()['count']
-
-    def _fetch_all_pages(self, url: str, query: Dict, total_count: int) -> List[Dict]:
-        """[HELPER] Обход всех страниц с пагинацией"""
-        results = []
-        page_size = query['page_size']
-        total_pages = (total_count + page_size - 1) // page_size
-
-        for page in range(total_pages):
-            query['page'] = page
-            response = self.session.get(
-                url,
-                headers=self.headers,
-                params={"q": json.dumps(query)},
-                timeout=self.config.timeout
-            )
-            response.raise_for_status()
-            results.extend(response.json().get('result', []))
+    def _fetch_total_count(self) -> int:
+        """[CONTRACT][HELPER] Получение общего кол-ва дашбордов в системе
+        @delegates:
+            - Сетевой запрос -> APIClient
+            - Обработка ответа -> собственный метод
+        @errors:
+            - SupersetAPIError при проблемах с API
+        """
+        query_params = {
+            'columns': ['id'],
+            'page': 0,
+            'page_size': 1
+        }
        
-        return results
+        try:
+            return self.network.fetch_paginated_count(
+                endpoint="/dashboard/",
+                query_params=query_params,
+                count_field="count"
+            )
+        except requests.exceptions.RequestException as e:
+            raise SupersetAPIError(f"Ошибка получения количества дашбордов: {str(e)}")
+
+    def _fetch_all_pages(self, query: Dict, total_count: int) -> List[Dict]:
+        """[HELPER] Обход всех страниц с пагинацией"""
+        """[CONTRACT] Получение всех данных с пагинированного API
+        @delegates:
+            - Сетевые запросы -> APIClient.fetch_paginated_data()
+        @params:
+            query: оригинальный query-объект (без page)
+            total_count: общее количество элементов
+        @return:
+            Список всех элементов
+        @errors:
+            - SupersetAPIError: проблемы с API
+            - ValueError: некорректные параметры пагинации
+        """
+        try:
+            if not query.get('page_size'):
+                raise ValueError("Отсутствует page_size в query параметрах")
+
+            return self.network.fetch_paginated_data(
+                endpoint="/dashboard/",
+                base_query=query,
+                total_count=total_count,
+                results_field="result"
+            )
+
+        except (requests.exceptions.RequestException, ValueError) as e:
+            error_ctx = {
+                "query": query,
+                "total_count": total_count,
+                "error": str(e)
+            }
+            self.logger.error("[PAGINATION_ERROR]", extra=error_ctx)
+            raise SupersetAPIError(f"Ошибка пагинации: {str(e)}") from e

    def _validate_import_file(self, zip_path: Union[str, Path]) -> None:
        """[HELPER] Проверка файла перед импортом"""
@@ -573,24 +502,3 @@ class SupersetClient:
        with zipfile.ZipFile(path) as zf:
            if not any(n.endswith('metadata.yaml') for n in zf.namelist()):
                raise DashboardNotFoundError("Архив не содержит metadata.yaml")
-
-    def _execute_import(self, file_obj: BinaryIO, file_name: str) -> Dict:
-        """[HELPER] Выполнение API-запроса импорта"""
-        url = f"{self.config.base_url}/dashboard/import/"
-        
-        files = {'formData': (file_name, file_obj, 'application/x-zip-compressed')}
-        headers = {k: v for k, v in self.headers.items() if k.lower() != 'content-type'}
-        
-        response = self.session.post(
-            url,
-            files=files,
-            data={'overwrite': 'true'},
-            headers=headers,
-            timeout=self.config.timeout * 2  # Увеличенный таймаут для импорта
-        )
-        
-        if response.status_code == 403:
-            raise PermissionDeniedError("Недостаточно прав для импорта")
-            
-        response.raise_for_status()
-        return response.json()