init
This commit is contained in:
Volobuev Andrey
274
superset_tool/client.py
Normal file
274
superset_tool/client.py
Normal file
@@ -0,0 +1,274 @@
|
||||
import requests
|
||||
from requests.exceptions import HTTPError
|
||||
import urllib3
|
||||
import json
|
||||
from typing import Dict, Optional, Tuple, List, Any
|
||||
from pydantic import BaseModel, Field
|
||||
from .utils.fileio import *
|
||||
from .exceptions import *
|
||||
from .models import SupersetConfig
|
||||
|
||||
|
||||
class SupersetClient:
|
||||
def __init__(self, config: SupersetConfig):
|
||||
self.config = config
|
||||
self.session = requests.Session()
|
||||
self._setup_session()
|
||||
self._authenticate()
|
||||
|
||||
def _setup_session(self):
|
||||
adapter = requests.adapters.HTTPAdapter(
|
||||
max_retries=3,
|
||||
pool_connections=10,
|
||||
pool_maxsize=100
|
||||
)
|
||||
|
||||
if not self.config.verify_ssl:
|
||||
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
||||
|
||||
self.session.mount('https://', adapter)
|
||||
self.session.verify = self.config.verify_ssl
|
||||
|
||||
def _authenticate(self):
|
||||
try:
|
||||
# Сначала логинимся для получения access_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": True
|
||||
},
|
||||
verify=self.config.verify_ssl
|
||||
)
|
||||
response.raise_for_status()
|
||||
self.access_token = response.json()["access_token"]
|
||||
|
||||
# Затем получаем 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"]
|
||||
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
|
||||
raise
|
||||
|
||||
|
||||
@property
|
||||
def headers(self):
|
||||
return {
|
||||
"Authorization": f"Bearer {self.access_token}",
|
||||
"X-CSRFToken": self.csrf_token,
|
||||
"Referer": self.config.base_url,
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
def get_dashboard(self, dashboard_id_or_slug: str ) -> Dict:
|
||||
"""
|
||||
Получаем информацию по дашборду (если передан dashboard_id_or_slug), либо по всем дашбордам, если параметр не передан
|
||||
Параметры:
|
||||
:dashboard_id_or_slug - id или короткая ссылка
|
||||
"""
|
||||
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.raise_for_status()
|
||||
return response.json()["result"]
|
||||
except requests.exceptions.RequestException as e:
|
||||
raise SupersetAPIError(f"Failed to get dashboard: {str(e)}") from e
|
||||
|
||||
|
||||
def get_dashboards(self, query: Optional[Dict] = None) -> Tuple[int, List[Dict]]:
|
||||
"""
|
||||
Получаем информацию по всем дашбордам с учетом пагинации.
|
||||
|
||||
Параметры:
|
||||
query (Optional[Dict]): Дополнительные параметры запроса, включая пагинацию и фильтры.
|
||||
|
||||
Возвращает:
|
||||
Tuple[int, List[Dict]]: Кортеж, содержащий общее количество дашбордов и список всех дашбордов.
|
||||
"""
|
||||
url = f"{self.config.base_url}/dashboard/"
|
||||
modified_query: Dict = {}
|
||||
all_results: List[Dict] = []
|
||||
total_count: int = 0
|
||||
current_page: int = 0
|
||||
|
||||
try:
|
||||
total_count = self.session.get(
|
||||
url,
|
||||
#q=modified_query,
|
||||
headers=self.headers,
|
||||
timeout=self.config.timeout
|
||||
).json()['count']
|
||||
except requests.exceptions.RequestException as e:
|
||||
raise SupersetAPIError(f"Ошибка при получении кол-ва дашбордов: {str(e)}") from e
|
||||
#Инициализация параметров запроса с учетом переданного query
|
||||
|
||||
if query:
|
||||
modified_query = query.copy()
|
||||
# Убедимся, что page_size установлен, если не передан
|
||||
modified_query.setdefault("page_size", 20)
|
||||
else:
|
||||
modified_query = {
|
||||
"columns": [
|
||||
"slug",
|
||||
"id",
|
||||
"changed_on_utc",
|
||||
"dashboard_title",
|
||||
"published"
|
||||
],
|
||||
"page": 0,
|
||||
"page_size": 20
|
||||
}
|
||||
|
||||
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)} ,
|
||||
timeout=self.config.timeout
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
all_results.extend(data.get("result", []))
|
||||
|
||||
current_page += 1
|
||||
# Проверка, достигли ли последней страницы
|
||||
return total_count, all_results
|
||||
except requests.exceptions.RequestException as e:
|
||||
raise SupersetAPIError(f"Ошибка при получении дашбордов: {str(e)}") from e
|
||||
|
||||
def export_dashboard(self, dashboard_id: int) -> Tuple[bytes, str]:
|
||||
"""Экспортирует дашборд из Superset в виде ZIP-архива и возвращает его содержимое с именем файла.
|
||||
|
||||
Параметры:
|
||||
:dashboard_id (int): Идентификатор дашборда для экспорта
|
||||
|
||||
Возвращает:
|
||||
Tuple[bytes, str]: Кортеж, содержащий:
|
||||
- bytes: Бинарное содержимое ZIP-архива с дашбордом
|
||||
- str: Имя файла (из заголовков ответа или в формате dashboard_{id}.zip)
|
||||
|
||||
Исключения:
|
||||
SupersetAPIError: Вызывается при ошибках:
|
||||
- Проблемы с сетью/соединением
|
||||
- Невалидный ID дашборда
|
||||
- Отсутствие прав доступа
|
||||
- Ошибки сервера (status code >= 400)
|
||||
|
||||
Пример использования:
|
||||
content, filename = client.export_dashboard(5)
|
||||
with open(filename, 'wb') as f:
|
||||
f.write(content)
|
||||
|
||||
Примечания:
|
||||
- Для экспорта используется API Endpoint: /dashboard/export/
|
||||
- Имя файла пытается извлечь из Content-Disposition заголовка
|
||||
- По умолчанию возвращает имя в формате dashboard_{id}.zip
|
||||
- Архив содержит JSON-метаданные и связанные элементы (датасеты, чарты)
|
||||
"""
|
||||
url = f"{self.config.base_url}/dashboard/export/"
|
||||
params = {"q": f"[{dashboard_id}]"}
|
||||
|
||||
try:
|
||||
response = self.session.get(
|
||||
url,
|
||||
headers=self.headers,
|
||||
params=params,
|
||||
timeout=self.config.timeout
|
||||
)
|
||||
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:
|
||||
raise SupersetAPIError(f"Export failed: {str(e)}") from e
|
||||
|
||||
|
||||
def import_dashboard(self, zip_path) -> Dict:
|
||||
"""Импортирует дашборд в Superset из ZIP-архива.
|
||||
|
||||
Параметры:
|
||||
zip_path (Path): Путь к ZIP-файлу с дашбордом для импорта
|
||||
|
||||
Возвращает:
|
||||
dict: Ответ API в формате JSON с результатами импорта
|
||||
|
||||
Исключения:
|
||||
RuntimeError: Вызывается при:
|
||||
- Ошибках сети/соединения
|
||||
- Невалидном формате ZIP-архива
|
||||
- Конфликте прав доступа
|
||||
- Ошибках сервера (status code >= 400)
|
||||
- Попытке перезаписи без соответствующих прав
|
||||
|
||||
Пример использования:
|
||||
result = client.import_dashboard(Path("my_dashboard.zip"))
|
||||
print(f"Импортирован дашборд: {result['title']}")
|
||||
|
||||
Примечания:
|
||||
- Использует API Endpoint: /dashboard/import/
|
||||
- Автоматически устанавливает overwrite=true для перезаписи существующих дашбордов
|
||||
- Удваивает стандартный таймаут для обработки длительных операций
|
||||
- Удаляет Content-Type заголовок (автоматически генерируется для multipart/form-data)
|
||||
- Архив должен содержать валидные JSON-метаданные в формате Superset
|
||||
- Ответ может содержать информацию о созданных/обновленных ресурсах (датасеты, чарты, владельцы)
|
||||
- При конфликте имен может потребоваться ручное разрешение через параметры импорта
|
||||
"""
|
||||
url = f"{self.config.base_url}/dashboard/import/"
|
||||
|
||||
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:
|
||||
files = {
|
||||
'formData': (
|
||||
zip_name, # Имя файла
|
||||
f, # Файловый объект
|
||||
'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
|
||||
)
|
||||
# Обрабатываем ответ
|
||||
try:
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
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
|
||||
Reference in New Issue
Block a user