用 Rust 扩展 SQLite
作为进程内数据库，SQLite 具有其他扩展机制，例如 用户定义函数（简称 UDF）。但是UDF有一些缺点：
- UDF 在 SQLite 连接中是当前生效的，而不是为所有连接共享；
- UDF 必须在程序中定义。这意味着您需要在与您的应用程序相同的作用域内使用该功能。
这就是 UDF 的用武之地。UDF 可以用任何可以编译为共享库或者 DLL 的编程语言编写。然后，您可以共享已编译的对象并从任何应用程序或编程语言中加载它们。在这篇文章中，我们将看到如何使用Rust编写 SQLite 可加载扩展。
借鉴
我们可以从 phiresky/sqlite-zstd 学到的 SQLite UDF 简化版本技术。这是一个在 SQLite 上启用 zstd 压缩的 SQLite 扩展，如果您想查看比这篇文章更高级的示例，我强烈建议您查看它。
Cargo.toml
- 需要依赖：https://github.com/Genomicsplc/rusqlite/tree/loadable-extensions 原因是 #910
- 合并版本为：https://github.com/litements/rusqlite/tree/loadable-extensions-release-2
- crate-type 标记为：["cdylib"]。这将告诉 rust 编译器我们正在构建一个共享库。
[package]
name = "sqlite-regex"
version = "0.1.0"
edition = "2021"
[features]
default = []
build_extension = ["rusqlite/bundled", "rusqlite/functions"]
[lib]
crate-type = ["cdylib"]
[dependencies]
regex = "1.5.4"
log = "0.4.14"
env_logger = "0.9.0"
anyhow = "1.0.54"
[dependencies.rusqlite]
package = "rusqlite"
git = "https://github.com/litements/rusqlite/"
branch = "loadable-extensions-release-2"
default-features = false
features = ["loadable_extension", "vtab", "functions", "bundled"]
编写扩展和入口函数
函数 ah 作用是会将一个 anyhow 错误转换为 rusqlite 错误；函数 init_logging 则将设置 env_logger。
#![allow(clippy::missing_safety_doc)]
use crate::ffi::loadable_extension_init;
use anyhow::Context as ACtxt;
use log::LevelFilter;
use regex::bytes::Regex;
use rusqlite::ffi;
use rusqlite::functions::{Context, FunctionFlags};
use rusqlite::types::{ToSqlOutput, Value, ValueRef};
use rusqlite::Connection;
use std::os::raw::c_int;
fn ah(e: anyhow::Error) -> rusqlite::Error {
rusqlite::Error::UserFunctionError(format!("{:?}", e).into())
}
fn init_logging(default_level: LevelFilter) {
let lib_log_env = "SQLITE_REGEX_LOG";
if std::env::var(lib_log_env).is_err() {
std::env::set_var(lib_log_env, format!("{}", default_level))
}
let logger_env = env_logger::Env::new().filter(lib_log_env);
env_logger::try_init_from_env(logger_env).ok();
}
当您尝试在 SQLite 中加载 UDF 时，它首先需要一个入口点函数。根据 sqlite3_load_extension C 语言函数文档， 如果没有提供入口，它将根据文件名进行猜测。如果我们调用已编译的扩展 regex_ext，它将尝试加载一个名为 sqlite3_regex_ext_init 的入口，因为该扩展具有文件名。 regex_ext.{so,dll,dylib}。如果您需要更大的灵活性，还有一个 SQL 函数来加载，它可以让您指定入口点。有了它，您可以执行以下操作：
SELECT load_extension('path/to/loadable/extension/regex_ext.[extension]', 'sqlite3_regex_init')
现在它将尝试找到一个称为 sqlite3_regex_init 入口点的函数，而不是 sqlite3_regex_ext_init .
入口函数：
#[no_mangle]
pub unsafe extern "C" fn sqlite3_regex_init(
db: *mut ffi::sqlite3,
_pz_err_msg: &mut &mut std::os::raw::c_char,
p_api: *mut ffi::sqlite3_api_routines,
) -> c_int {
loadable_extension_init(p_api);
match init(db) {
Ok(()) => {
log::info!("[regex-extension] init ok");
ffi::SQLITE_OK
}
Err(e) => {
log::error!("[regex-extension] init error: {:?}", e);
ffi::SQLITE_ERROR
}
}
}
init 函数
fn init(db_handle: *mut ffi::sqlite3) -> anyhow::Result<()> {
let db = unsafe { rusqlite::Connection::from_handle(db_handle)? };
load(&db)?;
Ok(())
}
这里 from_handle 是 rusqlite 的方法
将 rust 函数加载到 SQLite
fn add_functions(c: &Connection) -> anyhow::Result<()> {
let deterministic = FunctionFlags::SQLITE_DETERMINISTIC | FunctionFlags::SQLITE_UTF8;
c.create_scalar_function("regex_extract", 2, deterministic, |ctx: &Context| {
regex_extract(ctx).map_err(ah)
})?;
c.create_scalar_function("regex_extract", 3, deterministic, |ctx: &Context| {
regex_extract(ctx).map_err(ah)
})?;
Ok(())
}
我们正在使用 rusqlite 方法 create_scalar_function 。如果您阅读 SQLite 文档，您会看到sqlite3_create_function() 接收 5 个参数，第一个参数 db 已经隐含在我们的 rust 代码中，因为 create_scalar_function 是 Connection 对象上的一个方法 ，所以 db 信息已经在 self 上. 这意味着在代码中实现使用 4 个参数。剩下的第一个参数是我们想要在 SQLite 中注册函数的名称，如果我们传递 value "regex_extract"，我们将能够像regex_extract()在 SQL 查询中一样使用这个函数。第二个参数是函数接受的参数数量。
最后在 rust 中使用
fn regex_extract<'a>(ctx: &Context) -> anyhow::Result<ToSqlOutput<'a>> {
let arg_pat = 0;
let arg_input_data = 1;
let arg_cap_group = 2;
let empty_return = Ok(ToSqlOutput::Owned(Value::Null));
let pattern = match ctx.get_raw(arg_pat) {
ValueRef::Text(t) => t,
e => anyhow::bail!("regex pattern must be text, got {}", e.data_type()),
};
let re = Regex::new(std::str::from_utf8(pattern)?)?;
let input_value = match ctx.get_raw(arg_input_data) {
ValueRef::Text(t) => t,
ValueRef::Null => return empty_return,
e => anyhow::bail!("regex expects text as input, got {}", e.data_type()),
};
let cap_group: usize = if ctx.len() <= arg_cap_group {
// no capture group, use default
0
} else {
ctx.get(arg_cap_group).context("capture group")?
};
// let mut caploc = re.capture_locations();
// re.captures_read(&mut caploc, input_value);
if let Some(cap) = re.captures(input_value) {
match cap.get(cap_group) {
None => empty_return,
// String::from_utf8_lossy
Some(t) => {
let value = String::from_utf8_lossy(t.as_bytes());
return Ok(ToSqlOutput::Owned(Value::Text(value.to_string())));
}
}
} else {
empty_return
}
}
新建函数，该函数接收一个 rawrusqlite::functions::Context作为输入并将其拆分到函数体内。
构建 udf
cargo build --release
现在我们可以在 SQLite REPL 中加载扩展
SELECT load_extension('target/release/libsqlite_regex.dylib', 'sqlite3_regex_init');
如果成功的话你会得到如下信息
[2022-05-14T23:22:09Z INFO sqlite_regex] [regex-extension] init ok
脚本：
#!/usr/bin/env python3
import sqlite3
conn = sqlite3.connect("test.db", isolation_level=None)
print(f"Loading SQLite extension in connection: {conn}")
conn.enable_load_extension(True)
conn.execute(
"SELECT load_extension('target/release/libsqlite_regex.dylib', 'sqlite3_regex_init');"
)
conn.enable_load_extension(False)
print("Running tests...")
print("Testing pattern 'x(ab)' WITHOUT capture group")
row = conn.execute("SELECT regex_extract('x(ab)', 'xxabaa')").fetchone()
assert row[0] == "xab", row[0]
print("Testing pattern 'x(ab)' WITH capture group = 1")
row = conn.execute("SELECT regex_extract('x(ab)', 'xxabaa', 1)").fetchone()
assert row[0] == "ab", row[0]
print("Testing pattern 'x(ab)' WITH capture group = 0")
row = conn.execute("SELECT regex_extract('x(ab)', 'xxabaa', 0)").fetchone()
assert row[0] == "xab", row[0]
print("Testing pattern 'g(oog)+le' WITHOUT capture group")
row = conn.execute("SELECT regex_extract('g(oog)+le', 'googoogoogle')").fetchone()
assert row[0] == "googoogoogle", row[0]
print("Testing pattern 'g(oog)+le' WITH capture group = 1")
row = conn.execute("SELECT regex_extract('g(oog)+le', 'googoogoogle', 1)").fetchone()
assert row[0] == "oog", row[0]
print("Testing pattern '[Cc]at' WITHOUT capture group")
row = conn.execute("SELECT regex_extract('[Cc]at', 'cat')").fetchone()
assert row[0] == "cat", row[0]
print("Testing pattern '[Cc]at' WITHOUT capture group, expecting empty return")
row = conn.execute("SELECT regex_extract('[Cc]at', 'hello')").fetchone()
assert row[0] is None, row[0]
当然，新链接将报错.
以上更多例子和功能可以参考链接
- https://ricardoanderegg.com/posts/extending-sqlite-with-rust/
From 日报小组 侯盛鑫 坏姐姐
社区学习交流平台订阅：
1
共 0 条评论, 1 页
评论区
写评论还没有评论