ss-tools/examples/maintenance/maintenance-api-python.py

#!/usr/bin/env python3
"""
Example: External tool triggers maintenance via superset-tools API.

This script demonstrates how an ETL pipeline (Airflow DAG, cron job, CI/CD)
can start and end maintenance banners on Superset dashboards using the
superset-tools API Key authentication.

Requirements: Python 3.10+, requests (pip install requests)

Usage:
  # Start maintenance
  python maintenance-api-python.py start \
    --api-key ssk_M7xqaP2zL9vR4nF8kC1bH5jT6dW0yA3 \
    --base-url https://superset-tools.example.com \
    --tables public.messages,public.users \
    --duration-hours 4 \
    --message "Scheduled ETL refresh"

  # Same, but auto-end the maintenance at end_time (API param auto_end=true)
  python maintenance-api-python.py start \
    --api-key ssk_M7xqaP2zL9vR4nF8kC1bH5jT6dW0yA3 \
    --base-url https://superset-tools.example.com \
    --tables public.messages \
    --duration-hours 4 \
    --auto-end

  # End maintenance by event ID
  python maintenance-api-python.py end \
    --api-key ssk_M7xqaP2zL9vR4nF8kC1bH5jT6dW0yA3 \
    --base-url https://superset-tools.example.com \
    --event-id m-abc123

  # End ALL maintenance (use with caution!)
  python maintenance-api-python.py end-all \
    --api-key ssk_M7xqaP2zL9vR4nF8kC1bH5jT6dW0yA3 \
    --base-url https://superset-tools.example.com \
    --environment ss-dev

=============================== API SPECIFICATION ============================
СПЕЦИФИКАЦИЯ API ОБСЛУЖИВАНИЯ
--------------------------------
Базовый URL:     {BASE_URL}/api/maintenance
Аутентификация:  заголовок  X-API-Key: <API_KEY>   (или JWT-токен)

РАЗРЕШЕНИЯ API-КЛЮЧА (permissions), требуемые скриптом:
  maintenance:start      — создание события (POST /start)
  maintenance:end        — завершение конкретного события (POST /{id}/end)
  maintenance:end_all    — массовое завершение всех событий (POST /end-all)

Ответы мутационных эндпоинтов — это сырые тела схем (без обёртки), например
{ "task_id": "...", "maintenance_id": "...", "status": "pending" }.
Все мутационные эндпоинты возвращают HTTP 202 и task_id — операция
выполняется асинхронно планировщиком (TaskManager).

--- POST /api/maintenance/start ---
  HTTP 202 Accepted | Разрешение: maintenance:start
  Тело запроса (JSON):
    tables          [string]  ОБЯЗАТЕЛЬНО. Список таблиц (1..100).
    start_time      datetime  ОБЯЗАТЕЛЬНО. Начало окна (ISO 8601).
    end_time        datetime  Опционально. Должно быть позже start_time.
    auto_end        bool      По умолчанию false. Если true и задан end_time,
                              планировщик сам завершит обслуживание в end_time.
    message         string    Опционально. Текст баннера (макс. 500 симв.).
    environment_id  string    ОБЯЗАТЕЛЬНО. Целевое окружение (ss-dev и т.п.).
  Ответ 202:  { "task_id": "...", "maintenance_id": "...", "status": "pending" }
  Ошибки:
    400 — невалидные данные (end_time <= start_time; start_time слишком в
          прошлом, допуск 1 час; > 100 таблиц);
    401 — неверный/отозванный ключ;
    403 — недостаточно прав;
    404 — неизвестное environment_id;
    409 — идемпотентность (такое же tables/start_time/end_time уже активно):
          { "maintenance_id": "...", "status": "already_active" }.

--- POST /api/maintenance/{maintenance_id}/end ---
  HTTP 202 Accepted | Разрешение: maintenance:end | Тело: отсутствует.
  Ответ 202:  { "task_id": "...", "status": "pending" }
  Ошибки:
    404 — событие не найдено;
    401/403 — аутентификация/права.
  Идемпотентность: повторный end уже завершённого события — тоже 202,
  но с status == "already_completed" (считается успехом).

--- POST /api/maintenance/end-all ---
  HTTP 202 Accepted | Разрешение: maintenance:end_all
  ОСТОРОЖНО: снимает баннеры со ВСЕХ дашбордов.
  Тело (JSON, опционально): { "environment_id": "..." }
    Если не указан — используется окружение из ключа API.
  Ответ 202:  { "task_id": "...", "status": "pending" }

--- Прочие (read-only) эндпоинты модуля, требуют READ ---
  GET  /api/maintenance/events               — списки active и completed событий
  GET  /api/maintenance/dashboard-banners    — состояние баннеров по дашбордам
  POST /api/maintenance/preview-dashboards   — какие дашборды затронут таблицы
  GET  /api/maintenance/settings             — настройки обслуживания
  PUT  /api/maintenance/settings             — обновить настройки (только admin)

=========================== ИНСТРУКЦИЯ (РУССКИЙ) ==============================
Назначение: ставить и снимать баннер «ведутся технические работы» на
дашбордах Superset при обновлении таблиц (ETL, cron, CI/CD).

ПРЕДВАРИТЕЛЬНЫЕ ШАГИ:
  1. Установите зависимости:  pip install requests
  2. Создайте API-ключ superset-tools с нужными разрешениями
     (maintenance:start / maintenance:end / maintenance:end_all).
  3. Укажите --base-url (адрес сервера) и --api-key.

КОМАНДЫ:
  Обслуживание public.messages на 4 часа в ss-dev:
    python maintenance-api-python.py start \
        --api-key ssk_... --base-url https://superset-tools.example.com \
        --tables public.messages --environment ss-dev --duration-hours 4
  Несколько таблиц, с сообщением и автоснятием баннера:
    python maintenance-api-python.py start \
        --api-key ssk_... --tables public.messages,public.users \
        --environment ss-prod --duration-hours 4 \
        --message "Плановый ETL" --auto-end
  Завершить конкретное событие:
    python maintenance-api-python.py end \
        --api-key ssk_... --event-id m-abc123
  Аварийно снять ВСЕ баннеры в ss-dev:
    python maintenance-api-python.py end-all \
        --api-key ssk_... --environment ss-dev

О ПАРАМЕТРЕ --auto-end:
  start_time и end_time вычисляются автоматически (сейчас и сейчас+часы).
  Без --auto-end end_time справочный — баннер снимают вручную командой `end`.
  С --auto-end баннер снимут автоматически ровно в end_time.

ВЫХОДНОЙ КОД: 0 — успех; 1 — ошибка (HTTPError от сервера).

ЗАМЕЧАНИЯ ПО БЕЗОПАСНОСТИ:
  - Не передавайте ключ API через аргументы в реальных CI/cron, если это
    небезопасно; используйте защищённые переменные окружения/секреты.
  - `end-all` снимает баннеры со всех дашбордов — используйте с осторожностью.
================================================================================
"""

