我被 pgx 及其背后的 Rust 美学征服

tyrchen

发布于 2022-12-05 14:29:30

1.1K0

发布于 2022-12-05 14:29:30

知道我的人都了解，自 2018 年比较正式地学习 Rust 以来（在此要感谢张汉东老师的大力推荐），我慢慢被 Rust 征服，成为一名不折不扣的拥趸。我的业余项目，90% 都是用 Rust 写就的，另外 10% 基本被 typescript（前端）和 python（主要是 notebook）瓜分。我对 Rust 热爱也体现在我的公众号和 B 站上，近两年发布的内容，主要和 Rust 有关。然而，我很少直接吹捧 Rust，更多是通过 “show me the code” 来展示 Rust 的美妙。这个周末，在 reddit/rust 版，我无意发现了 pgx 这样一个使用 Rust 来撰写 postgres extension 的集成工具，在深入地了解其文档并写了几百行代码后，我立刻就被那种直击心灵的简约之美冲破了防线，不得不在此吹上一波。如此优雅地解决另一个生态系统（postgres）的扩展的问题，我就想说，除了 Rust，还有谁？

我相信，虽然我的读者大多在日常生活和工作中都使用过 postgres 来存储数据，也或多或少使用过 postgres extension 来扩展 postgres 的能力（比如 PostGIS，TimescaleDb），但研究过，甚至撰写过 posgres extension 的同学估计一只手都数得过来。无他 —— 你需要深入了解 postgres 的内部机理，掌握撰写 extension 的整套逻辑，妥善处理好内存管理和并发安全，并且有还算不错的 C 语言功底，才能写出一个简单的 extension。

那么，写一个 extension 究竟有多繁杂？

我们看在 postgres 里，一个非常简单的 generate_series 函数，它生成一个给定起止的列表。这样一个简单的功能，如果要用 extension 实现，核心代码大概就要 100 行，还不包括上百行的脚手架代码：

这段代码充斥了 SRF，memory context，function context 等一大堆你我并不熟悉的概念，想要写对，必然要狠狠花上一番功夫。

然而，使用 pgx 的话，包括脚手架在内的全部代码就下面这几行，核心代码就一句：

即便你没有写过 Rust，从它那简单直观的表述，你也可以清晰地了解到它想要达成的目标。

打个 90 后开发者可能无法理解的比喻，用 Rust (pgx) 之于 C 撰写 postgres extension，就好比用 VB 之于 MFC 编写 windows 应用，或者用 rails 之于 CGI 编写 web 应用。

然而，上述对比只强调了开发效率的成倍提升，却忽略了 VB/rails 潜在的性能上的损失。Rust 完全没有这个问题：

我们撰写的 my_generate_series 跟原生的，用 C 撰写的 generate_series 方法相比，效率在同一个量级，旗鼓相当。

解构 pgx 的魔法

那么，pgx 代码究竟施展了什么样的魔法，让 postgres extension 的撰写如此简单？

答案是 Rust 自身的诸多特性：内存和并发安全性，宏支持，以及和 C 的 ABI 的兼容。这些特性共同造就了 pgx 如此优雅的使用体验：

使用 pg_module_magic!() 来处理 extension 的脚手架代码。这个宏的背后是一大坨脚手架代码来设置 extension 的上下文。
使用 #[pg_extern] 来封装 Rust 函数，使其接口符合 postgres extension 的 C ABI，以及处理 Rust 数据结构和 postgres 内部数据结构的转换。
#[pg_extern] 和 default! 宏甚至可以帮助 pgx 工具链生成相关的 SQL 语句，这样当打包一个完整的 extension 时，你可以省却撰写这些 SQL 语句的痛苦。

不要忘了，Rust 还有无与伦比的正确性的保证。正如我曾经介绍过的，可以用 Rust 扩展 elixir 能力的 rustler 项目一样，你的 pgx 代码，只要编译通过，便（几乎）没有内存安全和并发安全问题；并且，如果在你的 extension 中，抛出致命异常（panic），postgres server 不会崩溃，只是执行这个操作的 transaction 被回滚而已。

