diff --git a/code/bot.py b/code/bot.py index b7654c6..919f394 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 +from discord.ext import commands, 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,7 @@ class MyBot(commands.Bot): ) async def setup_hook(self): - delete_old_temp_accounts.start() + delete_accounts_task.start() for ext in os.listdir("./code/cogs"): if ext.endswith(".py"): await self.load_extension(f"cogs.{ext[:-3]}") @@ -30,8 +29,11 @@ async def on_ready(): config.LOG.info(f"{bot.user} has connected to Discord.") -@tasks.loop(seconds=60) -async def delete_old_temp_accounts(): +@tasks.loop(minutes=1) +async def delete_accounts_task(): + from utils.jellyfin_delete import delete_accounts + + Base.metadata.create_all(bind=engine) delete_accounts() diff --git a/code/cogs/newaccount.py b/code/cogs/newaccount.py index 8e17cda..1fca84c 100644 --- a/code/cogs/newaccount.py +++ b/code/cogs/newaccount.py @@ -1,9 +1,10 @@ import discord from discord import app_commands from discord.ext import commands -import sqlite3 +from utils.database import Session from utils.jellyfin_create import create_jellyfin_account +from utils.models import JellyfinAccounts from utils.config import ( JELLYFIN_PUBLIC_URL, JELLYFIN_ENABLED, @@ -19,15 +20,15 @@ class NewAccount(commands.Cog): @app_commands.check(lambda inter: JELLYFIN_ENABLED) async def newaccount(self, interaction: discord.Interaction) -> None: """Create a new temporary Jellyfin account""" + # Defer in case it takes too long + await interaction.response.defer(ephemeral=True) # 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( @@ -39,9 +40,7 @@ class NewAccount(commands.Cog): ), color=0xD01B86, ) - return await interaction.response.send_message( - embed=embed, ephemeral=True - ) + return await interaction.followup.send(embed=embed) # Create a new Jellyfin account for the user response = create_jellyfin_account(interaction.user.id) @@ -54,17 +53,15 @@ class NewAccount(commands.Cog): ), color=0xD01B86, ) - await interaction.response.send_message( - embed=embed, ephemeral=True - ) + await interaction.followup.send(embed=embed) # Send the user their account information embed = discord.Embed( 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" @@ -84,9 +81,7 @@ class NewAccount(commands.Cog): ), color=0xD01B86, ) - return await interaction.response.send_message( - embed=embed, ephemeral=True - ) + return await interaction.followup.send(embed=embed) async def setup(bot): diff --git a/code/cogs/request.py b/code/cogs/request.py index 2ad71ba..eeeb467 100644 --- a/code/cogs/request.py +++ b/code/cogs/request.py @@ -31,6 +31,8 @@ class Request(commands.Cog): name: str, ) -> None: """Request a movie or tv show to be added to the library""" + # Could take a sec. so defer the response + await interaction.response.defer(ephemeral=True) # Get matching content from relevant service if form == "Movie": content_data = get_content( @@ -54,9 +56,7 @@ class Request(commands.Cog): ), color=0xD01B86, ) - return await interaction.response.send_message( - embed=embed, ephemeral=True - ) + return await interaction.followup.send(embed=embed, ephemeral=True) if content_data == "ALREADY ADDED": embed = discord.Embed( @@ -70,9 +70,7 @@ class Request(commands.Cog): ), color=0xD01B86, ) - return await interaction.response.send_message( - embed=embed, ephemeral=True - ) + return await interaction.followup.send(embed=embed, ephemeral=True) embed = discord.Embed( title="Results Found", @@ -103,9 +101,7 @@ class Request(commands.Cog): SONARR_QUALITY_PROFILE_ID, ) - await interaction.response.send_message( - embed=embed, view=view, ephemeral=True - ) + await interaction.followup.send(embed=embed, view=view, ephemeral=True) async def setup(bot): diff --git a/code/cogs/status.py b/code/cogs/status.py index 90d0c8b..ef044cc 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, @@ -20,17 +21,20 @@ class Status(commands.Cog): async def status(self, interaction: discord.Interaction) -> None: """Get the status of the movies you have requested""" # Defer the response - await interaction.response.defer() + 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,9 +46,7 @@ class Status(commands.Cog): ), color=0xD01B86, ) - return await interaction.response.send_message( - embed=embed, ephemeral=True - ) + return await interaction.followup.send(embed=embed) # Create template embed embed = discord.Embed( @@ -76,8 +78,8 @@ class Status(commands.Cog): embed.description += radarr_desc + sonarr_desc + non_queue_desc - # Send the embed - await interaction.response.send_message(embed=embed, ephemeral=True) + # Send the follow-up message + await interaction.followup.send(embed=embed) def unpack_content(self, requested_content: list) -> tuple: """ @@ -94,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 @@ -171,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( @@ -187,23 +189,26 @@ 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 += ( - f"\n**{title} ({release_year})** - Status: `NOT" - " FOUND" - f" ({int(data['statistics']['percentOfEpisodes'])}%" - " of eps.)`" - ) + # If data["statistics"] exists and is not None + if "statistics" in data and data["statistics"] != None: + if "percentOfEpisodes" in data["statistics"]: + description += ( + f"\n**{title} ({release_year})** - Status: `NOT" + " FOUND" + f" ({int(data['statistics']['percentOfEpisodes'])}%" + " of eps.)`" + ) # All other scenarios, download not found else: 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 668954f..89b5e89 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 """ @@ -147,7 +148,8 @@ class RequestButtonView(discord.ui.View): ), color=0xD01B86, ) - await interaction.response.send_message(embed=embed) + await interaction.response.edit_message(view=None) + await interaction.followup.send(embed=embed) # Alert the user that the content failed to be added else: embed = discord.Embed( @@ -158,33 +160,30 @@ class RequestButtonView(discord.ui.View): f" {self.service} library." ), ) - return await interaction.response.send_message(embed=embed) + await interaction.delete_original_response() + return await interaction.response.edit_message(embed=embed) # 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( @@ -200,4 +199,4 @@ class RequestButtonView(discord.ui.View): ), color=0xD01B86, ) - await interaction.response.send_message(embed=embed) + await interaction.response.edit_message(embed=embed, view=None) diff --git a/code/utils/database.py b/code/utils/database.py new file mode 100644 index 0000000..cd757ec --- /dev/null +++ b/code/utils/database.py @@ -0,0 +1,13 @@ +from sqlalchemy import create_engine +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import sessionmaker +import os + +if not os.path.exists("data"): + os.makedirs("data") + +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