cjdic2/src-core/src/db/search.rs

use rusqlite::{Connection, Transaction, params};
use wana_kana::utils::katakana_to_hiragana;

use crate::{
    CJDicError, Timer, YomitanProgress,
    db::{DBFILE, DBSCHEMA, DbChild, yomitan_writer::normalize_term},
    tokenizer::Tokenizer,
};

pub(super) fn normalize_reading(s: &str) -> String {
    katakana_to_hiragana(s)
}

pub struct SearchDatabase<'a> {
    conn: &'a mut Connection,
}

impl<'a> SearchDatabase<'a> {
    pub(crate) fn new(conn: &'a mut Connection) -> Self {
        Self { conn }
    }

    pub fn create_schema(&mut self) -> Result<(), CJDicError> {
        let tx = self.conn.transaction()?;
        SearchDatabase::create_schema_tx(&tx)?;
        tx.commit()?;
        Ok(())
    }

    fn create_schema_tx(tx: &'a Transaction) -> Result<(), CJDicError> {
        tx.execute_batch(
            &format!("
            CREATE TABLE IF NOT EXISTS {0}.terms (
                id          INTEGER PRIMARY KEY,    -- yomitan.terms.id OR yomitan dups OR -user.terms.id
                term        TEXT NOT NULL,          -- normalized
                reading     TEXT NOT NULL,          -- normalized
                max_score   INTEGER
            );

            CREATE INDEX IF NOT EXISTS {0}.terms_term       ON terms (term);
            CREATE INDEX IF NOT EXISTS {0}.terms_reading    ON terms (reading);
            CREATE INDEX IF NOT EXISTS {0}.terms_max_score  ON terms (max_score DESC);

            CREATE VIRTUAL TABLE IF NOT EXISTS {0}.terms_ft USING fts5 (
                term,
                content='',
                tokenize=unicode61
            );
        ", DBSCHEMA[DbChild::Search]),
        )?;

        Ok(())
    }

    pub fn regenerate_yomitan(
        &mut self,
        yomitan_schema: &str,
        tokenizer: Tokenizer,
        progress_callback: impl Fn(YomitanProgress),
    ) -> Result<(), CJDicError> {
        let total = self.conn.query_row(
            &format!("SELECT COALESCE(MAX(id), 0) FROM {}.terms", yomitan_schema),
            [],
            |r| r.get::<_, i64>(0),
        )? as usize;

        if total == 0 {
            return Ok(());
        }

        let message = &format!("Generating {}", DBFILE[DbChild::Search]);
        progress_callback(YomitanProgress {
            message: message.to_string(),
            current: 0,
            total,
            steps: 0,
        });

        let _timer = Timer::new(message.to_string());

        let tx = self.conn.transaction()?;
        {
            tx.execute_batch(&format!(
                "
                DELETE FROM {0}.terms;
                DROP TABLE {0}.terms_ft;
            ",
                DBSCHEMA[DbChild::Search]
            ))?;

            SearchDatabase::create_schema_tx(&tx)?;

            let mut stmt = tx.prepare(&format!(
                "SELECT t.id, t.term, t.reading, tr.max_score
                FROM {0}.terms t
                LEFT JOIN {0}.view_terms_term_rank tr ON tr.term = t.term AND tr.reading = t.reading
                ",
                yomitan_schema
            ))?;

            let rows = stmt.query_map([], |r| {
                Ok((
                    r.get::<_, i64>(0)?,
                    r.get::<_, String>(1)?,
                    r.get::<_, String>(2)?,
                    r.get::<_, Option<i64>>(3)?,
                ))
            })?;

            let mut stmt_write = tx.prepare(&format!(
                "INSERT INTO {}.terms (id, term, reading, max_score) VALUES (?1, ?2, ?3, ?4)",
                DBSCHEMA[DbChild::Search]
            ))?;

            let mut stmt_write_ft = tx.prepare(&format!(
                "INSERT INTO {}.terms_ft (rowid, term) VALUES (?1, ?2)",
                DBSCHEMA[DbChild::Search]
            ))?;

            for (i, row) in rows.enumerate() {
                let (id, term, reading, max_score) = row?;
                stmt_write.execute(params![
                    id,
                    normalize_term(&term),
                    normalize_reading(&reading),
                    max_score
                ])?;

                stmt_write_ft.execute(params![id, tokenizer.ja_normalize(term).join(" ")])?;

                if i % 10_000 == 0 {
                    progress_callback(YomitanProgress {
                        message: message.to_string(),
                        current: i,
                        total,
                        steps: i,
                    });
                }
            }
        }
        tx.commit()?;

        Ok(())
    }
}