lono/db.py

from BTrees.IIBTree import IIBTree
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()
    # 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()
    # end migrations
    commit()
    return conn


class Subscriber(Persistent):
    def __init__(self, user_id, name):
        self.uid = user_id
        self.name = name
        self.messages_forward = IIBTree()
        self.messages_reverse = IIBTree()

    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))