libsimple/
lib.rs

1#![doc = include_str!("../README.md")]
2#![cfg_attr(docsrs, feature(doc_cfg))]
3#![warn(missing_docs)]
4
5#[cfg(feature = "jieba")]
6use std::path::Path;
7
8use better_embedded::strategies::DefaultCheckStrategy;
9use rusqlite::ffi::{sqlite3_auto_extension, sqlite3_cancel_auto_extension};
10
11use crate::ffi::sqlite3_simple_init;
12
13pub mod ffi;
14
15/// Enable sqlite3_simple_init() as an auto extension.
16pub fn enable_auto_extension() -> rusqlite::Result<()> {
17    let res = unsafe { sqlite3_auto_extension(Some(sqlite3_simple_init)) };
18    ffi::check_err(res)
19}
20
21/// Disable sqlite3_simple_init() as an auto extension.
22pub fn disable_auto_extension() -> rusqlite::Result<()> {
23    let res = unsafe { sqlite3_cancel_auto_extension(Some(sqlite3_simple_init)) };
24    ffi::check_err(res)
25}
26
27/// Release dict files into directory.
28/// Only need to call this method once.
29///
30/// Then you may call [`set_dict`] for each connection.
31#[cfg(feature = "jieba")]
32#[cfg_attr(docsrs, doc(cfg(feature = "jieba")))]
33pub fn release_dict(directory: impl AsRef<Path>) -> std::io::Result<()> {
34    let directory = directory.as_ref().to_path_buf();
35    if !directory.is_dir() { std::fs::create_dir_all(&directory)?; }
36
37    macro_rules! embedded_file {
38        ($target: ident, $source: expr) => {
39            let file = include_bytes!(concat!("../cppjieba/dict/", $source));
40            let target = $target.join($source);
41            better_embedded::release_file_with_check(file, &target, DefaultCheckStrategy::config())?;
42        };
43    }
44    embedded_file!(directory, "jieba.dict.utf8");
45    embedded_file!(directory, "user.dict.utf8");
46    embedded_file!(directory, "hmm_model.utf8");
47    embedded_file!(directory, "idf.utf8");
48    embedded_file!(directory, "stop_words.utf8");
49
50    Ok(())
51}
52
53/// Only need to call once for each connection,
54/// but must call this function before using sql `jieba_query`.
55///
56/// You should call [`release_dict`] first.
57#[cfg(feature = "jieba")]
58#[cfg_attr(docsrs, doc(cfg(feature = "jieba")))]
59pub fn set_dict(connection: &rusqlite::Connection, directory: impl AsRef<Path>) -> rusqlite::Result<()> {
60    let directory = directory.as_ref();
61    let directory = directory.to_str()
62        .ok_or_else(|| rusqlite::Error::InvalidPath(directory.to_path_buf()))?;
63    connection.query_row("SELECT jieba_dict(?)", rusqlite::params![directory], |_| Ok(()))
64}
65
66#[cfg(test)]
67mod tests {
68    #[test]
69    fn test() -> anyhow::Result<()> {
70        crate::enable_auto_extension()?;
71        let dir = tempfile::tempdir()?;
72        crate::release_dict(&dir)?;
73
74        let conn = rusqlite::Connection::open_in_memory()?;
75        crate::set_dict(&conn, &dir)?;
76        conn.execute_batch("
77            PRAGMA key = '123456';
78            CREATE TABLE singer (id INTEGER, name TEXT);
79            CREATE VIRTUAL TABLE d USING fts5(id, name, tokenize = 'simple');
80            CREATE TRIGGER dtrigger AFTER INSERT ON singer BEGIN
81                INSERT INTO d(id, name) VALUES (new.id, new.name);
82            END;
83            INSERT INTO singer (id, name) VALUES (1, '中华人民共和国国歌');
84        ")?;
85        assert_eq!(1, conn.query_row(
86           "SELECT id FROM d WHERE name MATCH jieba_query('中华国歌')",
87           [], |row| row.get::<_, i64>(0)
88        )?);
89        Ok(())
90    }
91}
libsimple/lib.rs

libsimple/
lib.rs
