finished queue system, overwrites discord kill process to loop

cogs/music/main.py

@@ -1,20 +1,15 @@
-import discord
 from discord.ext import commands
 from discord.ext.commands.context import Context
+
 import cogs.music.util as util
 import cogs.music.queue as queue
+import cogs.music.translate as translate
 
 import datetime
 import pytz
 
-import yt_dlp
-
 from cogs.music.help import music_help
 
-ydl_opts = {
-        'format': 'bestaudio/best',
-        'outtmpl': 'downloads/%(title)s.%(ext)s',
-}
 
 class music(commands.Cog):
     def __init__(self, client):
@@ -27,6 +22,8 @@ class music(commands.Cog):
         help_command.cog = self
         self.help_command = help_command
 
+        queue.initialize_tables()
+
 
     @commands.command(
             help="Displays latency from the bot",
@@ -59,20 +56,26 @@ class music(commands.Cog):
     @commands.command(
             help="Queues a song into the bot",
             aliases=['p', 'qeue', 'q'])
+    @commands.check(util.in_server)
     async def play(self, ctx: Context, *, url=None):
         if url is None:
             raise commands.CommandError("Must provide a link or search query")
 
+        server = ctx.guild.id
+
+        #TODO potentially save requests before getting stream link
+        audio = translate.main(url)
+
+        #TODO make sure user isn't queuing in dm for some stupid reason
+        for song in audio:
+           await queue.add_song(server, song, ctx.author.id)
+
+        await ctx.message.add_reaction('👍')
 
         await util.join_vc(ctx)
 
+        if await queue.is_server_playing(ctx.guild.id):
+            return
 
-        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        await queue.update_server(server, True)
-            info = ydl.extract_info(url, download=True)
+        await queue.play(ctx)
-            filename = ydl.prepare_filename(info)
-
-        ctx.voice_client.play(discord.FFmpegPCMAudio(executable="ffmpeg", source=filename), after=self.test)
-
-
-    def test(self, error):
-        print("Hello")

cogs/music/queue.py

@@ -1,7 +1,17 @@
 import sqlite3
 
+import discord
+import asyncio
+
 db_path = "./data/music.db"
 
+FFMPEG_OPTS = {
+    'before_options':
+        '-reconnect 1 -reconnect_streamed 1 -reconnect_delay_max 5',
+
+    'options':
+        '-vn'
+}
 
 # Creates the tables if they don't exist
 def initialize_tables():
@@ -12,18 +22,24 @@ def initialize_tables():
     # Create servers table if it doesn't exist
     cursor.execute('''CREATE TABLE IF NOT EXISTS servers (
                         server_id TEXT PRIMARY KEY,
-                        is_playing INTEGER DEFAULT 0,
+                        is_playing INTEGER DEFAULT 0
-                    )''')
+                    );''')
+
+    # Set all to not playing
+    cursor.execute("UPDATE servers SET is_playing = 0;")
 
     # Create queue table if it doesn't exist
     cursor.execute('''CREATE TABLE IF NOT EXISTS queue (
-                        server_id TEXT,
+                        server_id TEXT NOT NULL,
                         song_link TEXT,
                         queued_by TEXT,
-                        index INTEGER,
+                        position INTEGER NOT NULL,
                         has_played INTEGER DEFAULT 0,
-                        PRIMARY KEY (server_id, order_num)
+                        PRIMARY KEY (position),
-                    )''')
+                        FOREIGN KEY (server_id) REFERENCES servers(server_id)
+                    );''')
+    # Clear all entries
+    cursor.execute("DELETE FROM queue;")
 
     # Commit the changes and close the connection
     conn.commit()
@@ -31,26 +47,17 @@ def initialize_tables():
 
 
 # Queue a song in the db
-def add_song(server_id, song_link, queued_by):
+async def add_song(server_id, song_link, queued_by):
     # Connect to db
     conn = sqlite3.connect(db_path)
     cursor = conn.cursor()
 
-    add_server(server_id, cursor, conn)
+    await add_server(server_id, cursor, conn)
 
