little-hesinde/calibre-db/src/search.rs

//! Provide search funcitonality for calibre.
//!
//! Because the calibre database can not be disturbed (it is treated as read-only)
//! it attaches a temporary database and inserts the relevant data into a
//! virtual table leveraging fts5 (https://www.sqlite.org/fts5.html). Full-text search is run on
//! that virtual table.

use r2d2::{Pool, PooledConnection};
use r2d2_sqlite::SqliteConnectionManager;
use rusqlite::named_params;
use tempfile::NamedTempFile;

use crate::data::{book::Book, error::DataStoreError};

/// A lot of joins but only run once at startup.
const SEARCH_INIT_QUERY: &str = "INSERT INTO search.fts(book_id, data) 
    SELECT b.id as book_id, 
      COALESCE(b.title, '') || ' ' || 
      COALESCE(a.name, '') || ' ' || 
      COALESCE(c.text, '') || ' ' ||
      COALESCE(GROUP_CONCAT(DISTINCT t.name), '') || ' ' ||
      COALESCE(GROUP_CONCAT(DISTINCT i.val), '') || ' ' ||
      COALESCE(GROUP_CONCAT(DISTINCT s.name), '') as data
    FROM main.books as b
      LEFT JOIN main.books_authors_link AS b2a ON b.id = b2a.book
      LEFT JOIN main.authors AS a ON b2a.author = a.id 
      LEFT JOIN main.comments AS c ON c.book = b.id
      LEFT JOIN main.books_tags_link AS b2t ON b.id = b2t.book
      LEFT JOIN main.tags AS t ON b2t.tag = t.id
      LEFT JOIN main.identifiers AS i ON i.book = b.id
      LEFT JOIN main.books_series_link AS b2s ON b.id = b2s.book
      LEFT JOIN main.series AS s ON b2s.series = s.id
    GROUP BY b.id";

/// Attach the fts temporary database to the read-only calibre database.
pub(crate) fn attach(pool: &Pool<SqliteConnectionManager>) -> Result<(), DataStoreError> {
    let conn = pool.get()?;
    let tmpfile = NamedTempFile::new()?;
    let (_, path) = tmpfile.keep()?;

    conn.execute(
        &format!("ATTACH DATABASE '{}' AS search", path.to_string_lossy()),
        [],
    )?;
    init(&conn)?;

    Ok(())
}

/// Initialise the fts virtual table.
fn init(conn: &PooledConnection<SqliteConnectionManager>) -> Result<(), DataStoreError> {
    conn.execute(
        "CREATE VIRTUAL TABLE search.fts USING fts5(book_id, data)",
        [],
    )?;
    conn.execute(SEARCH_INIT_QUERY, [])?;

    Ok(())
}

/// Run a full-text search with the parameter `query`.
pub(crate) fn search(
    query: &str,
    pool: &Pool<SqliteConnectionManager>,
) -> Result<Vec<Book>, DataStoreError> {
    let conn = pool.get()?;

    let mut stmt =
        conn.prepare("SELECT book_id FROM search.fts WHERE data MATCH (:query) ORDER BY rank")?;
    let params = named_params! { ":query": query };
    let books = stmt
        .query_map(params, |r| -> Result<u64, rusqlite::Error> { r.get(0) })?
        .filter_map(Result::ok)
        .filter_map(|id| Book::scalar_book(&conn, id).ok())
        .collect();

    Ok(books)
}