Add automatically adding youtube songs to store

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

Cargo.lock

@@ -193,7 +193,7 @@ dependencies = [
 
 [[package]]
 name = "bot-database"
-version = "0.3.0"
+version = "0.4.0"
 dependencies = [
  "chrono",
  "diesel",
@@ -2346,6 +2346,7 @@ dependencies = [
  "sysinfo",
  "thiserror",
  "tokio",
+ "trigram",
 ]
 
 [[package]]
@@ -2474,6 +2475,16 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "trigram"
+version = "0.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "324350103581dfc6869f2c9399c019f394ac49a7ab16af23206add07851c5347"
+dependencies = [
+ "lazy_static",
+ "regex",
+]
+
 [[package]]
 name = "try-lock"
 version = "0.2.3"

Cargo.toml

@@ -35,4 +35,5 @@ sysinfo = "0.16.5"
 reqwest = "0.11.2"
 chrono-tz = "0.5.3"
 sauce-api = "0.7.1"
-rustc_version_runtime = "0.2.0"
+rustc_version_runtime = "0.2.0"
+trigram = "0.4.4"

src/commands/music/mod.rs

@@ -35,10 +35,12 @@ use skip::SKIP_COMMAND;
 
 use crate::messages::music::now_playing::update_now_playing_msg;
 use crate::providers::music::queue::{MusicQueue, Song};
-use crate::providers::music::youtube_dl;
+use crate::providers::music::{add_youtube_song_to_database, youtube_dl};
 use crate::providers::settings::{get_setting, Setting};
 use crate::utils::context_data::{DatabaseContainer, Store};
 use crate::utils::error::{BotError, BotResult};
+use aspotify::Track;
+use bot_database::Database;
 
 mod clear_queue;
 mod current;
@@ -327,6 +329,7 @@ async fn get_songs_for_query(ctx: &Context, msg: &Message, query: &str) -> BotRe
             log::debug!("Query is youtube video");
             let mut song: Song = youtube_dl::get_video_information(&query).await?.into();
             added_one_msg(&ctx, msg, &mut song).await?;
+            add_youtube_song_to_database(&store, &database, &mut song).await?;
             songs.push(song);
         } else {
             log::debug!("Query is playlist with {} songs", songs.len());
@@ -335,17 +338,38 @@ async fn get_songs_for_query(ctx: &Context, msg: &Message, query: &str) -> BotRe
     } 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?;
+        let tracks = store.spotify_api.get_songs_in_playlist(&query).await?;
+
+        for track in tracks {
+            songs.push(
+                get_youtube_song_for_track(&database, track.clone())
+                    .await?
+                    .unwrap_or(track.into()),
+            )
+        }
+
         added_multiple_msg(&ctx, msg, &mut songs).await?;
     } 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?;
+        let tracks = store.spotify_api.get_songs_in_album(&query).await?;
+
+        for track in tracks {
+            songs.push(
+                get_youtube_song_for_track(&database, track.clone())
+                    .await?
+                    .unwrap_or(track.into()),
+            )
+        }
+
         added_multiple_msg(&ctx, msg, &mut songs).await?;
     } 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 track = store.spotify_api.get_song_name(&query).await?;
+        let mut song = get_youtube_song_for_track(&database, track.clone())
+            .await?
+            .unwrap_or(track.into());
         added_one_msg(ctx, msg, &mut song).await?;
         songs.push(song);
     } else {
@@ -409,3 +433,16 @@ pub async fn is_dj(ctx: &Context, guild: GuildId, user: &User) -> BotResult<bool
         Ok(true)
     }
 }
+
+/// Searches for a matching youtube song for the given track in the local database
+async fn get_youtube_song_for_track(database: &Database, track: Track) -> BotResult<Option<Song>> {
+    log::debug!("Trying to find track in database.");
+    if let Some(id) = track.id {
+        let entry = database.get_song(&id).await?;
+        log::trace!("Found entry is {:?}", entry);
+        Ok(entry.map(Song::from))
+    } else {
+        log::debug!("Track has no ID");
+        Ok(None)
+    }
+}

src/main.rs

@@ -4,12 +4,12 @@ use crate::utils::logging::init_logger;
 #[macro_use]
 extern crate bot_serenityutils;
 
-pub mod client;
+mod client;
 mod commands;
-pub mod handler;
+mod handler;
 mod messages;
 mod providers;
-pub mod utils;
+mod utils;
 
 pub static VERSION: &str = env!("CARGO_PKG_VERSION");
 

src/providers/music/mod.rs

@@ -1,5 +1,115 @@
+use crate::providers::music::queue::Song;
+use crate::utils::context_data::StoreData;
+use crate::utils::error::BotResult;
+use aspotify::{ArtistSimplified, Track};
+use bot_database::Database;
+use regex::Regex;
+use responses::VideoInformation;
+use youtube_dl::search_video_information;
+
 pub(crate) mod lyrics;
 pub(crate) mod queue;
 pub(crate) mod responses;
 pub(crate) mod spotify;
 pub(crate) mod youtube_dl;