-    # Grab current index
+    max_order_num = await get_max(server_id, cursor) + 1
-    cursor.execute(f"""
-                   SELECT MAX(index)
-                   FROM queue
-                   WHERE server_id = ?
-                   """, (server_id,))
-    result = cursor.fetchone()
-
-    # Highnest number or 0
-    max_order_num = result[0] + 1 if result[0] is not None else 0
 
     cursor.execute("""
-                   INSERT INTO queue (server_id, song_link, queued_by, index)
+                   INSERT INTO queue (server_id, song_link, queued_by, position)
                    VALUES (?, ?, ?, ?)
                    """, (server_id, song_link, queued_by, max_order_num))
 
@@ -59,7 +66,7 @@ def add_song(server_id, song_link, queued_by):
 
 
 # Add server to db if first time queuing
-def add_server(server_id, cursor, conn):
+async def add_server(server_id, cursor, conn):
     # Check if the server exists
     cursor.execute('''SELECT COUNT(*)
                    FROM servers
@@ -76,46 +83,73 @@ def add_server(server_id, cursor, conn):
 
 
 # set song as played and update indexes
-def mark_song_as_finished(server_id, order_num):
+async def mark_song_as_finished(server_id, order_num):
     # Connect to the database
     conn = sqlite3.connect(db_path)
     cursor = conn.cursor()
 
     # Update the song as finished
     cursor.execute('''DELETE FROM queue
-                   WHERE server_id = ? AND order_num = ?''',
+                   WHERE server_id = ? AND position = ?''',
                    (server_id, order_num))
-    #cursor.execute('''UPDATE queue
-    #               SET is_finished = 1
-    #               WHERE server_id = ? AND index = ?''',
-    #               (server_id, order_num))
-
-    # Get the order numbers of unplayed songs
-    cursor.execute('''SELECT index
-                   FROM queue
-                   WHERE server_id = ? AND is_finished = 0''', (server_id,))
-    unplayed_order_nums = [row[0] for row in cursor.fetchall()]
-
-    # Update the order numbers of unplayed songs
-    for new_order, old_order in enumerate(unplayed_order_nums, start=1):
-        cursor.execute('''UPDATE queue
-                       SET order_num = ?
-                       WHERE server_id = ? AND order_num = ?''',
-                       (new_order, server_id, old_order))
 
     # Close connection
     conn.commit()
     conn.close()
 
 
+# Grab max order from server
+async def get_max(server_id, cursor):
+    cursor.execute(f"""
+                   SELECT MAX(position)
+                   FROM queue
+                   WHERE server_id = ?
+                   """, (server_id,))
+    result = cursor.fetchone()
+
+    # Highnest number or 0
+    max_order_num = result[0] if result[0] is not None else -1
+
+    return max_order_num
+
+
+# Pop song from server
+async def pop(server_id):
+    # Connect to db
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    # JUST INCASE!
+    await add_server(server_id, cursor, conn)
+
+    max_order = await get_max(server_id, cursor)
+    if max_order == -1:
+        conn.commit()
+        conn.close()
+        return None
+
+    cursor.execute('''SELECT song_link
+                   FROM queue
+                   WHERE server_id = ? AND position = ?
+                   ''', (server_id, max_order))
+    result = cursor.fetchone()
+
+    conn.commit()
+    conn.close()
+
+    await mark_song_as_finished(server_id, max_order)
+
+    return result[0]
+
+
 # Sets the playing variable in a server to true or false
-def update_server(server_id, playing: bool):
+async def update_server(server_id, playing: bool):
     # Connect to database
     conn = sqlite3.connect(db_path)
     cursor = conn.cursor()
 
     # add server to db if not present
-    add_server(server_id, cursor, conn)
+    await add_server(server_id, cursor, conn)
 
     value = 1 if playing else 0
 
@@ -129,13 +163,14 @@ def update_server(server_id, playing: bool):
     conn.commit()
     conn.close()
 
-def is_server_playing(server_id):
+
+async def is_server_playing(server_id):
     # Connect to db
     conn = sqlite3.connect(db_path)
     cursor = conn.cursor()
 
     # add server to db if not present
