Add command to manage playlists

Signed-off-by: trivernis <trivernis@protonmail.com>

database/migrations/2021-04-09-095329_create_guild_playlists/down.sql

@@ -0,0 +1,2 @@
+-- This file should undo anything in `up.sql`
+DROP TABLE guild_playlists;

database/migrations/2021-04-09-095329_create_guild_playlists/up.sql

@@ -0,0 +1,7 @@
+-- Your SQL goes here
+CREATE TABLE guild_playlists (
+    guild_id BIGINT NOT NULL,
+    name VARCHAR(255) NOT NULL,
+    url VARCHAR(1024) NOT NULL,
+    PRIMARY KEY (guild_id, name)
+)

database/src/database.rs

@@ -78,4 +78,53 @@ impl Database {
 
         Ok(())
     }
+
+    /// Returns a list of all guild playlists
+    pub fn get_guild_playlists(&self, guild_id: u64) -> DatabaseResult<Vec<GuildPlaylist>> {
+        use guild_playlists::dsl;
+        log::debug!("Retrieving guild playlists for guild {}", guild_id);
+        let connection = self.pool.get()?;
+        let playlists: Vec<GuildPlaylist> = dsl::guild_playlists
+            .filter(dsl::guild_id.eq(guild_id as i64))
+            .load::<GuildPlaylist>(&connection)?;
+
+        Ok(playlists)
+    }
+
+    /// Returns a guild playlist by name
+    pub fn get_guild_playlist(
+        &self,
+        guild_id: u64,
+        name: &str,
+    ) -> DatabaseResult<Option<GuildPlaylist>> {
+        use guild_playlists::dsl;
+        log::debug!("Retriving guild playlist '{}' for guild {}", name, guild_id);
+        let connection = self.pool.get()?;
+
+        let playlists: Vec<GuildPlaylist> = dsl::guild_playlists
+            .filter(dsl::guild_id.eq(guild_id as i64))
+            .filter(dsl::name.eq(name))
+            .load::<GuildPlaylist>(&connection)?;
+
+        Ok(playlists.into_iter().next())
+    }
+
+    /// Adds a new playlist to the database overwriting the old one
+    pub fn add_guild_playlist(&self, guild_id: u64, name: &str, url: &str) -> DatabaseResult<()> {
+        use guild_playlists::dsl;
+        log::debug!("Inserting guild playlist '{}' for guild {}", name, guild_id);
+        let connection = self.pool.get()?;
+        insert_into(dsl::guild_playlists)
+            .values(GuildPlaylistInsert {
+                guild_id: guild_id as i64,
+                name: name.to_string(),
+                url: url.to_string(),
+            })
+            .on_conflict((dsl::guild_id, dsl::name))
+            .do_update()
+            .set(dsl::url.eq(url.to_string()))
+            .execute(&connection)?;
+
+        Ok(())
+    }
 }

database/src/models.rs

@@ -14,3 +14,18 @@ pub struct GuildSettingInsert {
     pub key: String,
     pub value: String,
 }
+
+#[derive(Queryable, Debug)]
+pub struct GuildPlaylist {
+    pub guild_id: i64,
+    pub name: String,
+    pub url: String,
+}
+
+#[derive(Insertable, Debug)]
+#[table_name = "guild_playlists"]
+pub struct GuildPlaylistInsert {
+    pub guild_id: i64,
+    pub name: String,
+    pub url: String,
+}

database/src/schema.rs

