telegram_bots/aggregator/models.py

101 lines
3.6 KiB
Python
Raw Normal View History

2019-01-27 21:20:09 +00:00
import os
from datetime import datetime
from tempfile import TemporaryDirectory
import pytz
from django.conf import settings
2019-01-27 22:02:25 +00:00
from django.db import models, transaction
2019-01-27 21:30:40 +00:00
from pyrogram import Chat as PyrogramChat, Message as PyrogramMessage
2019-01-27 21:20:09 +00:00
class AggregationSource(models.Model):
owner = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)
title = models.CharField(max_length=64)
2019-01-27 21:49:17 +00:00
chat_id = models.CharField(db_index=True, max_length=64)
2019-01-27 21:20:09 +00:00
def __str__(self):
return self.title
class Chat(models.Model):
2019-01-27 21:49:17 +00:00
chat_id = models.CharField(db_index=True, max_length=64)
2019-01-27 21:20:09 +00:00
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):
2019-01-27 21:32:34 +00:00
from aggregator.tasks import collect_new_messages
2019-01-27 21:30:40 +00:00
obj, created = Chat.objects.update_or_create(
2019-01-27 21:20:09 +00:00
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()
2019-01-27 21:30:40 +00:00
if created:
2019-01-27 22:02:25 +00:00
transaction.commit()
2019-01-27 21:30:40 +00:00
collect_new_messages.delay(obj.pk)
2019-01-27 21:20:09 +00:00
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',