scripts/wait_for_services.py - dev/test - ump-team/ump-infra

scripts/wait_for_services.py
124 строки4 KB
Dmitry Kochenov
fix: исправить 46 ошибок basedpyright и добавить type hint stubs
09 авг 2026, 15:25
09 авг 2026, 15:258a46eb7
#!/usr/bin/env python3
"""
Скрипт ожидания готовности сервисов (Postgres, Redis, SeaweedFS).

Запуск:
    uv run python scripts/wait_for_services.py

Таймауты:
    POSTGRES_TIMEOUT: 30 сек
    REDIS_TIMEOUT: 30 сек
    SEAWEEDFS_TIMEOUT: 30 сек

Exit codes:
    0 — все сервисы готовы
    1 — один или несколько сервисов не отвечают
"""

import os
import sys
import time

import psycopg2
import redis
import requests

POSTGRES_HOST = os.getenv("POSTGRES_HOST", "localhost")
POSTGRES_PORT = int(os.getenv("POSTGRES_PORT", 5432))
POSTGRES_USER = os.getenv("POSTGRES_USER", "app")
POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD", "app")
POSTGRES_DB = os.getenv("POSTGRES_DB", "app")
POSTGRES_TIMEOUT = int(os.getenv("POSTGRES_TIMEOUT", 30))

REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
REDIS_TIMEOUT = int(os.getenv("REDIS_TIMEOUT", 30))

SEAWEEDFS_HOST = os.getenv("SEAWEEDFS_HOST", "localhost")
SEAWEEDFS_PORT = int(os.getenv("SEAWEEDFS_PORT", 9333))
SEAWEEDFS_TIMEOUT = int(os.getenv("SEAWEEDFS_TIMEOUT", 30))


def wait_for_postgres():
    """Ожидание готовности Postgres."""
    print(f"Ожидание Postgres на {POSTGRES_HOST}:{POSTGRES_PORT}...")
    start = time.time()
    while time.time() - start < POSTGRES_TIMEOUT:
        try:
            conn = psycopg2.connect(
                host=POSTGRES_HOST,
                port=POSTGRES_PORT,
                user=POSTGRES_USER,
                password=POSTGRES_PASSWORD,
                database=POSTGRES_DB,
            )
            conn.close()
            print(f"✓ Postgres готов через {time.time() - start:.1f} сек")
            return True
        except psycopg2.OperationalError as e:
            print(f"  Postgres не готов: {e}")
            time.sleep(1)
    print("✗ Postgres не отвечает")
    return False


def wait_for_redis():
    """Ожидание готовности Redis."""
    print(f"Ожидание Redis на {REDIS_HOST}:{REDIS_PORT}...")
    start = time.time()
    while time.time() - start < REDIS_TIMEOUT:
        try:
            r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0)
            r.ping()
            print(f"✓ Redis готов через {time.time() - start:.1f} сек")
            return True
        except redis.ConnectionError as e:
            print(f"  Redis не готов: {e}")
            time.sleep(1)
    print("✗ Redis не отвечает")
    return False


def wait_for_seaweedfs():
    """Ожидание готовности SeaweedFS."""
    print(f"Ожидание SeaweedFS на {SEAWEEDFS_HOST}:{SEAWEEDFS_PORT}...")
    start = time.time()
    while time.time() - start < SEAWEEDFS_TIMEOUT:
        try:
            response = requests.get(f"http://{SEAWEEDFS_HOST}:{SEAWEEDFS_PORT}/cluster/health", timeout=5)
            if response.status_code == 200:
                print(f"✓ SeaweedFS готов через {time.time() - start:.1f} сек")
                return True
        except requests.exceptions.RequestException as e:
            print(f"  SeaweedFS не готов: {e}")
            time.sleep(1)
    print("✗ SeaweedFS не отвечает")
    return False


def main():
    """Запуск проверок всех сервисов."""
    print("=" * 60)
    print("ОЖИДАНИЕ ГОТОВНОСТИ СЕРВИСОВ")
    print("=" * 60)

    results = []
    results.append(("Postgres", wait_for_postgres()))
    results.append(("Redis", wait_for_redis()))
    results.append(("SeaweedFS", wait_for_seaweedfs()))

    print("=" * 60)
    all_ok = all(r[1] for r in results)
    if all_ok:
        print("✓ Все сервисы готовы!")
        return 0
    else:
        print("✗ Некоторые сервисы не отвечают:")
        for name, ok in results:
            status = "OK" if ok else "FAIL"
            print(f"  - {name}: {status}")
        return 1


if __name__ == "__main__":
    sys.exit(main())