LibreTunes/src/api/search.rs

use crate::prelude::*;

cfg_if! {
if #[cfg(feature = "ssr")] {
    use diesel::sql_types::*;
    use diesel::pg::Pg;
    use diesel::expression::AsExpression;

    // Define pg_trgm operators
    // Functions do not use indices for queries, so we need to use operators
    diesel::infix_operator!(Similarity, " % ", backend: Pg);
    diesel::infix_operator!(Distance, " <-> ", Float, backend: Pg);

    // Create functions to make use of the operators in queries
    fn trgm_similar<T: AsExpression<Text>, U: AsExpression<Text>>(left: T, right: U)
        -> Similarity<T::Expression, U::Expression> {
        Similarity::new(left.as_expression(), right.as_expression())
    }

    fn trgm_distance<T: AsExpression<Text>, U: AsExpression<Text>>(left: T, right: U)
        -> Distance<T::Expression, U::Expression> {
        Distance::new(left.as_expression(), right.as_expression())
    }
}
}

/// A simple type for search results
/// A vector of tuples containing the item and its score
pub type SearchResults<T> = Vec<(T, f32)>;

/// Turn `SearchResults` into just the results, discarding scores
pub fn remove_search_score<T>(results: SearchResults<T>) -> Vec<T> {
    results.into_iter().map(|(item, _score)| item).collect()
}

/// Search for albums by title
///
/// # Arguments
/// `query` - The search query. This will be used to perform a fuzzy search on the album titles
/// `limit` - The maximum number of results to return
///
/// # Returns
/// A Result containing a vector of albums if the search was successful, or an error if the search failed
#[api_fn(endpoint = "search_albums")]
pub async fn search_albums(
    query: String,
    limit: i64,
    db_conn: &mut PgPooledConn,
) -> BackendResult<SearchResults<frontend::Album>> {
    let album_ids = albums::table
        .filter(trgm_similar(albums::title, query.clone()))
        .order_by(trgm_distance(albums::title, query.clone()).desc())
        .limit(limit)
        .select(albums::id)
        .into_boxed();

    let mut albums_map: HashMap<i32, (frontend::Album, f32)> = HashMap::new();

    let album_artists: Vec<(backend::Album, Option<backend::Artist>, f32)> = albums::table
        .filter(albums::id.eq_any(album_ids))
        .left_join(
            album_artists::table
                .inner_join(artists::table)
                .on(albums::id.eq(album_artists::album_id)),
        )
        .select((
            albums::all_columns,
            artists::all_columns.nullable(),
            trgm_distance(albums::title, query.clone()),
        ))
        .load(db_conn)
        .context("Error loading album artists from database")?;

    for (album, artist, score) in album_artists {
        if let Some((stored_album, _score)) = albums_map.get_mut(&album.id) {
            if let Some(artist) = artist {
                stored_album.artists.push(artist);
            }
        } else {
            let albumdata = frontend::Album {
                id: album.id,
                title: album.title,
                artists: artist.map(|artist| vec![artist]).unwrap_or(vec![]),
                release_date: album.release_date,
                image_path: LocalPath::to_web_path_or_placeholder(album.image_path),
            };

            albums_map.insert(album.id, (albumdata, score));
        }
    }

    let mut albums: Vec<(frontend::Album, f32)> = albums_map.into_values().collect();
    albums.sort_by(|(_a, a_score), (_b, b_score)| b_score.total_cmp(a_score));

    Ok(albums)
}

/// Search for artists by name
///
/// # Arguments
/// `query` - The search query. This will be used to perform a fuzzy search on the artist names
/// `limit` - The maximum number of results to return
///
/// # Returns
/// A Result containing a vector of artists if the search was successful, or an error if the search failed
#[api_fn(endpoint = "search_artists")]
pub async fn search_artists(
    query: String,
    limit: i64,
    db_conn: &mut PgPooledConn,
) -> BackendResult<SearchResults<frontend::Artist>> {
    let artist_list = artists::table
        .filter(trgm_similar(artists::name, query.clone()))
        .order_by(trgm_distance(artists::name, query.clone()).desc())
        .limit(limit)
        .select((artists::all_columns, trgm_distance(artists::name, query)))
        .load::<(backend::Artist, f32)>(db_conn)
        .context("Error loading artists from database")?;

    let artist_data = artist_list
        .into_iter()
        .map(|(artist, score)| {
            (
                frontend::Artist {
                    id: artist.id,
                    name: artist.name,
                    image_path: LocalPath::to_web_path_or_placeholder(artist.image_path),
                },
                score,
            )
        })
        .collect();

    Ok(artist_data)
}

/// Search for songs by title
///
/// # Arguments
/// `query` - The search query. This will be used to perform a fuzzy search on the song titles
/// `limit` - The maximum number of results to return
///
/// # Returns
/// A Result containing a vector of songs if the search was successful, or an error if the search failed
#[api_fn(endpoint = "search_songs")]
pub async fn search_songs(
    query: String,
    limit: i64,
    db_conn: &mut PgPooledConn,
) -> BackendResult<SearchResults<i32>> {
    songs::table
        .filter(trgm_similar(songs::title, query.clone()))
        .order_by(trgm_distance(songs::title, query.clone()).desc())
        .limit(limit)
        .select((songs::id, trgm_distance(songs::title, query.clone())))
        .load(db_conn)
        .context("Failed to search songs")
}

/// Search for songs, albums, and artists by title or name
///
/// # Arguments
/// `query` - The search query. This will be used to perform a fuzzy search on the
/// song titles, album titles, and artist names
/// `limit` - The maximum number of results to return for each type
///
/// # Returns
/// A Result containing a tuple of vectors of albums, artists, and songs if the search was successful,
#[api_fn(endpoint = "search")]
pub async fn search(
    query: String,
    limit: i64,
) -> BackendResult<(
    SearchResults<frontend::Album>,
    SearchResults<frontend::Artist>,
    SearchResults<i32>,
)> {
    let albums = search_albums(query.clone(), limit);
    let artists = search_artists(query.clone(), limit);
    let songs = search_songs(query, limit);

    use tokio::join;

    let (albums, artists, songs) = join!(albums, artists, songs);

    let albums = albums.context("Error searching for albums")?;
    let artists = artists.context("Error searching for artists")?;
    let songs = songs.context("Error searching for songs")?;

    Ok((albums, artists, songs))
}