import argparse
import sys
from datetime import datetime, timezone, timedelta

try:
    import requests
except ImportError:
    print("Error: 'requests' library is required. Install with: pip install requests")
    sys.exit(1)


API_KEY_HEADER = "X-API-Key"


def start_maintenance(base_url: str, api_key: str, tables: list[str],
                      environment_id: str, duration_hours: int = 4,
                      message: str | None = None, auto_end: bool = False,
                      timeout: int = 30) -> dict:
    """
    Start a maintenance event.

    Args:
        base_url: superset-tools base URL (e.g. https://superset-tools.example.com)
        api_key: API key with 'maintenance:start' permission
        tables: List of table names to flag (e.g. ["public.messages"])
        environment_id: Superset environment (e.g. "ss-dev", "ss-prod")
        duration_hours: How long the maintenance window lasts
        message: Optional custom message for the banner
        auto_end: When True, sends auto_end=true so the scheduler removes the banner
            automatically at end_time. When False (default), end_time is informational
            only and maintenance must be ended manually via `end`.
        timeout: Request timeout in seconds

    Returns:
        API response dict with task_id and maintenance_id

    Raises:
        requests.HTTPError: On 4xx/5xx response with error details

    Note:
        Idempotent retry: identical (tables, start_time, end_time) against an active
        event returns 409 with status='already_active' and the existing maintenance_id
        — treated as success here.
    """
    now = datetime.now(timezone.utc)
    start_time = now.isoformat()
    end_time = (now + timedelta(hours=duration_hours)).isoformat()

    payload = {
        "tables": tables,
        "start_time": start_time,
        "end_time": end_time,
        "environment_id": environment_id,
    }
    if auto_end:
        payload["auto_end"] = True
    if message:
        payload["message"] = message

    print(f"[maintenance] Starting maintenance for tables: {tables}")
    print(f"[maintenance]   Window: {start_time} → {end_time}")
    print(f"[maintenance]   Environment: {environment_id}")
    if auto_end:
        print("[maintenance]   Auto-end: ENABLED (banner will be removed at end_time)")

    response = requests.post(
        f"{base_url}/api/maintenance/start",
        json=payload,
        headers={API_KEY_HEADER: api_key},
        timeout=timeout,
    )

    if response.status_code == 202:
        result = response.json()
        print(f"[maintenance] ✅ Event created (id={result['maintenance_id']})")
        print(f"[maintenance]    Task: {result.get('task_id', 'N/A')}")
        return result

    if response.status_code == 409:
        result = response.json()
        if result.get("status") == "already_active":
            print(f"[maintenance] ℹ️ Event already active (idempotent, id={result.get('maintenance_id', 'N/A')})")
            return result

    if response.status_code == 400:
        error = _format_error(response)
        print(f"[maintenance] ❌ Validation error: {error}")
    elif response.status_code == 401:
        print("[maintenance] ❌ Authentication failed: invalid or revoked API key")
    elif response.status_code == 403:
        print("[maintenance] ❌ Permission denied: API key lacks 'maintenance:start'")
    else:
        print(f"[maintenance] ❌ Unexpected error ({response.status_code}): {response.text[:500]}")

    response.raise_for_status()
    return {}  # unreachable


