telegram_bots/bots/management/commands/bot_worker.py

66 lines
2.3 KiB
Python
Raw Normal View History

2019-11-12 17:13:56 +00:00
import logging
import traceback
2020-12-11 10:51:27 +00:00
from concurrent.futures.thread import ThreadPoolExecutor
from multiprocessing.pool import ThreadPool
2019-11-12 17:13:56 +00:00
import sentry_sdk
from django.core.cache import cache
from django.core.management import BaseCommand
from telegram import TelegramError, Update
2019-11-12 17:13:56 +00:00
from telegram.error import TimedOut
from telegram.ext import CallbackContext
2019-11-12 17:13:56 +00:00
from bots.models import TelegramBot
class Command(BaseCommand):
def handle(self, *args, **options):
pool = ThreadPool(8)
def error_handler(update: Update, ctx: CallbackContext):
sentry_sdk.capture_exception(ctx.error)
logging.exception('Exception while processing update', exc_info=ctx.error)
2021-03-11 21:48:40 +00:00
initialized = False
2019-11-12 17:13:56 +00:00
dispatchers = []
2020-12-11 10:51:27 +00:00
def check_updates(dispatcher):
try:
updates = dispatcher.bot.get_updates(dispatcher.last_update_id)
except TimedOut:
return
except TelegramError as e:
sentry_sdk.capture_exception(e)
traceback.print_exc()
updates = []
for update in updates:
pool.apply_async(dispatcher.process_update, (update,))
dispatcher.last_update_id = update.update_id + 1
2019-11-12 17:13:56 +00:00
while True:
try:
2021-03-11 21:48:40 +00:00
if not initialized or cache.get('bots_reset'):
2019-11-12 17:13:56 +00:00
logging.warning('Reloading dispatchers')
dispatchers = []
for bot in TelegramBot.objects.filter(active=True):
try:
dispatcher = bot.build_dispatcher(error_handler)
2019-11-12 17:13:56 +00:00
dispatcher.last_update_id = 0
dispatchers.append(dispatcher)
except TelegramError:
pass
cache.delete('bots_reset')
2021-03-11 21:48:40 +00:00
initialized = True
2019-11-12 17:13:56 +00:00
2020-12-11 10:51:27 +00:00
with ThreadPoolExecutor() as executor:
for dispatcher in dispatchers:
executor.submit(check_updates, dispatcher)
2019-11-12 17:13:56 +00:00
except KeyboardInterrupt:
pool.terminate()
pool.join()
2019-11-12 17:13:56 +00:00
return
except Exception as e:
sentry_sdk.capture_exception(e)
traceback.print_exc()