-    add_server(server_id, cursor, conn)
+    await add_server(server_id, cursor, conn)
 
     cursor.execute("""SELECT is_playing
                    FROM servers
@@ -143,8 +178,55 @@ def is_server_playing(server_id):
                    (server_id,))
 
     result = cursor.fetchone()
+    print(result)
 
     conn.commit()
     conn.close()
 
-    return result
+    return True if result[0] == 1 else False
+
+
+# Delete all songs from a server
+async def clear(server_id):
+    # Connect to db
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    await add_server(server_id, cursor, conn)
+    await update_server(server_id, False)
+
+    # Delete all songs from the server
+    cursor.execute('''DELETE FROM queue WHERE server_id = ?''', (server_id,))
+
+    conn.commit()
+    conn.close()
+
+
+# Play and loop songs in server
+async def play(ctx):
+    server_id = ctx.guild.id
+
+    # Wait until song is stopped playing 
+    #while ctx.voice_client.is_playing():
+        #await asyncio.sleep(1)
+
+    # check next song
+    url = await pop(server_id)
+
+    # if no other song update server and return
+    if url is None:
+        await update_server(server_id, False)
+        return
+
+    # else play next song and call play again
+    ctx.voice_client.play(
+            AstroPlayer(ctx, url, FFMPEG_OPTS))
+
+class AstroPlayer(discord.FFmpegPCMAudio):
+    def __init__(self, ctx, source, options) -> None:
+        self.ctx = ctx
+        super().__init__(source, **options)
+
+    def _kill_process(self):
+        super()._kill_process
+        asyncio.run(play(self.ctx))

cogs/music/translate.py

@@ -1,8 +1,16 @@
 # Handles translating urls and search terms
 
+import yt_dlp as ytdlp
+
+ydl_opts = {
+        'format':           'bestaudio/best',
+        'quiet':            True,
+        'default_search':   'ytsearch',
+}
+
 def main(url):
 
-    url = url.lower()
+    #url = url.lower()
 
     # Check if link or search
     if url.startswith("https://") is False:
@@ -16,7 +24,8 @@ def main(url):
             return spotify_playlist(url)
 
     soundcloud_song = 'soundcloud' in url and 'sets' not in url
-    soundcloud_playlist = 'soundcloud' in url and 'sets' in url
+    # Not implemented yet
+    #soundcloud_playlist = 'soundcloud' in url and 'sets' in url
 
     youtube_song = 'watch?v=' in url or 'youtu.be/' in url
     youtube_playlist = 'playlist?list=' in url
@@ -27,24 +36,30 @@ def main(url):
     if youtube_playlist:
         return playlist_download(url)
 
-    return False
+    return []
 
 
 def search_song(search):
-    return None
+    with ytdlp.YoutubeDL(ydl_opts) as ydl:
+        info = ydl.extract_info(f"ytsearch1:{search}", download=False)
+        audio_url = info['entries'][0]['url'] # Get audio stream URL
+    return [audio_url]
 
 
 def spotify_song(url):
-    return None
+    return []
 
 
 def spotify_playlist(url):
-    return None
+    return []
 
 
 def song_download(url):
-    return None
+    with ytdlp.YoutubeDL(ydl_opts) as ydl:
+        info = ydl.extract_info(url, download=False)
+        audio_url = info['url'] # Get audio stream URL
+    return [audio_url]
 
 
 def playlist_download(url):
-    return None
+    return []

cogs/music/util.py

@@ -1,5 +1,3 @@
-import discord
-from discord.ext import commands
 from discord.ext.commands.context import Context
 from discord.ext.commands.converter import CommandError
 
@@ -45,3 +43,8 @@ async def leave_vc(ctx: Context):
 
     # Disconnect
     await ctx.voice_client.disconnect(force=False)
+
+
+# Check if command was entered in a server
+async def in_server(ctx: Context):
+    return ctx.guild != None 

main.py

@@ -6,4 +6,4 @@ import help
 client = Astro(command_prefix=config.get_prefix(), intents=discord.Intents.all())
 client.help_command = help.AstroHelp()
 
-client.run(config.get_login("dev"))
+client.run(config.get_login("live"))