From fbfe72c774bf725e84caee3ee3640e3d833cf509 Mon Sep 17 00:00:00 2001
From: bakatrouble <bakatrouble@gmail.com>
Date: Mon, 28 Jan 2019 00:30:40 +0300
Subject: [PATCH] load last 10 messages on chat add

---
 aggregator/client.py |  6 +++++-
 aggregator/models.py |  9 ++++++---
 aggregator/tasks.py  | 14 ++++++++++++++
 3 files changed, 25 insertions(+), 4 deletions(-)
 create mode 100644 aggregator/tasks.py

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)