使用 Rust 编写用户定义函数

向 MariaDB 或 MySQL 服务器实例添加功能的最直接方法之一是创建用户定义函数（UDF）。这些是加载自动态二进制文件的编译函数，它们比用 SQL 编写的函数性能更高、更灵活，并提供与内置函数相同的功能。

这些 UDF 通常用 C 或 C++ 编写，但现在有一个库可以轻松地用 Rust 编写它们。本博客讨论了编写此库的一些原因，然后提供了一个非常基础的用法示例，即使没有 Rust 语言经验也可以理解。

为什么选择 Rust？

可以用来生成编译后的动态库的任何语言都可以用于编写 MariaDB 的扩展，这通常是 C 或 C++（服务器源码本身的语言）。目前的这种方法没有任何问题，但使用 Rust 编写它们具有一些优势

• 编译器始终保证防止最常见的相关 CWE（特别是越界读/写、use after free、null 解引用和竞态条件）
• 充分利用类型安全可以强制保证代码正确性
• 智能指针防止内存泄漏（类似于 C++，但 Rust 的实现更直接一些）
• 出色的工具链；cargo 是 Rust 的默认构建系统。开箱即用，你可以获得
  • 编译 (cargo check / cargo build)
  • Linting (cargo clippy)
  • 测试 (cargo test，别名 cargo t)
  • 文档 (cargo doc，对于库通常是 cargo doc --document-private-items --open)
  • 依赖管理（在 Cargo.toml 中配置）

数据库是性能瓶颈很容易拖慢其所驱动的 Web 服务的应用程序。它们还必须非常小心地避免安全问题，因为像缓冲区越界读取这样容易被忽视的问题可能意味着敏感数据泄露。Rust 特别适合这个领域：它提供了与 C 和 C++ 相似或更好的性能，同时保证避免这些语言中最常见的安全陷阱。

如果你不熟悉这门语言，你可能会质疑 C 接口和底层编程如何能够在有这些保证的情况下实现。答案相当直接；任何可能引入 unsoundness 的操作（指针操作、潜在的线程不安全操作、内联汇编、C FFI）都是可能的，但只能在 unsafe {...} 块内部进行。这意味着这些少数隔离的代码行可以轻松识别并彻底验证；任何建立在这些块之上的安全 API 只要 unsafe 部分是健全的，就能保证是健全的。由于 udf 库为你处理了所有这些 unsafe 操作，因此几乎可以在安全代码中完全编写任何 UDF！

示例 UDF 演练

在本节中，我们将实现一个极其简单的用户定义函数，并介绍其编写、构建和使用方面。要跟着做，你需要一个 Rust 编译器，版本 ≥1.65（由于依赖于 GATs）；如有需要，运行 rustup update 更新到最新版本。如果你还没有 Rust，请从 https://rustup.rs/ 获取。如果你使用 IDE，请获取 rust-analyzer 语言服务器（非必需，但有益）。

工作区设置

第一步是创建一个新项目；Cargo 让这变得相当容易；

# Create a new Rust project called test-udf
$ cargo new --lib test-udf

$ cd test-udf

# Simply validate that the code compiles, same as `cargo c`
$ cargo check

# Run the unit test, same as `cargo t`
$ cargo test

上述命令创建了一个名为 test-udf 的目录，其中包含 Cargo.toml 和 src/lib.rs，然后验证一切正常工作（cargo new 命令会创建一个非常简单的示例函数和测试）。我们需要更新 Cargo.toml 来告诉 Cargo 生成正确类型的输出（动态库），并将 udf 用作依赖项

[package]
name = "test-udf"
version = "0.1.0"
edition = "2021"
publish = false # prevent accidentally publishing to crates.io

# Add this section to specify we want to create a C dynamic library
[lib]
crate-type = ["cdylib"]

[dependencies]
udf = "0.5" # our dependency on the `udf` crate

你可以删除 lib.rs 中的所有内容，我们的设置就完成了。

UDF 架构

让我们编写一个非常简单的 UDF，它对整数进行运行总计。

一个 UDF 通常需要向服务器提供三个符号