def end_maintenance(base_url: str, api_key: str, event_id: str,
                    timeout: int = 30) -> dict:
    """
    End a specific maintenance event and remove its banners.

    Args:
        base_url: superset-tools base URL
        api_key: API key with 'maintenance:end' permission
        event_id: The event ID returned by start_maintenance
        timeout: Request timeout in seconds

    Returns:
        API response dict

    Raises:
        requests.HTTPError: On error
    """
    print(f"[maintenance] Ending maintenance event: {event_id}")

    response = requests.post(
        f"{base_url}/api/maintenance/{event_id}/end",
        headers={API_KEY_HEADER: api_key},
        timeout=timeout,
    )

    if response.status_code == 202:
        result = response.json()
        if result.get("status") == "already_completed":
            print("[maintenance] ℹ️  Event was already completed (idempotent)")
            return result
        print(f"[maintenance] ✅ Event end scheduled (task={result.get('task_id', 'N/A')})")
        return result

    if response.status_code == 404:
        print(f"[maintenance] ❌ Event not found: {event_id}")
    elif response.status_code == 401:
        print("[maintenance] ❌ Authentication failed: invalid or revoked API key")
    elif response.status_code == 403:
        print("[maintenance] ❌ Permission denied: API key lacks 'maintenance:end'")
    else:
        print(f"[maintenance] ❌ Unexpected error ({response.status_code}): {response.text[:500]}")

    response.raise_for_status()
    return {}


def end_all_maintenance(base_url: str, api_key: str,
                        environment_id: str | None = None,
                        timeout: int = 30) -> dict:
    """
    End ALL active maintenance events in the given environment.
    USE WITH CAUTION — this removes banners from all dashboards.

    Args:
        base_url: superset-tools base URL
        api_key: API key with 'maintenance:end_all' permission
        environment_id: If set, only end events in this environment.
                        If None, uses the API key's default scope.
        timeout: Request timeout in seconds

    Returns:
        API response dict

    Raises:
        requests.HTTPError: On error
    """
    print(f"[maintenance] Ending ALL maintenance events "
          f"(env={environment_id or 'key-scope'})")
    print("[maintenance] ⚠️  This will remove banners from ALL affected dashboards!")

    payload = {}
    if environment_id:
        payload["environment_id"] = environment_id

    response = requests.post(
        f"{base_url}/api/maintenance/end-all",
        json=payload,
        headers={API_KEY_HEADER: api_key},
        timeout=timeout,
    )

    if response.status_code == 202:
        result = response.json()
        print(f"[maintenance] ✅ All events end scheduled (task={result.get('task_id', 'N/A')})")
        return result

    if response.status_code == 401:
        print("[maintenance] ❌ Authentication failed: invalid or revoked API key")
    elif response.status_code == 403:
        print("[maintenance] ❌ Permission denied: API key lacks 'maintenance:end_all'")
    elif response.status_code == 400:
        print(f"[maintenance] ❌ {_format_error(response)}")
    else:
        print(f"[maintenance] ❌ Unexpected error ({response.status_code}): {response.text[:500]}")

    response.raise_for_status()
    return {}


