AiOpsWechat/celery_app.py

#from celery import Celery
import celery.schedules
from redbeat import RedBeatSchedulerEntry
from datetime import timedelta
from celery.schedules import crontab

from celery_config import *

from services.redis_service import RedisService
from services.kafka_service import KafkaService
from services.biz_service import BizService

#from config import load_config,conf
from urllib.parse import quote
import asyncio,os,random,sys 
#from tasks import add_friends_task
#from common.log import logger

import logging

logger = logging.getLogger('redbeat')



scheduled_task_sync_wx_info_interval = 6000
scheduled_task_add_contacts_from_chatrooms_interval=10
environment = os.environ.get('environment', 'default') 
# if environment != 'default':
#     scheduled_task_sync_wx_info_interval = 60*11
#     scheduled_task_add_contacts_from_chatrooms_interval = 3600*24
# if environment != 'test':
#     scheduled_task_add_contacts_from_chatrooms_interval = 60*10


if environment == 'production':
    scheduled_task_sync_wx_info_interval = 60*10
    scheduled_task_add_contacts_from_chatrooms_interval = 60*11
elif environment == 'test':
    scheduled_task_sync_wx_info_interval = 60*10
    scheduled_task_add_contacts_from_chatrooms_interval = 60*11
else:
    scheduled_task_sync_wx_info_interval = 1000
    scheduled_task_add_contacts_from_chatrooms_interval=6
    

# 定义定时任务列表 (任务 ID, 任务名称, 执行间隔秒, 任务参数)
# tasks_schedule = [
#     #("redbeat:scheduled_task", "tasks.scheduled_task", 3, []),
#     #("redbeat:scheduled_task_sync_wx", "tasks.scheduled_task_sync_wx", 15, []),
#     ("redbeat:scheduled_task_sync_wx_info", "tasks.scheduled_task_sync_wx_info", scheduled_task_sync_wx_info_interval, [redis_config, kafka_config, gewe_config]),
#     #("redbeat:add_friends_task", "tasks.add_friends_task", random.randint(1, 10), [redis_config]),  # 10分钟执行一次
# ]


# # 注册 RedBeat 任务
# for task_id, task_name, interval, task_args in tasks_schedule:
#     redbeat_entry = RedBeatSchedulerEntry(
#         name=task_id,
#         task=task_name,
#         schedule=celery.schedules.schedule(timedelta(seconds=interval)),
#         args=task_args, 
#         app=celery_app
#     )
#     redbeat_entry.save()


# # 为 add_friends_task 生成随机的分钟和小时
random_minute = random.randint(0, 59)
random_hour = random.randint(0, 23)  # 可以根据需要调整小时范围


# 定义定时任务列表 (任务 ID, 任务名称, 执行间隔或调度, 任务参数)
tasks_schedule = [
    # 其他任务保持不变
    ("redbeat:scheduled_task_sync_wx_info", "tasks.scheduled_task_sync_wx_info", 
     celery.schedules.schedule(timedelta(seconds=scheduled_task_sync_wx_info_interval)), 
     [redis_config, kafka_config, gewe_config]),
    
    # # add_friends_task 使用 crontab 随机时间
    # ("redbeat:add_friends_task", "tasks.add_friends_task", 
    #  crontab(minute=random_minute, hour='*'),  # 每小时在随机分钟执行
    #  # 如果想每天在随机时间执行，可以用: crontab(minute=random_minute, hour=random_hour)
    #  [redis_config]),
     ("redbeat:scheduled_task_add_contacts_from_chatrooms", "tasks.scheduled_task_add_contacts_from_chatrooms", 
     celery.schedules.schedule(timedelta(seconds=scheduled_task_add_contacts_from_chatrooms_interval)), 
     [redis_config, kafka_config, gewe_config]),
]

# # 注册 RedBeat 任务
for task_id, task_name, schedule_obj, task_args in tasks_schedule:
    redbeat_entry = RedBeatSchedulerEntry(
        name=task_id,
        task=task_name,
        schedule=schedule_obj,  # 现在使用 schedule_obj 而不是创建新的 schedule
        args=task_args,
        app=celery_app
    )
    redbeat_entry.save()
    # 如果是 add_friends_task，打印其随机调度信息
    if task_name == "tasks.add_friends_task":
        if isinstance(schedule_obj, crontab):
            print(f"已注册 `{task_name}` 任务，将在每小时的第 {random_minute} 分钟执行")
            logger.info(f"已注册 `{task_name}` 任务，将在每小时的第 {random_minute} 分钟执行")
    else:
        print('scheduled_task_sync_wx_info 定时任务载入成功！')
        logger.info(f"scheduled_task_sync_wx_info 定时任务执行载入成功！")

def setup_schedule():
    # 初始化第一次执行（例如：5秒后执行）
    initial_run_in = 5
    entry = RedBeatSchedulerEntry(
        name='random-task',
        task='tasks.random_scheduled_task',
        schedule=timedelta(seconds=initial_run_in),
        app=celery_app
    )
    entry.save()
    print(f"Initial task scheduled in {initial_run_in} seconds")

setup_schedule()