这些能力，其它编程语言只具备一部分：它们或者效率不高，或者表现力不强，或者无法保证程序的正确性，或者用繁文缛节恶心死你（我发誓不是在说 java）：

真的有必要写自己的 postgres extension 呢？

很多时候，我们不去做一件事，或者想不到做这样的事情有什么意义的时候，往往可能因为我们没有能力去做。当我们被赋能的时候，无穷的想象力就会同时喷薄而出。

更好的 ID 系统？

做数据库设计的时候，我们最头疼的问题是如何设计一个有意义、高性能且保证一定随机性的 ID。自增 ID 缺乏随机性，且会暴露数据细节（黑客可以通过 id 的规律爬到大量数据）；UUID4 具备足够随机性，但无法排序。如果我想把 mongodb 的 ObjectId 或者 uuid7（可排序）引入 postgres 可以么？如果我想把应用程序内部定义的某个 ID 结构映射到 postgres 可以么？

可以！如果你想 postgres 支持 uuid7，只需要引入相应的 crate，然后写上四行代码：

#[pg_extern]
fn uuid7() -> String {
    uuid7::uuid7().into()
}

你可以仔细读一读下面的 psql 的输入输出，感受一下这几行代码带来的全新世界：

在这个例子里，我们为 postgres 引入了可排序的 uuid7。于是，我们可以在创建 test1 table 时，将其作为主键的缺省值，我可以像之前那样为 test1 插入数据，此时，生成的 id 就使用了 uuid7。为了证明它的可排序性，我生成了一张 ids table，并用 between and 寻找两个 id 中间的所有 id。

这是一个无比简单又无比实用的 extension。当然，uuid7() 这个 postgres 函数的返回值可以优化，我这里为展示方便，简单地返回了 string，效率还不算最好。

如果你没有被震撼到，那么容许我偷偷提点一句：你可以用整个 Rust 生态里的各种库来满足你对 postgres extension 的需求。这太 TM 作弊了。

更方便地定义 postgres 数据类型？

接下来我们来个更加震撼的：通过 pgx 和 serde，你可以很方便地将 Rust 类型映射到 Postgres 类型。请看图：

熟悉 Rust 的同学对这些派生宏的用法并不陌生，它们为数据结构实现了各种各样的 trait。比如这个 PhoneNumber 结构，它可以在 Rust 代码中进行相等的判定（PartialEq），比较的判定（PartialOrd），以及作为 hash map 的 key（Hash）。进一步的，它还实现了 postgres 类型的这些操作。我们没写几行代码，就在 postgres 中生成了下面一大堆以 phonenumber_ 为前缀的函数：

还进一步生成了一大堆 SQL 操作符的定义（上百行 SQL，这里只截取等号的定义）：

这真的是对那些吭哧吭哧用 C 写 extension 的程序员的降维打击啊！

我们来回复一下心情。

想想看，原本在数据库中你是怎么存储电话号码的？字符串？ok，如果让你把北京的电话号码查询出来，你该如何去做？使用 like 查询？或者把表结构更改成更利于查询的结构（把区号独立出来）？

现在，通过自定义类型 PhoneNumber，你可以用在数据库中用更好的数据结构来表达你的数据，且无痛支持原生的 SQL 操作符。一个简单的 Rust 数据结构的定义，辅以一些宏修饰，就达到了几百行 C 代码的效果。

空间和时间，我一个都不想放弃？

既然我们在拿着榔头（pgx）到处找钉子的路上越走越远，那么，我们来个更加疯狂的想法。假设你做了一款神奇的区块链应用，你用数据库存储用户的钱包地址和公钥的关系。一般而言，钱包地址是公钥派生出来的，如果我们想从钱包地址查询到公钥，那么就需要创建表，把二者都储存起来。这样虽然满足了查询的需求，但数据包含没有意义的冗余。有没有可能只存公钥，不存钱包地址就能完成这个查询呢？可是 Postgres 并不知道它们是如何映射的啊？