@@ -1,3 +1,11 @@
+table! {
+    guild_playlists (guild_id, name) {
+        guild_id -> Int8,
+        name -> Varchar,
+        url -> Varchar,
+    }
+}
+
 table! {
     guild_settings (guild_id, key) {
         guild_id -> Int8,
@@ -5,3 +13,8 @@ table! {
         value -> Nullable<Varchar>,
     }
 }
+
+allow_tables_to_appear_in_same_query!(
+    guild_playlists,
+    guild_settings,
+);

src/commands/music/mod.rs

@@ -19,7 +19,9 @@ use leave::LEAVE_COMMAND;
 use pause::PAUSE_COMMAND;
 use play::PLAY_COMMAND;
 use play_next::PLAY_NEXT_COMMAND;
+use playlists::PLAYLISTS_COMMAND;
 use queue::QUEUE_COMMAND;
+use save_playlist::SAVE_PLAYLIST_COMMAND;
 use shuffle::SHUFFLE_COMMAND;
 use skip::SKIP_COMMAND;
 
@@ -27,7 +29,7 @@ use crate::providers::music::queue::{MusicQueue, Song};
 use crate::providers::music::{
     get_video_information, get_videos_for_playlist, search_video_information,
 };
-use crate::utils::context_data::Store;
+use crate::utils::context_data::{DatabaseContainer, Store};
 use crate::utils::error::{BotError, BotResult};
 use regex::Regex;
 use std::sync::atomic::{AtomicIsize, AtomicUsize, Ordering};
@@ -40,13 +42,26 @@ mod leave;
 mod pause;
 mod play;
 mod play_next;
+mod playlists;
 mod queue;
+mod save_playlist;
 mod shuffle;
 mod skip;
 
 #[group]
 #[commands(
-    join, leave, play, queue, skip, shuffle, current, play_next, clear, pause
+    join,
+    leave,
+    play,
+    queue,
+    skip,
+    shuffle,
+    current,
+    play_next,
+    clear,
+    pause,
+    save_playlist,
+    playlists
 )]
 #[prefixes("m", "music")]
 pub struct Music;