• 一个 init 调用，用于验证参数类型并执行内存分配
• 一个 process 调用，每行运行一次并产生结果
• 一个 deinit 调用，用于释放设置阶段分配的任何内存

Rust 中常见的这种接口被分组到 trait 中。BasicUdf trait 在这里很有趣，它提供了 init 和 process 的接口（deinit 会自动处理）。这个 trait 应该在表示要在 process 调用（每行调用一次）之间共享的数据的结构体上实现。在这种情况下，数据就是我们当前的累计总和。

struct RunningTotal(i64);

我这里使用的是“tuple struct”语法，这意味着你可以用数字访问字段（some_struct.0, some_struct.1），而不是通过名称访问（some_struct.field）。这只是一个方便的写法，因为我们只有一个字段，但你完全可以使用标准结构体（它们在底层是相同的）

struct RunningTotal {
  total: i64
}

无论哪种情况，我们的结构体名称都会被转换为蛇形命名法，以作为 SQL 函数的名称（所以 Runningtotal 变成 running_total(...)）。这将成为我们实现的基础。

基本结构

现在我们需要做三件事

• 导入所需的类型和函数。udf 有一个 prelude 模块，其中包含最常用的导入项，所以我们可以直接导入其中的所有内容；
• 为我们的结构体实现 BasicUdf trait；以及
• 添加 #[register] 宏以创建正确的符号。

最少的编译代码如下所示

use udf::prelude::*;

struct RunningTotal(i64);

#[register]
impl BasicUdf for  RunningTotal {
    type Returns<'a> = i64;

    fn init(cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
        todo!()
    }

    fn process<'a>(
        &'a mut self,
        cfg: &UdfCfg<Process>,
        args: &ArgList<Process>,
        error: Option<NonZeroU8>,
    ) -> Result<Self::Returns<'a>, ProcessError> {
        todo!()
    }
}

（提示：如果你只输入 impl BasicUdf for RunningTotal {}，然后在括号内打开快速修复 (VSCode 上按 ctrl+.)，它会提供自动填充函数签名。）

哇！函数签名！BasicUdf 的文档详细介绍了这里的一切作用，但我们来简单分解一下

type Returns<'a> = i64;

这只是我们指定 UDF 返回类型的地方。关于可能的返回类型的更多信息，请参阅文档，但我们只使用 i64 来表示非空整数。（忽略 <'a> —— 它仅在返回引用时相关）

fn init(cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
    todo!()
}

这是我们的初始化函数，它接受一个配置对象 cfg 和一个参数列表 args。这些类型上的 Init 只是一个标记，表明它们在哪里被使用（它与可用的方法相关）。

这返回一个 Result，这是一个内置的 enum，表示一个值，其 Ok 和 Success 有不同的类型（Rust 枚举是“带标签的联合”或“和类型”）。因此，从我们的函数签名中，我们可以知道成功的函数调用类型将是 Self（即 RunningTotal，它会被保存供以后使用），而错误将是一个 String（它会显示给用户）。这很有道理，对吧？

todo!() 是一个内置宏，它匹配编译所需的任何类型签名。如果我们真的尝试运行它会失败，但我们至少可以验证我们的代码结构没有错误

$ cargo c
warning: `test-udf` (lib) generated 5 warnings (run `cargo fix --lib -p test-udf` to apply 5 suggestions)
    Finished dev [unoptimized + debuginfo] target(s) in 0.07s

有关于未使用参数的警告，但基本结构都已经设置好了！

UDF 实现：init

现在我们有了基本结构，让我们看看如何获得一些结果。

我们 init 函数的主要目标是验证参数。让我们看看实现，然后分解它

fn init(_cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
    if args.len() != 1 {
        return Err(format!("expected 1 argument; got {}", args.len()));
    }

    // Coerce everything to an integer
    args.get(0).unwrap().set_type_coercion(SqlType::Int);

    Ok(Self(0))
}

第一部分检查参数数量

if args.len() != 1 {
    return Err(format!("expected 1 argument; got {}", args.len()));
}

参数数量应为一个。如果不是，它会创建一个格式化的错误消息字符串并返回它（Err(something) 是构造 Result 枚举错误变体的方式）。

第二个逻辑块

args.get(0).unwrap().set_type_coercion(SqlType::Int);

