Небольшие изменения

app/taskiq/broker.py

@@ -1,14 +1,11 @@
 import asyncio
-from datetime import datetime, time
-import os
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from taskiq import ScheduledTask, TaskiqScheduler
-from taskiq_redis import ListQueueBroker, ListRedisScheduleSource, RedisAsyncResultBackend, RedisStreamBroker, RedisScheduleSource
+from taskiq_redis import ListQueueBroker, ListRedisScheduleSource, RedisAsyncResultBackend
 
-import logging
-logging.getLogger("taskiq.scheduler").setLevel(logging.DEBUG)
-logging.getLogger("taskiq.broker").setLevel(logging.DEBUG)
+from app.models.replication import ReplicationSchedule
+from app.services.task_tracker import task_tracker
 
 # ---------- Настройка результата ----------
 result_backend = RedisAsyncResultBackend(
@@ -20,8 +17,6 @@ result_backend = RedisAsyncResultBackend(
 broker = ListQueueBroker(
     url="redis://127.0.0.1:6379/0",
     queue_name="migration_queue",          # имя очереди
-    #consumer_group_name="migration_group", # группа потребителей
-    #maxlen=1000,                           # максимальная длина очереди
 ).with_result_backend(result_backend)
 
 # ---------- Источник расписаний (динамический, на Redis) ----------
@@ -52,36 +47,6 @@ scheduler = SchedulerWithStartup(
     sources=[schedule_source],
 )
 
-# ---------- Задача для миграции ----------
-@broker.task(
-    task_name="migrate_table_task",  # Короткое имя для удобства
-    retry_on_error=True,
-    max_retries=3,
-    retry_delay=60
-)
-async def migrate_table_task(table_name: str, full_reload: bool = False) -> Dict[str, Any]:
-    """
-    Асинхронная задача миграции.
-    Здесь вызывается синхронный код миграции (через asyncio.to_thread, чтобы не блокировать).
-    """
-    from app.services.migrator import migrator
-    from app.core.logging import migration_logger
-
-    migration_logger.info(f"Запуск миграции для {table_name} (full_reload={full_reload})")
-
-    try:
-        # Запускаем синхронную функцию в отдельном потоке, чтобы не блокировать event loop
-        result = await asyncio.to_thread(
-            migrator.run_migration,
-            tables=[table_name],
-            full_reload=full_reload,
-            send_email=True
-        )
-        complex.info(f"Миграция {table_name} завершена")
-        return {"success": True, "table": table_name, "result": result}
-    except Exception as e:
-        migration_logger.error(f"Ошибка миграции {table_name}: {e}")
-        raise  # для retry
 
 # Функция синхронизации расписаний из БД в Redis (общая для startup и API)
 async def refresh_schedules():
@@ -98,43 +63,145 @@ async def refresh_schedules():
         await r.delete(*keys)
     await r.close()
 
-    for s in schedules:
-        if not s.enabled:
-            continue
-        hour = s.schedule_time.hour
-        minute = s.schedule_time.minute
-        days = s.days_list
-        cron = f"{minute} {hour} * * {','.join(map(str, days))}" if len(days) < 6 else f"{minute} {hour} * * *"
-        # await migrate_table_task.schedule_by_cron(
-        #     schedule_source,
-        #     cron,
-        #     s.table_name,
-        #     s.full_reload
-        # )
-        task = ScheduledTask(
-            task_name="migrate_table_task",
-            labels={},
-            cron=cron,
-            cron_offset='Asia/Tokyo',  # UTC+9
-            args=[s.table_name],
-            kwargs={"full_reload": s.full_reload}
+    # Задача для проверки расписания
+    checker_task = ScheduledTask(
+        task_name="check_schedules_task",
+        labels={},
+        cron="* * * * *",  # ← Каждую минуту
+        cron_offset='Asia/Tokyo',
+        args=[],
+        kwargs={}
+    )
+    await schedule_source.add_schedule(checker_task)
+
+    # for s in schedules:
+    #     if s.enabled is not True:
+    #         continue
+    #     hour = s.schedule_time.hour
+    #     minute = s.schedule_time.minute
+    #     days = s.days_list
+    #     cron = f"{minute} {hour} * * {','.join(map(str, days))}" if len(days) < 6 else f"{minute} {hour} * * *"
+
+    #     task = ScheduledTask(
+    #         task_name="migrate_table_task",
+    #         labels={},
+    #         cron=cron,
+    #         cron_offset='Asia/Tokyo',  # UTC+9
+    #         args=[],
+    #         kwargs={
+    #             "table_name": s.table.table_name,
+    #             "schedule_id": s.id,
+    #             "metadata_id": s.metadata_id,
+    #             "life_table_name": s.table.life_table_name,
+    #             "uses_life": s.table.has_use_life(),
+    #             "full_reload": s.full_reload
+    #         }
+    #     )
+    #     await schedule_source.add_schedule(task)
+
+    migration_logger.info("Расписание запущено")
+
+# ---------- Задача для миграции ----------
+@broker.task(
+    task_name="migrate_table_task",
+    retry_on_error=True,
+    max_retries=3,
+    retry_delay=60
+)
+async def migrate_table_task(
+    table_name: str, 
+    schedule_id: int,
+    metadata_id: int, 
+    life_table_name: Optional[str] = None,
+    uses_life: bool = False,
+    full_reload: bool = False,
+    batch_id: Optional[str] = None,
+    send_email: bool = False 
+    ) -> Dict[str, Any]:
+    """
+    Асинхронная задача миграции.
+    Здесь вызывается синхронный код миграции (через asyncio.to_thread, чтобы не блокировать).
+    """
+    from app.services.migrator import migrator
+    from app.core.logging import migration_logger
+
+    migration_logger.info(f"Запуск миграции: {table_name} (life_table_name={life_table_name}, uses_life={uses_life}, full_reload={full_reload}, batch_id={batch_id})")
+
+    try:
+        # Запускаем синхронную функцию в отдельном потоке, чтобы не блокировать event loop
+        result = await asyncio.to_thread(
+            migrator.run_migration,
+            table_name=table_name,
+            schedule_id=schedule_id,
+            metadata_id=metadata_id,
+            life_table_name=life_table_name,
+            uses_life=uses_life,
+            full_reload=full_reload,
+            send_email=False
         )
-        await schedule_source.add_schedule(task)
-        # schedules_itms = await schedule_source.get_schedules()
-        # for itm in schedules_itms:
-        #     print(itm)
-    migration_logger.info("Расписания обновлены")
 
-@broker.task
-async def test_task(message: str = "Hello, world!"):
-    """Простая тестовая задача."""
-    import logging
-    logging.basicConfig(level=logging.INFO)
-    logger = logging.getLogger(__name__)
-    logger.info(f"Тестовая задача выполнена: {message}")
-    return {"status": "ok", "message": message}
+        task_result = {"success": result, "table": table_name, "schedule_id": schedule_id}
 
-# Startup для планировщика (выполняется при запуске через CLI)
-# @broker.on_event('on_startup')
-# async def on_startup():
-    # await refresh_schedules()
+        # Если есть batch_id — отмечаем выполнение
+        if batch_id:
+            completed = await task_tracker.mark_completed(batch_id, task_result)
+            status = await task_tracker.get_batch_status(batch_id)
+            migration_logger.info(f"Задача завершена: {completed}/{status['total']} в батче {batch_id}")
+            
+            # Если все задачи выполнены — запускаем отправку email
+            if completed >= status["total"] and send_email:
+                await send_batch_email_task.kiq(batch_id=batch_id)
+        
+        return task_result
+    except Exception as e:
+        migration_logger.error(f"Ошибка миграции {table_name}: {e}")
+
+        # Даже при ошибке отмечаем задачу
+        if batch_id:
+            await task_tracker.mark_completed(
+                batch_id, 
+                {"success": False, "table": table_name, "error": str(e)}
+            )
+        raise  # для retry
+
+# ---------- Задача для отправки email ----------
+@broker.task(task_name="send_batch_email_task")
+async def send_batch_email_task(batch_id: str) -> Dict[str, Any]:
+    """Отправить сводный email после завершения всех задач"""
+    from app.core.logging import migration_logger
+    from app.services.task_tracker import task_tracker
+    from app.utils.email_sender import email_sender
+
+    migration_logger.info(f"📧 Отправка сводного email для батча {batch_id}")
+    
+    status = await task_tracker.get_batch_status(batch_id)
+    
+    # Формируем отчёт
+    successful = sum(1 for r in status["results"] if r.get("success"))
+    failed = len(status["results"]) - successful
+    
+    email_data = {
+        "batch_id": batch_id,
+        "total": status["total"],
+        "completed": status["completed"],
+        "successful": successful,
+        "failed": failed,
+        "results": status["results"]
+    }
+    
+    # Отправляем email
+    await email_sender.send_migration_summary_email(email_data)
+    
+    migration_logger.info(f"Сводный email отправлен: {successful} успешно, {failed} ошибок")
+    
+    return {"email_sent": True, "batch_id": batch_id}
+
+# ---------- Задача для проверки расписания ----------
+@broker.task(task_name="check_schedules_task")
+async def check_schedules_task():
+    """
+    Задача-планировщик: проверяет расписания и запускает батчи.
+    Запускается по cron (каждую минуту).
+    """
+    from app.services.scheduler import scheduler
+    await scheduler.check_and_run_schedules()

app/taskiq/tasks.py

@@ -0,0 +1,6 @@
+# app/tasks.py
+import asyncio
+from typing import Any, Dict, Optional
+from app.services.task_tracker import task_tracker
+from app.taskiq.broker import broker
+