@@ -240,8 +255,11 @@ async fn play_next_in_queue(
 
 /// Returns the list of songs for a given url
 async fn get_songs_for_query(ctx: &Context, msg: &Message, query: &str) -> BotResult<Vec<Song>> {
+    let guild_id = msg.guild_id.unwrap();
+    let mut query = query.to_string();
     lazy_static::lazy_static! {
         // expressions to determine the type of url
+        static ref PLAYLIST_NAME_REGEX: Regex = Regex::new(r"^pl:(\S+)$").unwrap();
         static ref YOUTUBE_URL_REGEX: Regex = Regex::new(r"^(https?(://))?(www\.)?(youtube\.com/watch\?.*v=.*)|(/youtu.be/.*)|(youtube\.com/playlist\?.*list=.*)$").unwrap();
         static ref SPOTIFY_PLAYLIST_REGEX: Regex = Regex::new(r"^(https?(://))?(www\.|open\.)?spotify\.com/playlist/.*").unwrap();
         static ref SPOTIFY_ALBUM_REGEX: Regex = Regex::new(r"^(https?(://))?(www\.|open\.)?spotify\.com/album/.*").unwrap();
@@ -250,12 +268,25 @@ async fn get_songs_for_query(ctx: &Context, msg: &Message, query: &str) -> BotRe
     let mut songs = Vec::new();
     let data = ctx.data.read().await;
     let store = data.get::<Store>().unwrap();
+    let database = data.get::<DatabaseContainer>().unwrap();
 
     log::debug!("Querying play input {}", query);
-    if YOUTUBE_URL_REGEX.is_match(query) {
+    if let Some(captures) = PLAYLIST_NAME_REGEX.captures(&query) {
+        log::debug!("Query is a saved playlist");
+        let pl_name: &str = captures.get(1).unwrap().as_str();
+        log::trace!("Playlist name is {}", pl_name);
+        let playlist_opt = database.get_guild_playlist(guild_id.0, pl_name)?;
+        log::trace!("Playlist is {:?}", playlist_opt);
+
+        if let Some(playlist) = playlist_opt {
+            log::debug!("Assigning url for saved playlist to query");
+            query = playlist.url;
+        }
+    }
+    if YOUTUBE_URL_REGEX.is_match(&query) {
         log::debug!("Query is youtube video or playlist");
         // try fetching the url as a playlist
-        songs = get_videos_for_playlist(query)
+        songs = get_videos_for_playlist(&query)
             .await?
             .into_iter()
             .map(Song::from)
@@ -264,32 +295,32 @@ async fn get_songs_for_query(ctx: &Context, msg: &Message, query: &str) -> BotRe
         // if no songs were found fetch the song as a video
         if songs.len() == 0 {
             log::debug!("Query is youtube video");
-            let mut song: Song = get_video_information(query).await?.into();
+            let mut song: Song = get_video_information(&query).await?.into();
             added_one_msg(&ctx, msg, &mut song).await?;
             songs.push(song);
         } else {
             log::debug!("Query is playlist with {} songs", songs.len());
             added_multiple_msg(&ctx, msg, &mut songs).await?;
         }
-    } else if SPOTIFY_PLAYLIST_REGEX.is_match(query) {
+    } else if SPOTIFY_PLAYLIST_REGEX.is_match(&query) {
         // search for all songs in the playlist and search for them on youtube
         log::debug!("Query is spotify playlist");
-        songs = store.spotify_api.get_songs_in_playlist(query).await?;
+        songs = store.spotify_api.get_songs_in_playlist(&query).await?;
         added_multiple_msg(&ctx, msg, &mut songs).await?;
-    } else if SPOTIFY_ALBUM_REGEX.is_match(query) {
+    } else if SPOTIFY_ALBUM_REGEX.is_match(&query) {
         // fetch all songs in the album and search for them on youtube
         log::debug!("Query is spotify album");
-        songs = store.spotify_api.get_songs_in_album(query).await?;
+        songs = store.spotify_api.get_songs_in_album(&query).await?;
         added_multiple_msg(&ctx, msg, &mut songs).await?;
-    } else if SPOTIFY_SONG_REGEX.is_match(query) {
+    } else if SPOTIFY_SONG_REGEX.is_match(&query) {
         // fetch the song name and search it on youtube
         log::debug!("Query is a spotify song");
-        let mut song = store.spotify_api.get_song_name(query).await?;
+        let mut song = store.spotify_api.get_song_name(&query).await?;
         added_one_msg(ctx, msg, &mut song).await?;
         songs.push(song);
     } else {
         log::debug!("Query is a youtube search");
-        let mut song: Song = search_video_information(query.to_string())
+        let mut song: Song = search_video_information(query.clone())
             .await?
             .ok_or(BotError::Msg(format!("Noting found for {}", query)))?
             .into();

src/commands/music/playlists.rs

@@ -0,0 +1,27 @@
+use crate::utils::context_data::get_database_from_context;
+use serenity::client::Context;
+use serenity::framework::standard::macros::command;
+use serenity::framework::standard::CommandResult;
+use serenity::model::channel::Message;
+
+#[command]
+#[only_in(guilds)]
+#[description("Displays a list of all saved playlists")]
+#[usage("")]
+async fn playlists(ctx: &Context, msg: &Message) -> CommandResult {
+    let guild = msg.guild(&ctx.cache).await.unwrap();
+    log::debug!("Displaying playlists for guild {}", guild.id);
+    let database = get_database_from_context(ctx).await;
+
+    let playlists = database.get_guild_playlists(guild.id.0)?;
+    msg.channel_id
+        .send_message(ctx, |m| {
+            m.embed(|e| {
+                e.title("Saved Playlists")
+                    .fields(playlists.into_iter().map(|p| (p.name, p.url, true)))
+            })
+        })
+        .await?;
+
+    Ok(())
+}

src/commands/music/save_playlist.rs

@@ -0,0 +1,34 @@
+use crate::utils::context_data::get_database_from_context;
+use serenity::client::Context;
+use serenity::framework::standard::macros::command;
+use serenity::framework::standard::{Args, CommandResult};
+use serenity::model::channel::Message;
+
+#[command]
+#[only_in(guilds)]
+#[description("Adds a playlist to the guilds saved playlists")]
+#[usage("<name> <url/query>")]
+#[example("anime https://www.youtube.com/playlist?list=PLqaM77H_o5hykROCe3uluvZEaPo6bZj-C")]
+#[min_args(2)]
+#[aliases("add-playlist", "save-playlist")]
+#[allowed_roles("DJ")]
+async fn save_playlist(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
+    let guild = msg.guild(&ctx.cache).await.unwrap();
+    let name: String = args.single().unwrap();
+    let url: &str = args.remains().unwrap();
+    log::debug!(
+        "Adding playlist '{}' with url '{}' to guild {}",
+        name,
+        url,
+        guild.id
+    );
+    let database = get_database_from_context(ctx).await;
+
+    database.add_guild_playlist(guild.id.0, &*name, url)?;
+
+    msg.channel_id
+        .say(ctx, format!("Playlist **{}** saved", name))
+        .await?;
+
+    Ok(())
+}