diff options
Diffstat (limited to 'code/utils')
| -rw-r--r-- | code/utils/config.py | 24 | ||||
| -rw-r--r-- | code/utils/content_view.py | 47 | ||||
| -rw-r--r-- | code/utils/database.py | 9 | ||||
| -rw-r--r-- | code/utils/jellyfin_create.py | 21 | ||||
| -rw-r--r-- | code/utils/jellyfin_delete.py | 60 | ||||
| -rw-r--r-- | code/utils/models.py | 31 |
6 files changed, 108 insertions, 84 deletions
diff --git a/code/utils/config.py b/code/utils/config.py index 1e7fdc6..ba1fac5 100644 --- a/code/utils/config.py +++ b/code/utils/config.py @@ -4,7 +4,6 @@ import sys import os import logging import requests -import sqlite3 from colorlog import ColoredFormatter @@ -108,10 +107,9 @@ schema = { def load_config() -> None: """ - Load DB, then load and validate the config file + Load the config file and validate it If the file does not exist, generate it """ - database_setup() if os.path.exists("/.dockerenv"): file_path = "config/config.yaml" else: @@ -157,26 +155,6 @@ jellyfin: ) -def database_setup() -> None: - """ - Create the database if it does not exist - """ - if not os.path.exists("data"): - os.makedirs("data") - db = sqlite3.connect("data/cordarr.db") - cursor = db.cursor() - cursor.execute( - "CREATE TABLE IF NOT EXISTS requests (title TEXT, release_year TEXT," - " local_id INTEGER, tmdbid INTEGER, tvdbid INTEGER, user_id INTEGER)" - ) - cursor.execute( - "CREATE TABLE IF NOT EXISTS jellyfin_accounts (user_id INTEGER," - " jellyfin_user_id INTEGER, deletion_time DATETIME)" - ) - db.commit() - db.close() - - def validate_config(contents) -> None: """ Validate the contents of the config file and assign variables diff --git a/code/utils/content_view.py b/code/utils/content_view.py index 7a982c0..4f5673e 100644 --- a/code/utils/content_view.py +++ b/code/utils/content_view.py @@ -1,6 +1,7 @@ import discord -import sqlite3 +from utils.models import Requests +from utils.database import Session from utils.content_add import add_content """ @@ -163,30 +164,26 @@ class RequestButtonView(discord.ui.View): ) # Keep track of the requests for the `/status` command - db = sqlite3.connect("data/cordarr.db") - cursor = db.cursor() - cursor.execute( - "INSERT INTO requests (title, release_year, local_id, tmdbid," - " tvdbid, user_id) VALUES (?, ?, ?, ?, ?, ?)", - ( - self.content_info["title"], - self.content_info["year"], - local_id, - ( - self.content_info["contentId"] - if self.service == "radarr" - else None - ), - ( - None - if self.service == "radarr" - else self.content_info["contentId"] - ), - interaction.user.id, - ), - ) - db.commit() - db.close() + with Session() as session: + session.add( + Requests( + title=self.content_info["title"], + release_year=self.content_info["year"], + local_id=local_id, + tmdbid=( + self.content_info["contentId"] + if self.service == "radarr" + else None + ), + tvdbid=( + None + if self.service == "radarr" + else self.content_info["contentId"] + ), + user_id=interaction.user.id, + ) + ) + session.commit() @discord.ui.button(label="Don't Request", style=discord.ButtonStyle.danger) async def dont_request_button( diff --git a/code/utils/database.py b/code/utils/database.py new file mode 100644 index 0000000..cb2ecf7 --- /dev/null +++ b/code/utils/database.py @@ -0,0 +1,9 @@ +from sqlalchemy import create_engine +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import sessionmaker + +database_url = "sqlite:///data/cordarr.db" + +engine = create_engine(database_url) +Session = sessionmaker(bind=engine) +Base = declarative_base() diff --git a/code/utils/jellyfin_create.py b/code/utils/jellyfin_create.py index e860c2b..03c6c8d 100644 --- a/code/utils/jellyfin_create.py +++ b/code/utils/jellyfin_create.py @@ -1,10 +1,11 @@ import datetime import requests import random -import sqlite3 from wonderwords import RandomWord from string import ascii_lowercase, digits +from utils.database import Session +from utils.models import JellyfinAccounts from utils.config import ( JELLYFIN_URL, JELLYFIN_HEADERS, @@ -67,14 +68,14 @@ def create_jellyfin_account(user_id): return False # Add the information to the database - db = sqlite3.connect("data/cordarr.db") - cursor = db.cursor() - cursor.execute( - "INSERT INTO jellyfin_accounts (user_id, jellyfin_user_id," - " deletion_time) VALUES (?, ?, ?)", - (user_id, jellyfin_user_id, deletion_time), - ) - db.commit() - db.close() + with Session() as session: + session.add( + JellyfinAccounts( + user_id=user_id, + jellyfin_user_id=jellyfin_user_id, + deletion_time=deletion_time, + ) + ) + session.commit() return username, password diff --git a/code/utils/jellyfin_delete.py b/code/utils/jellyfin_delete.py index 66af00b..0e65201 100644 --- a/code/utils/jellyfin_delete.py +++ b/code/utils/jellyfin_delete.py @@ -1,8 +1,9 @@ import datetime -import sqlite3 import requests -from utils.config import JELLYFIN_URL, JELLYFIN_HEADERS +from utils.database import Session +from utils.models import JellyfinAccounts +from utils.config import LOG, JELLYFIN_URL, JELLYFIN_HEADERS def delete_accounts(): @@ -10,29 +11,36 @@ def delete_accounts(): Delete Jellyfin accounts that have passed their deletion time """ # Get all expired Jellyfin accounts - db = sqlite3.connect("data/cordarr.db") - cursor = db.cursor() - cursor.execute( - "SELECT jellyfin_user_id FROM jellyfin_accounts WHERE" - " deletion_time < ?", - (datetime.datetime.now(),), - ) - jellyfin_user_ids = cursor.fetchall() - - # Delete the Jellyfin accounts - for jellyfin_user_id in jellyfin_user_ids: - request = requests.delete( - f"{JELLYFIN_URL}/Users/{jellyfin_user_id[0]}", - headers=JELLYFIN_HEADERS, + with Session() as session: + jellyfin_user_ids = ( + session.query(JellyfinAccounts.jellyfin_user_id) + .filter(JellyfinAccounts.deletion_time < datetime.datetime.now()) + .all() ) - # If 204 - account deleted - # If 404 - account not found - # Either way, remove account from database - if request.status_code in (404, 204): - cursor.execute( - "DELETE FROM jellyfin_accounts WHERE jellyfin_user_id = ?", - (jellyfin_user_id,), - ) - db.commit() - db.close() + # Delete each account + for jellyfin_user_id in jellyfin_user_ids: + print(f"Deleting account {jellyfin_user_id[0]}") + try: + response = requests.delete( + f"{JELLYFIN_URL}/Users/{jellyfin_user_id[0]}", + headers=JELLYFIN_HEADERS, + ) + response.raise_for_status() + # Get the account and delete it + account = ( + session.query(JellyfinAccounts) + .filter( + JellyfinAccounts.jellyfin_user_id + == jellyfin_user_id[0] + ) + .first() + ) + session.delete(account) + except: + LOG.error( + "Failed deleting Jellyfin account w/ ID" + f" {jellyfin_user_id[0]}" + ) + # Commit changes + session.commit() diff --git a/code/utils/models.py b/code/utils/models.py new file mode 100644 index 0000000..3dbd6be --- /dev/null +++ b/code/utils/models.py @@ -0,0 +1,31 @@ +from sqlalchemy import ( + Column, + Integer, + String, + DateTime, + BigInteger, +) + + +from utils.database import Base + + +class Requests(Base): + __tablename__ = "requests" + + id = Column(Integer, primary_key=True, index=True) + title = Column(String) + release_year = Column(Integer) + local_id = Column(Integer) + tmdbid = Column(Integer) + tvdbid = Column(Integer) + user_id = Column(BigInteger) + + +class JellyfinAccounts(Base): + __tablename__ = "jellyfin_accounts" + + id = Column(Integer, primary_key=True, index=True) + user_id = Column(BigInteger) + jellyfin_user_id = Column(String) + deletion_time = Column(DateTime) |