Backup

2025-04-08 16:38:58 +03:00
parent 6ffc432b42
commit 625b50a6d2
4 changed files with 115 additions and 102 deletions
--- a/superset_tool/client.py
+++ b/superset_tool/client.py
@@ -17,7 +17,7 @@ class SupersetClient:
        self.session = requests.Session()
        self._setup_session()
        self._authenticate()
-    
+
    def _setup_session(self):
        adapter = requests.adapters.HTTPAdapter(
            max_retries=3,
@@ -28,7 +28,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

@@ -48,11 +48,12 @@ class SupersetClient:
            )
            response.raise_for_status()
            self.access_token = response.json()["access_token"]
-            self.logger.info(f"Токен Bearer {self.access_token} получен c {login_url}")
-            
+            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}"},
@@ -61,14 +62,15 @@ class SupersetClient:
            response.raise_for_status()

            self.csrf_token = response.json()["result"]
-            self.logger.info(f"Токен CSRF {self.csrf_token} получен c {csrf_url}")
+            self.logger.info(
+                f"Токен CSRF {self.csrf_token} получен c {csrf_url}")
        except HTTPError as e:
            if e.response.status_code == 401:
                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 AuthenticationError(
+                    f"{error_msg}. Проверь данные аутенфикации") from e
            raise
-        

    @property
    def headers(self):
@@ -79,14 +81,15 @@ class SupersetClient:
            "Content-Type": "application/json"
        }

-    def get_dashboard(self, dashboard_id_or_slug: str ) -> Dict:
+    def get_dashboard(self, dashboard_id_or_slug: str) -> Dict:
        """
-        Получаем информацию по дашборду (если передан dashboard_id_or_slug), либо по всем дашбордам, если параметр не передан
+        Получаем информацию по дашборду (если передан dashboard_id_or_slug)
        Параметры:
            :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,
@@ -94,12 +97,14 @@ class SupersetClient:
                timeout=self.config.timeout
            )
            response.raise_for_status()
-            self.logger.info(f"ОК - Получили информацию по дашборду с {url}")
+            self.logger.info(f"ОК - Получили информацию по дашборду с {response.url}")
+
            return response.json()["result"]
        except requests.exceptions.RequestException as e:
-            self.logger.error(f"Ошибка при получении информации о дашборде: {str(e)}", exc_info=True)
-            raise SupersetAPIError(f"Ошибка при получении информации о дашборде:  {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]]:
        """
@@ -117,20 +122,25 @@ class SupersetClient:
        all_results: List[Dict] = []
        total_count: int = 0
        current_page: int = 0
-
+        q_param = '{ "columns": [  "id"  ], "page": 0, "page_size": 20}'
        try:
-            total_count = self.session.get(
-                url,
-                #q=modified_query,
+            response = self.session.get(
+                url=f"{url}?q={q_param}",
+                #params={"q": json.dumps(default_query)},  # Передаем такой body, иначе на prodta отдает 1 дашборд
                headers=self.headers,
                timeout=self.config.timeout
-            ).json()['count']
-            self.logger.info(f"ОК - Получили кол-во дашбордов ({total_count}) с {url}")
+            )
+            total_count = response.json()['count']
+            self.logger.info(
+                f"ОК - Получили кол-во дашбордов ({total_count}) с {url}")
+            self.logger.info(f"Запрос - {response.url}")
        except requests.exceptions.RequestException as e:
-           self.logger.error(f"Ошибка при получении кол-ва дашбордов: {str(e)}", exc_info=True)
-           raise SupersetAPIError(f"Ошибка при получении кол-ва дашбордов: {str(e)}") from e
-        #Инициализация параметров запроса с учетом переданного query
-        
+            self.logger.error(
+                f"Ошибка при получении кол-ва дашбордов: {str(e)}", exc_info=True)
+            raise SupersetAPIError(
+                f"Ошибка при получении кол-ва дашбордов: {str(e)}") from e
+        # Инициализация параметров запроса с учетом переданного query
+
        if query:
            modified_query = query.copy()
            # Убедимся, что page_size установлен, если не передан
@@ -150,18 +160,15 @@ class SupersetClient:

        page_size = modified_query["page_size"]

-        
-
        total_pages = (total_count + page_size - 1) // page_size

-       
        try:
            while current_page < total_pages:
                modified_query["page"] = current_page
                response = self.session.get(
                    url,
                    headers=self.headers,
-                    params={"q": json.dumps(modified_query)} ,
+                    params={"q": json.dumps(modified_query)},
                    timeout=self.config.timeout
                )
                response.raise_for_status()