+
+/// Searches for a youtube video for the specified song
+pub(crate) async fn song_to_youtube_video(song: &Song) -> BotResult<Option<VideoInformation>> {
+    let artist = song.author().clone();
+    let title = song.title().clone();
+    let match_query = format!("{} - {}", artist, title);
+
+    let queries = vec![
+        format!("{} - {} lyrics", artist, title),
+        format!("{} - {} audio only", artist, title),
+        format!("{} by {}", title, artist),
+        match_query.clone(),
+    ];
+
+    let mut last_result = None;
+    for query in queries {
+        let result = search_video_information(query).await?;
+
+        if let Some(video) = result {
+            if trigram::similarity(&video.title, &match_query) >= 0.4
+                || (trigram::similarity(&video.title, &title) >= 0.3
+                    && trigram::similarity(&video.uploader, &artist) >= 0.3)
+            {
+                return Ok(Some(video));
+            }
+            last_result = Some(video);
+        }
+    }
+
+    Ok(last_result)
+}
+
+/// Adds a youtube song to the database of songs
+pub async fn add_youtube_song_to_database(
+    store: &StoreData,
+    database: &Database,
+    song: &mut Song,
+) -> BotResult<()> {
+    match search_for_song_variations(store, song).await {
+        Ok(Some(track)) => {
+            log::debug!("Song found on spotify. Inserting metadata");
+            let artists = artists_to_string(track.artists);
+            let url = song.url().await.unwrap();
+
+            if let Some(id) = track.id {
+                database
+                    .add_song(&id, &artists, &track.name, &track.album.name, &url)
+                    .await?;
+            }
+        }
+        Err(e) => log::error!("Failed to search for song on spotify {:?}", e),
+        _ => {}
+    }
+
+    Ok(())
+}
+
+/// Searches for multiple queries on spotify
+async fn search_for_song_variations(
+    store: &StoreData,
+    song: &mut Song,
+) -> BotResult<Option<Track>> {
+    static COMMON_AFFIXES: &str =
+        r"feat\.(\s\w+)|official(\svideo)?|remastered|revisited|(with\s)?lyrics";
+    lazy_static::lazy_static! {
+        static ref COMMON_ADDITIONS: Regex = Regex::new(format!(r"(?i)\[|\]|\(?[^\w\s]*\s?({})[^\w\s]*\s?\)?", COMMON_AFFIXES).as_str()).unwrap();
+    }
+    let mut query = song
+        .title()
+        .replace(|c| c != ' ' && !char::is_alphanumeric(c), "");
+    query = COMMON_ADDITIONS.replace_all(&query, " ").to_string();
+
+    log::debug!("Searching for youtube song");
+    if let Some(track) = store.spotify_api.search_for_song(&query).await? {
+        let similarity = trigram::similarity(
+            &format!(
+                "{} {}",
+                artists_to_string(track.artists.clone()),
+                track.name
+            ),
+            &query,
+        );
+        if similarity > 0.3 {
+            log::debug!("Result is similar enough ({}). Returning track", similarity);
+            return Ok(Some(track));
+        }
+        log::debug!("Result is not similar enough");
+    }
+    log::debug!("No result found");
+
+    Ok(None)
+}
+
+/// Creates a string from a vector of artists
+pub fn artists_to_string(artists: Vec<ArtistSimplified>) -> String {
+    artists
+        .into_iter()
+        .map(|a| a.name)
+        .collect::<Vec<String>>()
+        .join("&")
+}

src/providers/music/queue.rs

@@ -6,7 +6,8 @@ use songbird::tracks::TrackHandle;
 use bot_coreutils::shuffle::Shuffle;
 
 use crate::providers::music::responses::{PlaylistEntry, VideoInformation};
-use crate::providers::music::youtube_dl;
+use crate::providers::music::song_to_youtube_video;
+use bot_database::models::YoutubeSong;
 use bot_serenityutils::core::MessageHandle;
 use std::sync::Arc;
 use tokio::sync::RwLock;
@@ -125,11 +126,7 @@ impl Song {
             Some(url)
         } else {
             log::debug!("Lazy fetching video for title");
-            let information =
-                youtube_dl::search_video_information(format!("{} - {}", self.author, self.title))
-                    .await
-                    .ok()
-                    .and_then(|i| i)?;
+            let information = song_to_youtube_video(&self).await.ok()??;
             self.url = Some(information.webpage_url.clone());
             self.thumbnail = information.thumbnail;
             self.author = information.uploader;
@@ -208,3 +205,14 @@ impl From<TrackSimplified> for Song {
         }
     }
 }