def _format_error(response: requests.Response) -> str:
    """Extract human-readable error from API response."""
    try:
        data = response.json()
        if "detail" in data:
            detail = data["detail"]
            if isinstance(detail, list):
                return "; ".join(d.get("msg", str(d)) for d in detail)
            return str(detail)
        return response.text[:200]
    except Exception:
        return response.text[:200]


def main():
    parser = argparse.ArgumentParser(
        description="superset-tools Maintenance Banner CLI",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Start 4-hour maintenance on public.messages
  %(prog)s start --api-key ssk_... --tables public.messages \\\\
      --environment ss-dev --duration-hours 4

  # End maintenance by event ID
  %(prog)s end --api-key ssk_... --event-id m-abc123

  # Emergency: end ALL maintenance in ss-dev
  %(prog)s end-all --api-key ssk_... --environment ss-dev
        """,
    )

    # Common options are defined on each subparser (via parents) so they can be
    # placed AFTER the subcommand — matching the usage examples above. argparse
    # does not reliably accept parent-parser options after the subcommand.
    common_opts = argparse.ArgumentParser(add_help=False)
    common_opts.add_argument(
        "--base-url",
        default="http://localhost:8000",
        help="superset-tools base URL (default: %(default)s)",
    )
    common_opts.add_argument(
        "--api-key",
        required=True,
        help="API key with required permissions",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    # start
    start_parser = subparsers.add_parser("start", parents=[common_opts], help="Start maintenance")
    start_parser.add_argument(
        "--tables",
        required=True,
        help="Comma-separated table names (e.g. 'public.messages,public.users')",
    )
    start_parser.add_argument(
        "--environment",
        required=True,
        dest="environment_id",
        help="Superset environment ID (e.g. 'ss-dev', 'ss-prod')",
    )
    start_parser.add_argument(
        "--duration-hours",
        type=int,
        default=4,
        help="Maintenance window in hours (default: %(default)s)",
    )
    start_parser.add_argument(
        "--message",
        default=None,
        help="Optional message for the banner",
    )
    start_parser.add_argument(
        "--auto-end",
        action="store_true",
        help="Also send auto_end=true so the scheduler removes the banner automatically "
        "at end_time. Without it, end_time is informational only.",
    )

    # end
    end_parser = subparsers.add_parser("end", parents=[common_opts], help="End a specific maintenance event")
    end_parser.add_argument(
        "--event-id",
        required=True,
        help="Maintenance event ID (from start response)",
    )

    # end-all
    end_all_parser = subparsers.add_parser(
        "end-all", parents=[common_opts], help="End ALL maintenance events (use with caution)"
    )
    end_all_parser.add_argument(
        "--environment",
        dest="environment_id",
        default=None,
        help="Limit to specific environment (optional)",
    )

    args = parser.parse_args()

    try:
        if args.command == "start":
            tables = [t.strip() for t in args.tables.split(",")]
            start_maintenance(
                base_url=args.base_url,
                api_key=args.api_key,
                tables=tables,
                environment_id=args.environment_id,
                duration_hours=args.duration_hours,
                message=args.message,
                auto_end=args.auto_end,
            )
        elif args.command == "end":
            end_maintenance(
                base_url=args.base_url,
                api_key=args.api_key,
                event_id=args.event_id,
            )
        elif args.command == "end-all":
            end_all_maintenance(
                base_url=args.base_url,
                api_key=args.api_key,
                environment_id=args.environment_id,
            )
    except requests.HTTPError:
        # Specific error message was already printed by the command function.
        sys.exit(1)
    except requests.RequestException as exc:
        # Network-level failure (connection refused, timeout, DNS, ...).
        print(f"[maintenance] ❌ Request failed: {exc}")
        sys.exit(1)


if __name__ == "__main__":
    main()