用 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 日报小组侯盛鑫坏姐姐

社区学习交流平台订阅：

Rust.cc 论坛: 支持 rss
微信公众号：Rust 语言中文社区

ITPub博客

【Rust日报】2022-05-20 - 用 Rust 扩展 SQLite