preparation.py

#!/usr/bin/env python3
"""
Функции подготовки для бенчмаркинга облачных PostgreSQL

Этот модуль содержит основные функции для подготовки и настройки баз данных
перед выполнением бенчмаркинга.
"""

import os
import sys
import subprocess
import time
import logging
import psycopg2

from config import benchmark_config
from preparation_utils import (
    test_network_latency,
    create_custom_workload_files,
    cleanup_custom_workload_files,
    vacuum_database,
    drop_pgbench_tables
)

logger = logging.getLogger('pg_benchmark')

def verify_prerequisites() -> str:
    """
    Проверка установки и доступности pgbench.
    
    Returns:
        str: версия pgbench или вызывает исключение, если не найден
    """
    try:
        result = subprocess.run(
            ["pgbench", "--version"], 
            capture_output=True, 
            text=True, 
            check=True
        )
        version = result.stdout.strip()
        logger.info(f"pgbench найден: {version}")
        return version
    except subprocess.CalledProcessError as e:
        logger.error(f"Ошибка выполнения pgbench: {e}")
        raise RuntimeError("Ошибка выполнения pgbench. Проверьте правильность установки.")
    except FileNotFoundError:
        logger.error("pgbench не найден в PATH")
        raise RuntimeError("pgbench не найден. Пожалуйста, установите инструменты клиента PostgreSQL.")

def test_connectivity(config: dict) -> tuple:
    """
    Проверка соединения с сервисом базы данных.
    
    Args:
        config: Словарь конфигурации базы данных
        
    Returns:
        Tuple[bool, float]: Статус соединения и задержка в миллисекундах
    """
    start_time = time.time()
    connection = None
    success = False
    
    try:
        connection_string = (
            f"host={config['host']} "
            f"port={config['port']} "
            f"dbname={config['database']} "
            f"user={config['user']} "
            f"password={config['password']} "
            f"connect_timeout={benchmark_config['connection_timeout']}"
        )
        
        if benchmark_config['use_ssl']:
            connection_string += " sslmode=require"
        
        logger.info(f"Проверка соединения с {config['name']} на {config['host']}")
        
        for attempt in range(benchmark_config['connection_retries']):
            try:
                connection = psycopg2.connect(connection_string)
                cursor = connection.cursor()
                cursor.execute("SELECT version();")
                version = cursor.fetchone()[0]
                cursor.close()
                
                end_time = time.time()
                latency_ms = (end_time - start_time) * 1000
                
                logger.info(f"Успешное соединение с {config['name']}")
                logger.info(f"Версия PostgreSQL: {version}")
                logger.info(f"Задержка соединения: {latency_ms:.2f} мс")
                
                success = True
                break
            except psycopg2.OperationalError as e:
                logger.warning(f"Попытка соединения {attempt+1} не удалась: {e}")
                if attempt < benchmark_config['connection_retries'] - 1:
                    time.sleep(2)  # Ожидание перед повторной попыткой
                else:
                    logger.error(f"Не удалось соединиться после {benchmark_config['connection_retries']} попыток")
                    latency_ms = -1
    except Exception as e:
        logger.error(f"Непредвиденная ошибка во время проверки соединения: {e}")
        latency_ms = -1
    finally:
        if connection:
            connection.close()
    
    return success, latency_ms

def prepare_database(config: dict, scale_factor: int) -> tuple:
    """
    Подготовка базы данных для бенчмаркинга.
    
    Args:
        config: Словарь конфигурации базы данных
        scale_factor: Фактор масштабирования для инициализации pgbench
        
    Returns:
        Tuple[bool, float]: Статус подготовки и время в секундах
    """
    start_time = time.time()
    
    # Создание команды pgbench для инициализации
    pgbench_cmd = [
        "pgbench",
        "-h", config["host"],
        "-p", str(config["port"]),
        "-U", config["user"],
        "-d", config["database"],
        "-i",  # Инициализация
        "-s", str(scale_factor)  # Фактор масштабирования
    ]
    
    # Установка переменной окружения PGPASSWORD
    env = os.environ.copy()
    env["PGPASSWORD"] = config["password"]
    
    try:
        logger.info(f"Инициализация таблиц pgbench для {config['name']} с фактором масштабирования {scale_factor}")
        result = subprocess.run(
            pgbench_cmd,
            env=env,
            capture_output=True,
            text=True,
            check=True
        )
        
        end_time = time.time()
        preparation_time = end_time - start_time
        
        logger.info(f"Подготовка базы данных завершена за {preparation_time:.2f} секунд")
        logger.debug(f"Вывод инициализации: {result.stdout}")
        
        return True, preparation_time
    except subprocess.CalledProcessError as e:
        logger.error(f"Ошибка подготовки базы данных: {e}")
        logger.error(f"Вывод ошибки: {e.stderr}")
        return False, -1