使用 pgx 我们可以创建一个 wallet 函数，声明它是 immutable 的，然后在其内部进行公钥到地址的转换（假设我们的钱包使用 blake3 做 hash）。注意，这里我为了演示方便，都是用了 base64 字符串而不是字节流：

有了 wallet 这个函数，我们就可以只使用公钥创建里面只有一个字段 pk 的查询表 keys，然后这样生成 index：

create index keys_pk_blake3 on keys (wallet(pk));

由此，如果我们要从用户的钱包地址找到公钥，可以这样查询：

select pk from keys where wallet(pk) = 'wallet_addr';

既节省了空间，又不影响效率。完整交互如下图所示（建议仔细观看）：

还有很多很多可以做的…

我们还能做很多事情。

比如可以生成复杂的 trigger。使用 SQL 处理 trigger 有很强的局限性，但写代码处理那就是另一片天地了。以我们上一篇谈到的交易系统为例，当股票的新的 OHLC 数据来临时，我们可以根据一个不断更新的中间状态计算出各种技术分析的数据，写入另一个表中。这样，在数据库侧，你就可以完成很多操作，避免在应用程序和数据之间来回地写入。

你也可以更好地索引数据。比如，使用 tantivy 做数据库中若干字段的搜索引擎 —— 我不知道这样做技术上的难度有多大，但 pgx 的创立者 ZomboDB 便构建了 extension，用 elasticsearch 取代 postgres 内置的全文索引，非常有想象力。

甚至，你可以在 Postgres 服务器内部向外发送 HTTP 请求（WTF），或者读写文档（WTF）。如果你嫌每次更新都需要重新加载 extension，你也可以尝试在某个 extension 中集成一个 wasm 运行时，或者 JS 运行时，让它可以动态加载某些功能或者执行某些脚本（WTF）。当然，这些想法过于天马行空，未必是好主意，但它们至少展示了一些有趣的可能性。

可是这样，我的应用就不是数据库无关的啦？

过去 10-20 年，随着 rails / django / phoenix 这样的胖 web 框架的崛起，使得我们沉迷于数据层使用 ORM 带来的「巨大好处」：数据库无关 —— 你只需要改改配置，就可以「轻松」在 sqlite3 / mysql / postgres / mssql 之间无缝迁移。诚然，本地测试使用 sqlite3，线上应用使用 postgres，这是 ORM 带来的好处，但可能也是唯一的好处。数据库的迁移从来就不是无缝的，即便你不使用任何 ORM 支持之外的功能，你也很难「无缝」地把生产环境中的数据从一个数据库迁移到另一个数据库。所以，数据库无关，很多时候是个自欺欺人的伪命题。

那么，撰写并使用 postgres extension 是一个好的选择么？

这取决于你能多高效地，并且正确地撰写这些扩展。之前我们做 web 应用，都尽量精简数据库内部的逻辑，这是出于这样一种考虑：当逻辑在你熟悉的代码中时，它更加容易被撰写，测试，学习以及维护。我相信没有人会认为传统的 postgres extension 的代码好维护，也很少有人有兴趣深入学习它。然而，pgx 逆转了这一切，就像上面展示的代码片段那样，你可以轻松地以和普通 Rust 代码没有太多区别的方式撰写、测试、打包以及加载你自己的扩展。我觉得这是无与伦比的体验 —— 如果你有一些 Rust 经验，我非常推荐你自己亲身使用 pgx 工具链，用 cargo pgx new 创建一个脚手架项目，什么也不用写，然后 cargo pgx run pg14 体验一下，那种感觉，就像哈利波特第一次随海格来到对角巷，见识到一个全新的世界一样。