import os from datetime import datetime from tempfile import TemporaryDirectory import pytz from django.conf import settings from django.db import models, transaction from pyrogram.types import Chat as PyrogramChat, Message as PyrogramMessage class AggregationSource(models.Model): owner = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE) title = models.CharField(max_length=64) chat_id = models.BigIntegerField(db_index=True) def __str__(self): return self.title class Chat(models.Model): chat_id = models.BigIntegerField(db_index=True) title = models.TextField() username = models.CharField(max_length=64, null=True, blank=True) photo = models.ImageField(null=True, blank=True) photo_id = models.CharField(max_length=64, null=True, blank=True) @classmethod def from_obj(cls, chat: PyrogramChat, client): from aggregator.tasks import collect_new_messages obj, created = Chat.objects.update_or_create( chat_id=chat.id, defaults={ 'title': chat.title or '{} {}'.format(chat.first_name, chat.last_name).rstrip(), 'username': chat.username, } ) if chat.photo is None: if obj.photo_id is not None: obj.photo = obj.photo_id = None obj.save() else: photo_file_id = chat.photo.small_file_id if photo_file_id != obj.photo_id: with TemporaryDirectory() as d: path = client.download_media(photo_file_id, os.path.join(d, ''), block=True) with open(path, 'rb') as f: obj.photo.save(os.path.basename(path), f, save=True) obj.photo_id = chat.photo.small_file_id obj.save() if created: transaction.on_commit(lambda: collect_new_messages.delay(obj.pk)) return obj def __str__(self): return '{} (chat_id="{}")'.format(self.title, self.chat_id) class MessageManager(models.Manager): def active_messages(self): return self.get_queryset().filter(deleted=False, replaced_by__isnull=True) class Message(models.Model): chat = models.ForeignKey(Chat, on_delete=models.CASCADE, related_name='messages') message_id = models.PositiveIntegerField(db_index=True) text = models.TextField(blank=True) date = models.DateTimeField() edit_date = models.DateTimeField(null=True, blank=True) deleted = models.BooleanField(default=False) replaced_by = models.ForeignKey('self', on_delete=models.CASCADE, null=True, blank=True) objects = MessageManager() @classmethod def from_obj(cls, message: PyrogramMessage, client): tz = pytz.timezone('UTC') chat = Chat.from_obj(message.chat, client) try: old = Message.objects.active_messages().get(chat=chat, message_id=message.message_id) except Message.DoesNotExist: old = None obj = Message.objects.create( chat=chat, message_id=message.message_id, text=message.text.html if message.text else '', date=tz.localize(datetime.utcfromtimestamp(message.date)), edit_date=tz.localize(datetime.utcfromtimestamp(message.edit_date)) if message.edit_date else None, ) if old is not None: old.replaced_by = obj old.save() return obj def __str__(self): return 'id="{}" (chat_id="{}")'.format(self.message_id, self.chat.chat_id) class Meta: ordering = 'message_id', 'edit_date',