+
+impl From<YoutubeSong> for Song {
+    fn from(song: YoutubeSong) -> Self {
+        Self {
+            title: song.title,
+            author: song.artist,
+            url: Some(song.url),
+            thumbnail: None,
+        }
+    }
+}

src/providers/music/spotify.rs

@@ -1,6 +1,5 @@
-use aspotify::{Client, ClientCredentials, PlaylistItem, PlaylistItemType};
+use aspotify::{Client, ClientCredentials, ItemType, PlaylistItemType, Track};
 
-use crate::providers::music::queue::Song;
 use crate::utils::error::{BotError, BotResult};
 
 pub struct SpotifyApi {
@@ -21,8 +20,26 @@ impl SpotifyApi {
         Self { client }
     }
 
+    /// Searches for a song on spotify
+    pub async fn search_for_song(&self, query: &str) -> BotResult<Option<Track>> {
+        log::debug!("Searching for song '{}' on spotify", query);
+        let types = vec![ItemType::Track];
+        let result = self
+            .client
+            .search()
+            .search(query, types, false, 1, 0, None)
+            .await?;
+        log::trace!("Result is {:?}", result);
+        let tracks = result
+            .data
+            .tracks
+            .ok_or(BotError::from("Failed to get search spotify results"))?;
+
+        Ok(tracks.items.into_iter().next())
+    }
+
     /// Returns the songs for a playlist
-    pub async fn get_songs_in_playlist(&self, url: &str) -> BotResult<Vec<Song>> {
+    pub async fn get_songs_in_playlist(&self, url: &str) -> BotResult<Vec<Track>> {
         log::debug!("Fetching spotify songs from playlist '{}'", url);
         let id = self.get_id_for_url(url)?;
         let mut playlist_tracks = Vec::new();
@@ -42,17 +59,9 @@ impl SpotifyApi {
             url
         );
 
-        let songs = playlist_tracks
-            .into_iter()
-            .filter_map(|item| item.item)
-            .filter_map(|t| match t {
-                PlaylistItemType::Track(t) => Some(Song::from(t)),
-                PlaylistItemType::Episode(_) => None,
-            })
-            .collect();
-        log::trace!("Songs are {:?}", songs);
+        log::trace!("Songs are {:?}", playlist_tracks);
 
-        Ok(songs)
+        Ok(playlist_tracks)
     }
 
     /// Returns the tracks of a playlist with pagination
@@ -61,43 +70,66 @@ impl SpotifyApi {
         id: &str,
         limit: usize,
         offset: usize,
-    ) -> BotResult<Vec<PlaylistItem>> {
+    ) -> BotResult<Vec<Track>> {
         log::trace!(
             "Fetching songs from spotify playlist: limit {}, offset {}",
             limit,
             offset
         );
-        let tracks = self
+        let page = self
             .client
             .playlists()
             .get_playlists_items(id, limit, offset, None)
             .await?
             .data;
+
+        let tracks: Vec<Track> = page
+            .items
+            .into_iter()
+            .filter_map(|item| item.item)
+            .filter_map(|t| match t {
+                PlaylistItemType::Track(t) => Some(t),
+                PlaylistItemType::Episode(_) => None,
+            })
+            .collect();
         log::trace!("Tracks are {:?}", tracks);
 
-        Ok(tracks.items)
+        Ok(tracks)
     }
 
     /// Returns all songs for a given album
-    pub async fn get_songs_in_album(&self, url: &str) -> BotResult<Vec<Song>> {
+    pub async fn get_songs_in_album(&self, url: &str) -> BotResult<Vec<Track>> {
         log::debug!("Fetching songs for spotify album '{}'", url);
         let id = self.get_id_for_url(url)?;
         let album = self.client.albums().get_album(&*id, None).await?.data;
         log::trace!("Album is {:?}", album);
-        let song_names: Vec<Song> = album.tracks.items.into_iter().map(Song::from).collect();
-        log::debug!("{} songs found in album '{}'", song_names.len(), url);
 
-        Ok(song_names)
+        let simple_tracks: Vec<String> = album
+            .tracks
+            .items
+            .into_iter()
+            .filter_map(|t| t.id)
+            .collect();
+        let tracks = self
+            .client
+            .tracks()
+            .get_tracks(simple_tracks, None)
+            .await?
+            .data;
+
+        log::trace!("Tracks are {:?}", tracks);
+
+        Ok(tracks)
     }
 
     /// Returns song entity for a given spotify url
-    pub async fn get_song_name(&self, url: &str) -> BotResult<Song> {
+    pub async fn get_song_name(&self, url: &str) -> BotResult<Track> {
         log::debug!("Getting song for {}", url);
         let id = self.get_id_for_url(url)?;
         let track = self.client.tracks().get_track(&*id, None).await?.data;
         log::trace!("Track info is {:?}", track);
 
-        Ok(track.into())
+        Ok(track)
     }
 
     /// Returns the id for a given spotify URL