100 lines
3.6 KiB
Python
100 lines
3.6 KiB
Python
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',
|