telegram_bots/aggregator/models.py

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',