maubot-tmdb/tmdb/database.py
Scott Wallace cc89e95403
Refactoring:
* Remove the movie poster parsing
* Fix some code issues
* Increase blurb trunctation length
2024-11-28 09:04:03 +00:00

113 lines
3.6 KiB
Python

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