使用 .get(0) 尝试获取第一个参数。这返回一个 Option<SqlArg>，这是另一个内置的枚举类型，类似于 Result。

Option<T> 有两种可能的变体：Some(T) 表示存在一个类型为 T 的值，而 None 表示没有值。unwrap() 用于从 Some() 值中获取内部值，如果值为 None 则会 panic，所以我们用它来获取索引为 0 的参数（第一个参数）。

请注意，在 UDF 中 panicking 是一个非常糟糕的主意，因为它可能导致服务器崩溃。但是，我们这里已经验证了只有一个参数，所以 unwrapping 不会失败，因此在这里可以使用。

一旦我们获得了第一个参数表示，我们就可以调用 set_type_coercion。这会指示服务器将第一个参数强制转换为整数，否则取消函数调用并返回错误。

Ok(Self(0))

最后一行简单地创建一个内部值为 0 的 Self 实例，并将其包装在 Ok 中表示成功。在 Rust 中，任何块的最后一行，如果它不是以分号结尾，就是该块的值，所以这个语句就是我们的返回值。这就是 init 阶段所需的全部内容。

UDF 实现：process

process 函数的函数体也相当简单

fn process<'a>(
    &'a mut self,
    _cfg: &UdfCfg<Process>,
    args: &ArgList<Process>,
    _error: Option<NonZeroU8>,
) -> Result<i64, ProcessError> {
    // Get the value as an integer and add it to our total
    self.0 += args.get(0).unwrap().value().as_int().unwrap_or(0);

    // The result is just our running total
    Ok(self.0)
}

第一行包含了大部分逻辑，并使用组合器来保持代码
简洁。它做了以下几件事

• args.get(0).unwrap(): 这获取第一个参数，如上所述。我们在这里再次可以放心使用 unwrap，因为我们在 init 中已经验证了参数（init 和 process 在调用时接收相同数量的参数，并且我们在 init 中执行了验证）
• .value() 获取参数的值，这是一个 SqlResult。这是一个枚举，包含 String、Real、Int 和 Decimal 等变体。.as_int() 是一个方便的函数，返回一个 Option。如果值是一个非空整数，它将返回 Some(i64)；任何其他情况都将返回 None。因为我们在 init 中设置了类型强制转换，我们可以合理地预期所有值都将是可能为 null 的整数。
• unwrap_or(0) 在值为 Some 时与 unwrap() 作用相同，但在值为 None 时替换为指定的 0。这意味着任何 null 值都将使用 0，这样它们就不会影响我们的运行总和。

我们将此值添加到结构体的内部值 self.0 中，它代表我们当前的运行总和，然后使用 Ok(self.0) 返回。至此，我们的 process 函数就完成了！

单元测试

udf crate 提供了无需将其加载到 SQL 中即可彻底测试 UDF 实现的功能。如果你不熟悉，值得看看 Rust 的单元测试基础知识进行简要概述，但基本上任何标记为 #[test] 的函数都将作为单元测试运行。

我们需要更新 udf 依赖项以使用 mock feature，以便我们可以访问该 feature-gated 模块。将依赖项行更改为如下所示

udf = { version = "0.5", features = ["mock"] }

并在我们的测试 UDF 下面添加以下内容

#[cfg(test)]
mod tests {
    // Use our parent module
    use super::*;
    use udf::mock::*;

    #[test]
    fn test_basic() {
        // Create a mock `UdfCfg` and mock `UdfArgs`
        let mut cfg = MockUdfCfg::new();

        // Create an array of mock row arguments. We will "run" our function with each
        // one and verify the result
        let mut row_args = [
            // Each entry here acts as a row. Format for the macro is
            // `([type] value, "attribute name", nullable)`
            mock_args![(Int None, "", false)],
            mock_args![(10, "", false)],
            mock_args![(Int None, "", false)],
            mock_args![(-20, "", false)],
        ];

        // This is the expected output after each of the above calls
        let outputs = [0i64, 10, 10, -10];

        // Run the `init` function on our mock data
        let mut rt = RunningTotal::init(cfg.as_init(), row_args[0].as_init()).unwrap();

        // Iterate through our list of arguments
        for (arglist, outval) in row_args.iter_mut().zip(outputs.iter()) {
            // Run the process function and verify the result
            let res = rt.process(cfg.as_process(), arglist.as_process(), None);
            assert_eq!(res, Ok(*outval));
        }
    }
}

