140 lines
5.8 KiB
Python
140 lines
5.8 KiB
Python
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() |