diff --git a/aggregator/client.py b/aggregator/client.py index e721ab3..429fcd3 100644 --- a/aggregator/client.py +++ b/aggregator/client.py @@ -15,7 +15,11 @@ Session.notice_displayed = True def get_client(): config._reload_maybe() - session_path = os.path.relpath(default_storage.path(config.pyrogram_session.replace('.session', ''))) + pyrogram_sesssion = config.pyrogram_session.replace('.session', '') + if not pyrogram_sesssion or not config.pyrogram_app_id or config.pyrogram_app_hash: + raise RuntimeError('Pyrogram is not configured') + session_path = os.path.relpath(default_storage.path(pyrogram_sesssion)) + return Client(session_path, config.pyrogram_app_id, config.pyrogram_app_hash) diff --git a/aggregator/models.py b/aggregator/models.py index 1060d60..9c54e1b 100644 --- a/aggregator/models.py +++ b/aggregator/models.py @@ -5,8 +5,9 @@ from tempfile import TemporaryDirectory import pytz from django.conf import settings from django.db import models -from pyrogram import Chat as PyrogramChat, Message as PyrogramMessage, ChatPhoto as PyrogramChatPhoto -from pyrogram.api.types.chat_photo import ChatPhoto as MTProtoChatPhoto +from pyrogram import Chat as PyrogramChat, Message as PyrogramMessage + +from aggregator.tasks import collect_new_messages class AggregationSource(models.Model): @@ -27,7 +28,7 @@ class Chat(models.Model): @classmethod def from_obj(cls, chat: PyrogramChat, client): - obj, _ = Chat.objects.update_or_create( + obj, created = Chat.objects.update_or_create( chat_id=chat.id, defaults={ 'title': chat.title or '{} {}'.format(chat.first_name, chat.last_name).rstrip(), @@ -47,6 +48,8 @@ class Chat(models.Model): obj.photo.save(os.path.basename(path), f, save=True) obj.photo_id = chat.photo.small_file_id obj.save() + if created: + collect_new_messages.delay(obj.pk) return obj def __str__(self): diff --git a/aggregator/tasks.py b/aggregator/tasks.py new file mode 100644 index 0000000..fd76906 --- /dev/null +++ b/aggregator/tasks.py @@ -0,0 +1,14 @@ +from celery_once import QueueOnce + +from aggregator.client import get_client +from config.celery import app + +from .models import Chat +from .client import collect_new_messages as _collect_new_messages + + +@app.task(base=QueueOnce, once={'keys': ['chat_id'], 'graceful': True}) +def collect_new_messages(chat_id): + chat = Chat.objects.get(pk=chat_id) + with get_client() as client: + _collect_new_messages(client, chat)