@@ -170,11 +177,13 @@ class SupersetClient:
                current_page += 1

            self.logger.info(f"ОК - Получили информацию по дашбордам с {url}")
-                # Проверка, достигли ли последней страницы
+            # Проверка, достигли ли последней страницы
            return total_count, all_results
        except requests.exceptions.RequestException as e:
-            self.logger.error(f"Ошибка при получении информации о дашбордах: {str(e)}", exc_info=True)
-            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, logger: Optional[SupersetLogger] = None) -> Tuple[bytes, str]:
        """Экспортирует дашборд из Superset в виде ZIP-архива и возвращает его содержимое с именем файла.
@@ -209,7 +218,7 @@ class SupersetClient:
        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(
                url,
@@ -218,15 +227,15 @@ class SupersetClient:
                timeout=self.config.timeout
            )
            response.raise_for_status()
-            
-            filename = get_filename_from_headers(response.headers) or f"dashboard_{dashboard_id}.zip"
+
+            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
-        

    def import_dashboard(self, zip_path) -> Dict:
        """Импортирует дашборд в Superset из ZIP-архива.
@@ -260,8 +269,9 @@ 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'}
-        
+        headers_without_content_type = {
+            k: v for k, v in self.headers.items() if k.lower() != 'content-type'}
+
        zip_name = zip_path.name
        # Подготавливаем данные для multipart/form-data
        with open(zip_path, 'rb') as f:
@@ -272,15 +282,15 @@ class SupersetClient:
                    'application/x-zip-compressed'  # MIME-тип из curl
                )
            }
-            
+
            # Отправляем запрос
            response = self.session.post(
-                        url,
-                        files=files,
-                        data={'overwrite': 'true'},
-                        headers=headers_without_content_type,
-                        timeout=self.config.timeout * 2  # Longer timeout for imports
-                    )
+                url,
+                files=files,
+                data={'overwrite': 'true'},
+                headers=headers_without_content_type,
+                timeout=self.config.timeout * 2  # Longer timeout for imports
+            )
        # Обрабатываем ответ
        try:
            response.raise_for_status()
@@ -291,4 +301,4 @@ class SupersetClient:
            error_detail = f"{e.response.status_code} {e.response.reason}"
            if e.response.text:
                error_detail += f"\nТело ответа: {e.response.text}"
-            raise RuntimeError(f"Ошибка импорта: {error_detail}") from e
+            raise RuntimeError(f"Ошибка импорта: {error_detail}") from e
--- a/superset_tool/utils/fileio.py
+++ b/superset_tool/utils/fileio.py
@@ -13,24 +13,35 @@ from ..utils.logger import SupersetLogger

@contextmanager
 def create_temp_file(
-    content: bytes,
+    content: Optional[bytes] = None,
    suffix: str = ".zip",
+    mode: str = 'wb',
    logger: Optional[SupersetLogger] = None
 ):
-    """Контекстный менеджер для создания временных файлов с логированием"""
+    """Расширенный контекстный менеджер для временных файлов/директорий"""
    logger = logger or SupersetLogger(name="fileio", console=False)
    try:
-        logger.debug(f"Создание временного файла с суффиксом {suffix}")
-        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
-            tmp.write(content)
-            tmp.flush()
-        yield Path(tmp.name)
+        logger.debug(f"Создание временного ресурса с суффиксом {suffix}")
+        
+        # Для директорий
+        if suffix.startswith('.dir'):
+            with tempfile.TemporaryDirectory(suffix=suffix) as tmp_dir:
+                logger.debug(f"Создана временная директория: {tmp_dir}")
+                yield Path(tmp_dir)
+        
+        # Для файлов
+        else:
+            with tempfile.NamedTemporaryFile(suffix=suffix, mode=mode, delete=False) as tmp:
+                if content:
+                    tmp.write(content)
+                    tmp.flush()
+                logger.debug(f"Создан временный файл: {tmp.name}")
+                yield Path(tmp.name)
    except Exception as e:
-        logger.error(
-            f"Ошибка создания временного файла: {str(e)}", exc_info=True)
+        logger.error(f"Ошибка создания временного ресурса: {str(e)}", exc_info=True)
        raise
    finally:
-        if Path(tmp.name).exists():
+        if 'tmp' in locals() and Path(tmp.name).exists() and not suffix.startswith('.dir'):
            Path(tmp.name).unlink()
            logger.debug(f"Временный файл {tmp.name} удален")