replicator/app/services/scheduler.py

# app/services/scheduler.py

import asyncio
from datetime import datetime
from typing import List, Optional
import time as time_module

from app.core.logging import migration_logger
from app.core.config import settings
from app.services.migrator import migrator
from app.repository.replication_metadata_repo import replication_metadata_repo
from app.utils.email_sender import email_sender


class MigrationScheduler:
    """Планировщик миграций на базе PostgreSQL"""
    
    def __init__(self):
        self.running = False
        self.repo = replication_metadata_repo
        
        # Инициализируем расписания по умолчанию
        self._init_default_schedules()
    
    def _init_default_schedules(self):
        """Инициализация расписаний по умолчанию"""
        self.repo.init_default_schedules(settings.TABLES_TO_COPY)
        migration_logger.info("Расписания по умолчанию инициализированы")
    
    def set_schedule(self, table_name: str, schedule_time: str = "00:00", 
                    days: Optional[List[str]] = None, full_reload: bool = False, 
                    enabled: bool = True):
        """Установить расписание для таблицы"""
        if table_name not in settings.TABLES_TO_COPY:
            raise ValueError(f"Таблица {table_name} не найдена в списке для миграции")
        
        # Валидация формата времени
        try:
            datetime.strptime(schedule_time, "%H:%M")
        except ValueError:
            raise ValueError(f"Неверный формат времени: {schedule_time}. Используйте HH:MM")
        
        schedule = self.repo.set_schedule(
            table_name=table_name,
            schedule_time=schedule_time,
            days=days,
            full_reload=full_reload,
            enabled=enabled
        )
        
        if schedule:
            days_str = ', '.join(schedule.days_display) if days else 'все дни'
            migration_logger.info(
                f"Установлено расписание для {table_name}: "
                f"{schedule_time} [{days_str}] (full_reload={full_reload})"
            )
    
    def get_schedule(self, table_name: str):
        """Получить расписание для таблицы"""
        return self.repo.get_schedule(table_name)
    
    def get_all_schedules(self):
        """Получить все расписания"""
        return self.repo.get_all_schedules()
    
    def disable_schedule(self, table_name: str):
        """Отключить расписание"""
        self.repo.disable_schedule(table_name)
    
    def enable_schedule(self, table_name: str):
        """Включить расписание"""
        self.repo.enable_schedule(table_name)
    
    def delete_schedule(self, table_name: str):
        """Удалить расписание"""
        self.repo.delete_schedule(table_name)
    
    def get_due_tables(self, current_time: Optional[datetime] = None) -> List:
        """Получить таблицы для запуска сейчас"""
        due = self.repo.get_due_schedules(current_time)
        return due
    
    async def run_due_migrations(self):
        """Запустить миграции по расписанию"""
        due_schedules = self.get_due_tables()
        
        if not due_schedules:
            return
        
        migration_logger.info(f"Найдено {len(due_schedules)} таблиц для миграции по расписанию")
        
        for schedule in due_schedules:
            try:
                days_str = ', '.join(schedule.days_display)
                migration_logger.info(
                    f"Запуск миграции по расписанию для {schedule.table_name} "
                    f"в {schedule.schedule_time.strftime('%H:%M')} [{days_str}]"
                )
                
                # Запускаем миграцию
                result = await asyncio.to_thread(
                    migrator.run_migration,
                    tables=[schedule.table_name],
                    full_reload=schedule.full_reload,
                    send_email=True
                )
                
                # Обновляем время последнего запуска
                self.repo.update_schedule_last_run(schedule.table_name)
                
                # Логируем успешный запуск
                self.repo.log_operation(
                    table_name=schedule.table_name,
                    operation='SCHEDULED',
                    records_count=0,
                    status='SUCCESS'
                )
                
                migration_logger.info(f"Миграция по расписанию для {schedule.table_name} завершена")
                
            except Exception as e:
                error_msg = f"Ошибка при миграции по расписанию для {schedule.table_name}: {e}"
                migration_logger.error(error_msg)
                
                # Логируем ошибку
                self.repo.log_operation(
                    table_name=schedule.table_name,
                    operation='SCHEDULED',
                    records_count=0,
                    status='ERROR',
                    error_message=str(e)[:500]
                )
                
                email_sender.send_error_notification(
                    error_message=error_msg,
                    table_name=schedule.table_name
                )
    
    def start_scheduler(self):
        """Запустить планировщик"""
        self.running = True
        migration_logger.info("Планировщик миграций запущен")
        
        while self.running:
            try:
                # Проверяем, есть ли таблицы для миграции
                asyncio.run(self.run_due_migrations())
                
                # Ждем 60 секунд до следующей проверки
                time_module.sleep(60)
                
            except KeyboardInterrupt:
                self.stop_scheduler()
                break
            except Exception as e:
                migration_logger.error(f"Ошибка в планировщике: {e}")
                time_module.sleep(60)
    
    def stop_scheduler(self):
        """Остановить планировщик"""
        self.running = False
        migration_logger.info("Планировщик миграций остановлен")


# Глобальный экземпляр
scheduler = MigrationScheduler()