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


# load_config()

# KAFKA_BOOTSTRAP_SERVERS = conf().get("kafka_bootstrap_servers")
# KAFKA_TOPIC = 'topic.ai.ops.wx'
# KAFKA_GROUP_ID = 'ai-ops-wx'

# redis_host=conf().get("redis_host")
# redis_port=conf().get("redis_port")
# redis_password=conf().get("redis_password") 
# redis_db=conf().get("redis_db") 
# encoded_password = quote(redis_password) 
# # 配置 Celery
# celery_app = Celery(
#     "worker",
#     broker=f"redis://:{encoded_password}@{redis_host}:{redis_port}/{redis_db}",
#     backend=f"redis://:{encoded_password}@{redis_host}:{redis_port}/{redis_db}",
#     include=['tasks']
# )

# # 配置 redbeat 作为 Celery Beat 调度器
# celery_app.conf.update(
#     timezone="Asia/Shanghai",  # 设定时区
#     beat_scheduler="redbeat.RedBeatScheduler",  # 使用 RedBeat 作为调度器
#     redbeat_redis_url=f"redis://:{encoded_password}@{redis_host}:{redis_port}/{redis_db}"  # redbeat 存储任务调度信息的 Redis
#     ,
#     redbeat_lock_timeout=60,  # 避免多个 Beat 实例冲突
#     beat_max_loop_interval=5  # 让 Celery Beat 每 5 秒检查一次任务
# )

# task_name = "tasks.scheduled_task"
# # 任务执行间隔（每 10 秒执行一次）
# schedule = celery.schedules.schedule(timedelta(seconds=3))

# # RedBeat 任务唯一 ID
# redbeat_entry = RedBeatSchedulerEntry(
#     name="redbeat:scheduled_task",  # 任务 ID
#     task=task_name,  # 任务名称
#     schedule=schedule,  # 任务调度时间
#     args=[],
#     app=celery_app
        
# )

# # 保存任务到 Redis
# redbeat_entry.save()


# tasks_schedule = [
#     ("redbeat:scheduled_task", "tasks.scheduled_task", 3),
#     ("redbeat:scheduled_task_sync_wx", "tasks.scheduled_task_sync_wx", 15),
# ]

# # 创建并保存 RedBeat 任务
# for task_id, task_name, interval in tasks_schedule:
#     redbeat_entry = RedBeatSchedulerEntry(
#         name=task_id,
#         task=task_name,
#         schedule=celery.schedules.schedule(timedelta(seconds=interval)),
#         args=[],
#         app=celery_app
#     )
#     redbeat_entry.save()

# # 获取配置文件中的 redis_config、kafka_config、gewe_config
# redis_config = {
#     'host': redis_host,
#     'port': redis_port,
#     'password': redis_password,
#     'db': redis_db,
# }


# kafka_config = {
#     'bootstrap_servers': KAFKA_BOOTSTRAP_SERVERS,
#     'topic': KAFKA_TOPIC,
#     'group_id': KAFKA_GROUP_ID,
# }
# gewe_config = {
#     'api_url': "http://api.geweapi.com/gewe",
# }

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
    

# 定义定时任务列表 (任务 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():
    # 初始化第一次执行（例如：15秒后执行）
    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()