maubot-tmdb/tmdb/database.py

"""
This file is part of tmdb-bot.

tmdb-bot is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General  Public License version 3 as published by
the Free Software Foundation.

tmdb-bot is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General  Public License for more details.

You should have received a copy of the GNU Affero General  Public License
along with tmdb-bot.  If not, see <https://www.gnu.org/licenses/>.
"""

# pylint: disable=missing-class-docstring,missing-function-docstring

from time import time

from sqlalchemy import Column, Float, Integer, MetaData, String, Table, select
from sqlalchemy.engine.base import Engine


class Database:
    db: Engine

    def __init__(self, db: Engine) -> None:
        self.db = db
        meta = MetaData()
        meta.bind = db

        self.language = Table(
            "tmdb_language",
            meta,
            Column("id", Integer, primary_key=True, autoincrement=True),
            Column("user_id", String(255), nullable=False),
            Column("language", String(255), nullable=False),
        )
        self.tmdb_poster_size = Table(
            "tmdb_poster_size",
            meta,
            Column("id", Integer, primary_key=True, autoincrement=True),
            Column("user_id", String(255), nullable=False),
            Column("size", String(255), nullable=False),
        )
        self.tmdb_messages = Table(
            "tmdb_messages",
            meta,
            Column("timestamp", Float, primary_key=True),
            Column("event_id", String(64), nullable=False),
            Column("result_json", String(255), nullable=False),
        )
        meta.create_all(db)

    def set_message(self, event_id, result_json):
        with self.db.begin() as tx:
            timestamp = time()
            tx.execute(
                self.tmdb_messages.insert().values(
                    timestamp=timestamp, event_id=event_id, result_json=result_json
                )
            )

    def get_message(self, event_id):
        rows = self.db.execute(
            select([self.tmdb_messages.c.result_json]).where(
                self.tmdb_messages.c.event_id == event_id
            )
        )
        row = rows.fetchone()
        if row:
            return row["result_json"]
        return None

    def set_language(self, user_id, language):
        with self.db.begin() as tx:
            tx.execute(self.language.delete().where(self.language.c.user_id == user_id))
            tx.execute(
                self.language.insert().values(user_id=user_id, language=language)
            )

    def get_language(self, user_id):
        rows = self.db.execute(
            select([self.language.c.language]).where(self.language.c.user_id == user_id)
        )
        row = rows.fetchone()
        if row:
            return row["language"]
        return None

    def set_poster_size(self, user_id, size):
        with self.db.begin() as tx:
            tx.execute(
                self.tmdb_poster_size.delete().where(
                    self.tmdb_poster_size.c.user_id == user_id
                )
            )
            tx.execute(
                self.tmdb_poster_size.insert().values(user_id=user_id, size=size)
            )

    def get_poster_size(self, user_id):
        rows = self.db.execute(
            select([self.tmdb_poster_size.c.size]).where(
                self.tmdb_poster_size.c.user_id == user_id
            )
        )
        row = rows.fetchone()
        if row:
            return row["size"]
        return None