lono/db.py

51 lines
1.6 KiB
Python
Raw Normal View History

2019-08-19 19:05:45 +00:00
from datetime import datetime
2019-03-15 23:27:42 +00:00
from BTrees.IIBTree import IIBTree
2019-03-03 23:23:12 +00:00
from ZODB import DB
from ZODB.Connection import Connection
from ZODB.FileStorage import FileStorage
from persistent import Persistent
from persistent.mapping import PersistentMapping
from transaction import commit
from telegram import Message, Chat
def get_conn(read_only=False) -> Connection:
storage = FileStorage('users.fs', read_only=read_only)
db = DB(storage)
conn = db.open()
if not hasattr(conn.root, 'subscribers'):
conn.root.subscribers = PersistentMapping()
2019-03-15 23:27:42 +00:00
# migration 1
if not hasattr(conn.root, 'counter'):
conn.root.counter = 0
for user in conn.root.subscribers.values(): # type: Subscriber
if not hasattr(user, 'messages_forward') or not isinstance(user.messages_forward, IIBTree):
user.messages_forward = IIBTree()
user.messages_reverse = IIBTree()
2019-08-19 19:05:45 +00:00
# migration 2
if not hasattr(conn.root, 'last_media'):
conn.root.last_media = datetime.now()
2019-03-15 23:27:42 +00:00
# end migrations
commit()
2019-03-03 23:23:12 +00:00
return conn
class Subscriber(Persistent):
def __init__(self, user_id, name):
self.uid = user_id
self.name = name
2019-03-15 23:27:42 +00:00
self.messages_forward = IIBTree()
self.messages_reverse = IIBTree()
2019-03-03 23:23:12 +00:00
def update_from_message(self, m: Message):
self.name = Subscriber.get_name(m.chat)
@classmethod
def get_name(cls, chat: Chat):
return f'{chat.first_name or ""} {chat.last_name or ""} {chat.title or ""}'.strip()
@classmethod
def from_chat(cls, chat: Chat):
return cls(chat.id, cls.get_name(chat))