我们来检查一下结果

running 1 test
test tests::test_basic ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

太棒了！

有关更多示例单元测试，请参阅本博客文章的仓库。通过良好的单元测试，我们可以高度自信 UDF 按照预期执行，而无需将其加载到服务器中。

加载函数

对我们 UDF 的最终测试是在服务器中实际运行它。构建一个可以加载到 MariaDB 中的 .so 或 .dll 文件很容易 – cargo build --release 会生成输出，位于 target/release。但是，如果你有 Docker，在隔离环境中进行测试甚至更容易

FROM rust:1.66 AS build

WORKDIR /build

COPY . .

# Use Docker buildkit to cache our build directory for quicker recompilations
RUN --mount=type=cache,target=/usr/local/cargo/registry \
    --mount=type=cache,target=/build/target \
    cargo build --release \
    && mkdir /output \
    && cp target/release/libMY_CRATE_NAME.so /output

FROM mariadb:10.10

COPY --from=build /output/* /usr/lib/mysql/plugin/

这个 docker 镜像使用了缓存，这是 Docker buildkit 的一个功能。确保你使用的 Docker 版本较新，或已正确设置环境变量以启用它。或者，移除缓存指示符。

请务必更新 MY_CRATE_NAME 处的文件名（如果遵循了前面的建议，则为 libtest_udf.so）。以下命令运行并构建我们的镜像

# Build the image
$ docker build . --tag mdb-blog-udf

# Run the image and name it mariadb_blog_udf for convenience
$ docker run --rm -d -e MARIADB_ROOT_PASSWORD=example --name mariadb_blog_udf mdb-blog-udf

# Enter the SQL console
$ docker exec -it mariadb_blog_udf mysql -pexample

让我们加载我们的函数并测试它，首先使用太多参数，然后使用正确数量的参数

MariaDB [(none)]> CREATE FUNCTION running_total RETURNS integer SONAME 'libtest_udf.so';
Query OK, 0 rows affected (0.003 sec)

MariaDB [(none)]> select running_total(1, 2, 3);
ERROR 1123 (HY000): Can't initialize function 'running_total'; xpected 1 argument; got 3

MariaDB [(none)]> select running_total(10);
+-------------------+
| running_total(10) |
+-------------------+
|                10 |
+-------------------+
1 row in set (0.000 sec)

到目前为止一切顺利！现在来一个稍微难一点的测试

MariaDB [(none)]> create database db; use db; create table t1 (val int);
Query OK, 1 row affected (0.000 sec)

Database changed
Query OK, 0 rows affected (0.023 sec)

MariaDB [db]> insert into t1(val) values (1),(2),(3),(NULL),(-100),(50),(123456789);
Query OK, 7 rows affected (0.002 sec)
Records: 7  Duplicates: 0  Warnings: 0

MariaDB [db]> select val, running_total(val) from t1;
+-----------+--------------------+
| val       | running_total(val) |
+-----------+--------------------+
|         1 |                  1 |
|         2 |                  3 |
|         3 |                  6 |
|      NULL |                  6 |
|      -100 |                -94 |
|        50 |                -44 |
| 123456789 |          123456745 |
+-----------+--------------------+
7 rows in set (0.000 sec)

完美！

总结

我们成功编写了一个简单的 UDF，它验证参数并在调用之间存储一些数据，所有这些只需几行代码。我们还进行了单元测试和（非自动化）集成测试，这些是确保程序按预期工作的简单步骤。

如前所述，本示例的代码位于本博客文章的仓库中。对于希望进一步探索的人来说，通过实现 AggregateUdf trait，该代码可以轻松地转换为聚合 UDF。

有用链接

• udf 库仓库页面和文档
• MariaDB UDF 使用文档
• 用 Rust 编写的简单 UDF 套件

尝试示例或尝试编写自己的函数；欢迎加入我们的 Zulip 讨论结果。祝你编程愉快！