From b0ea7ab93564f1b2f004f7ea74783508f12f4ff6 Mon Sep 17 00:00:00 2001 From: Parker Date: Tue, 21 Jan 2025 20:38:33 -0600 Subject: [PATCH] Fixes + Use SQLAlchemy --- code/bot.py | 10 ++---- code/cogs/newaccount.py | 31 ++++++++++------- code/cogs/status.py | 63 +++++++++++++++++++---------------- code/utils/config.py | 24 +------------ code/utils/content_view.py | 47 ++++++++++++-------------- code/utils/database.py | 9 +++++ code/utils/jellyfin_create.py | 21 ++++++------ code/utils/jellyfin_delete.py | 60 ++++++++++++++++++--------------- code/utils/models.py | 31 +++++++++++++++++ requirements.txt | 3 +- 10 files changed, 165 insertions(+), 134 deletions(-) create mode 100644 code/utils/database.py create mode 100644 code/utils/models.py diff --git a/code/bot.py b/code/bot.py index b7654c6..b067f73 100644 --- a/code/bot.py +++ b/code/bot.py @@ -1,10 +1,9 @@ import discord from discord.ext import commands -from discord.ext import tasks import os +from utils.database import Base, engine import utils.config as config -from utils.jellyfin_delete import delete_accounts class MyBot(commands.Bot): @@ -15,7 +14,6 @@ class MyBot(commands.Bot): ) async def setup_hook(self): - delete_old_temp_accounts.start() for ext in os.listdir("./code/cogs"): if ext.endswith(".py"): await self.load_extension(f"cogs.{ext[:-3]}") @@ -30,11 +28,7 @@ async def on_ready(): config.LOG.info(f"{bot.user} has connected to Discord.") -@tasks.loop(seconds=60) -async def delete_old_temp_accounts(): - delete_accounts() - - if __name__ == "__main__": + Base.metadata.create_all(bind=engine) config.load_config() bot.run(config.BOT_TOKEN) diff --git a/code/cogs/newaccount.py b/code/cogs/newaccount.py index 8e17cda..f3f5ac8 100644 --- a/code/cogs/newaccount.py +++ b/code/cogs/newaccount.py @@ -1,9 +1,11 @@ import discord from discord import app_commands -from discord.ext import commands -import sqlite3 +from discord.ext import commands, tasks +from utils.database import Session from utils.jellyfin_create import create_jellyfin_account +from utils.jellyfin_delete import delete_accounts +from utils.models import JellyfinAccounts from utils.config import ( JELLYFIN_PUBLIC_URL, JELLYFIN_ENABLED, @@ -15,19 +17,20 @@ class NewAccount(commands.Cog): def __init__(self, bot): self.bot = bot + def cog_load(self): + self.delete_accounts_loop.start() + @app_commands.command() @app_commands.check(lambda inter: JELLYFIN_ENABLED) async def newaccount(self, interaction: discord.Interaction) -> None: """Create a new temporary Jellyfin account""" # Make sure the user doesn't already have an account - db = sqlite3.connect("data/cordarr.db") - cursor = db.cursor() - cursor.execute( - "SELECT * FROM jellyfin_accounts WHERE user_id = ?", - (interaction.user.id,), - ) - account = cursor.fetchone() - db.close() + with Session() as session: + account = ( + session.query(JellyfinAccounts) + .filter(JellyfinAccounts.user_id == interaction.user.id) + .first() + ) # Account already allocated if account: embed = discord.Embed( @@ -63,8 +66,8 @@ class NewAccount(commands.Cog): title="Jellyfin Account Information", description=( # fmt: off - "Here is your temporary account information.\n\n", - f"**Server URL:** `[{JELLYFIN_PUBLIC_URL}]({JELLYFIN_PUBLIC_URL})`\n" + "Here is your temporary account information.\n\n" + f"**Server URL:** `{JELLYFIN_PUBLIC_URL}`\n" f"**Username:** `{response[0]}`\n" f"**Password:** `{response[1]}`\n\n" "Your account will be automatically deleted in" @@ -88,6 +91,10 @@ class NewAccount(commands.Cog): embed=embed, ephemeral=True ) + @tasks.loop(minutes=1) + async def delete_accounts_loop(self): + delete_accounts() + async def setup(bot): await bot.add_cog(NewAccount(bot)) diff --git a/code/cogs/status.py b/code/cogs/status.py index 3a64146..769fc8f 100644 --- a/code/cogs/status.py +++ b/code/cogs/status.py @@ -2,8 +2,9 @@ import discord from discord import app_commands from discord.ext import commands import requests -import sqlite3 +from utils.models import Requests +from utils.database import Session from utils.config import ( RADARR_HOST_URL, RADARR_HEADERS, @@ -22,15 +23,18 @@ class Status(commands.Cog): # Defer the response await interaction.response.defer(ephemeral=True) - db = sqlite3.connect("data/cordarr.db") - cursor = db.cursor() - cursor.execute( - "SELECT title, release_year, local_id, tmdbid, tvdbid FROM" - " requests WHERE user_id = ?", - (interaction.user.id,), - ) - requested_content = cursor.fetchall() - db.close() + with Session() as session: + requested_content = ( + session.query( + Requests.title, + Requests.release_year, + Requests.local_id, + Requests.tmdbid, + Requests.tvdbid, + ) + .filter(Requests.user_id == interaction.user.id) + .all() + ) # No content requested if len(requested_content) == 0: @@ -42,7 +46,7 @@ class Status(commands.Cog): ), color=0xD01B86, ) - return await interaction.followup.send(embed=embed, ephemeral=True) + return await interaction.followup.send(embed=embed) # Create template embed embed = discord.Embed( @@ -75,7 +79,7 @@ class Status(commands.Cog): embed.description += radarr_desc + sonarr_desc + non_queue_desc # Send the follow-up message - await interaction.edit_original_response(embed=embed, ephemeral=True) + await interaction.followup.send(embed=embed) def unpack_content(self, requested_content: list) -> tuple: """ @@ -92,18 +96,18 @@ class Status(commands.Cog): sonarr_content_info = {} for content in requested_content: - title, (release_year), local_id, tmdbid, tvdbid = content + title, release_year, local_id, tmdbid, tvdbid = content if tmdbid is not None: - radarr_content_info[int(local_id)] = { + radarr_content_info[local_id] = { "title": title, - "release_year": int(release_year), - "tmdbid": int(tmdbid), + "release_year": release_year, + "tmdbid": tmdbid, } else: - sonarr_content_info[int(local_id)] = { + sonarr_content_info[local_id] = { "title": title, - "release_year": int(release_year), - "tvdbid": int(tvdbid), + "release_year": release_year, + "tvdbid": tvdbid, } return radarr_content_info, sonarr_content_info @@ -169,7 +173,7 @@ class Status(commands.Cog): for content in requested_content: title, release_year, local_id, tmdbid, _ = content # If not in queue - if int(local_id) not in added_ids: + if local_id not in added_ids: # Pull the movie data from the service if tmdbid is not None: data = requests.get( @@ -185,15 +189,16 @@ class Status(commands.Cog): # If the movie has a file, then it has finished downloading if data.get("hasFile", True): # Remove from database - db = sqlite3.connect("data/cordarr.db") - cursor = db.cursor() - cursor.execute( - "DELETE FROM requests WHERE user_id = ? AND" - " local_id = ?", - (user_id, int(local_id)), - ) - db.commit() - db.close() + with Session() as session: + request = ( + session.query(Requests) + .filter(Requests.user_id == user_id) + .filter(Requests.local_id == local_id) + .first() + ) + session.delete(request) + session.commit() + # If series and only a portion of episodes have been downloaded if data.get("statistics").get("percentOfEpisodes"): description += ( 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) diff --git a/requirements.txt b/requirements.txt index 97b43e4..f6e03c3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ wonderwords==2.2.0 PyYAML==6.0.2 jsonschema==4.23.0 jsonschema-specifications==2024.10.1 -discord.py==2.4.0 \ No newline at end of file +discord.py==2.4.0 +SQLAlchemy==2.0.37 \ No newline at end of file