first commit

app/taskiq/broker.py

@@ -0,0 +1,140 @@
+import asyncio
+from datetime import datetime, time
+import os
+from typing import Any, Dict
+
+from taskiq import ScheduledTask, TaskiqScheduler
+from taskiq_redis import ListQueueBroker, ListRedisScheduleSource, RedisAsyncResultBackend, RedisStreamBroker, RedisScheduleSource
+
+import logging
+logging.getLogger("taskiq.scheduler").setLevel(logging.DEBUG)
+logging.getLogger("taskiq.broker").setLevel(logging.DEBUG)
+
+# ---------- Настройка результата ----------
+result_backend = RedisAsyncResultBackend(
+    redis_url="redis://127.0.0.1:6379/0",
+    result_ex_time=86400,          # результаты хранятся 24 часа
+)
+
+# ---------- Настройка брокера (очередь с подтверждениями) ----------
+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) ----------
+schedule_source = ListRedisScheduleSource(
+    url="redis://127.0.0.1:6379/0",
+    prefix="migration_schedule",            # префикс для ключей в Redis
+)
+
+
+# ---------- Планировщик с кастомным startup ----------
+class SchedulerWithStartup(TaskiqScheduler):
+    """Планировщик с синхронизацией расписаний при запуске."""
+    
+    async def startup(self) -> None:
+        """Запуск планировщика с синхронизацией расписаний."""
+        from app.scheduler import sync_schedules_to_redis
+        from app.core.logging import migration_logger
+        
+        migration_logger.info("Запуск планировщика taskiq...")
+        # Синхронизируем расписания из БД в Redis
+        await sync_schedules_to_redis()
+        migration_logger.info("Планировщик готов к работе")
+        await super().startup()
+
+
+scheduler = SchedulerWithStartup(
+    broker=broker,
+    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():
+    from app.services.scheduler import scheduler as db_scheduler
+    from app.core.logging import migration_logger
+    migration_logger.info("Обновление расписаний в Redis...")
+    # Получаем все активные расписания из БД
+    schedules = db_scheduler.get_all_schedules()
+    # Очищаем старые ключи в Redis (можно через прямой redis)
+    import redis.asyncio as redis
+    r = await redis.from_url("redis://127.0.0.1:6379/0")
+    keys = await r.keys("migration_schedule:*")
+    if keys:
+        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}
+        )
+        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}
+
+# Startup для планировщика (выполняется при запуске через CLI)
+# @broker.on_event('on_startup')
+# async def on_startup():
+    # await refresh_schedules()