How We Built Postgres Compatibility in Rust: pgwire and DataFusion

发表于 更新于 温馨的阅读时光

原文：How We Built Postgres Compatibility in Rust: pgwire and DataFusion

背景：Greptime

中文名是格睿科技，其实应该反过介绍，因为这是一家中国公司，自研数据库开发，用于解决可观测性特点等等，官网

本文的作者正是这家公司的CTO

让我们回到本文的重点，他的这篇博客

引入：两种模式

Pg很火，rust也很火，那么如何结合二者？作者认为主要有两种方式：

  • 自底向上的，主要用pgrx库以插件的形式把rust生态引入PostgreSQL，典型案例比如ParadaDB
  • 自顶向下的，通过模拟协议和接口实现另类的“Postgres”，作者的Greptime DB就属于此类

接口：Postgres Protocol

狭义上,Postgres Protocol指的是基于TCP协议的和pg server端通信用的应用层协议，这也是事实上连接到数据库时在TCP层发生的事

总体而言，协议主要包含以下五部分：

  • Startup: 在连接建立时，握手交换信息和各种身份认证机制
  • Simple Query: 基于文本的查询和回应
  • Extended Query: 扩展查询，即预编译，可以把语句和参数分开传
  • Copy: 导入/导出大量数据
  • Cancel: 处理终止信号

实际上这五类是作者总结概括之后的，实际上会分的更加琐碎

为了实现比较完整的pg兼容性，狭义上的这些协议只是第一层，还需要支持广义的协议，包括：

  • Query Language: SQL方言和关键函数
  • Data Type System: 建立一个到pg独有的类型的映射
  • pg_catalog Metadata: 对元数据信息系统的支持

架构：优缺点

相比于自己设计一套协议，复用pg的有很多优点：

  • 已经证明的可靠性
  • 解锁无尽的工具: 不同语言的驱动、管理工具、BI(Business Intelligence) Tools等等直接用PostgreSQL的已经非常完善的生态即可
  • 允许把Federated Queries作为一个FDW的数据源: FDW(Foreign Data Wrapper)，外部数据包装器，所以这个特性简单的说就是可以把GreptimeDB挂载到Pg上，甚至可以对两个数据库的表做JOIN操作

当然也有缺点：

  • Row-Oriented: Pg是关系型数据库，其协议设计时也是面向行传输的，而GreptimeDB却是按列处理的，这就需要一次转换
  • 取消查询机制: 查询取消机制和pg的进程模型绑定，所以很麻烦且有安全风险
  • 对于拓展查询的类型猜测: 拓展查询需要数据库在不实际执行时对表达式里的参数做类型推断，SQLite和DuckDB就没有提供这种API，因而做不到这种层面的兼容

实现：pgwire

GreptimeDB 用了作者写的pgwire库和DataFusion生态完成了实现

传输层协议支持

如果把pg的应用层协议类比http，那么pgwire就像hyper甚至axum。用户可以为不同程度的pg兼容性选择性的使用

以simple query为例，这个子协议甚至不假设用户的输入是SQL。pgwire提供了两个水平的api：

  • Low-level API: SimpleQueryHandler::on_query: 重点关注消息级别的处理和语义。
  • High-level API: SimpleQueryHandler::do_query: 重点关注数据级别的处理和语义。
#[async_trait]
impl SimpleQueryHandler for EchoHandler {
    async fn do_query<C>(&self, _client: &mut C, query: &str) ->
    PgWireResult<Vec<Response>>
    where
        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
        C::Error: Debug,
        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>
    {
        let query = query.to_string();
        let f1 = FieldInfo::new("input".into(), None, None, Type::VARCHAR,
        FieldFormat::Text);
        let schema = Arc::new(vec![f1]);
        let data = vec![Some(query)];
        let mut encoder = DataRowEncoder::new(schema.clone());
        let data_row_stream = stream::iter(data).map(move |r| {
            encoder.encode_field(&r)?;
            Ok(encoder.take_row())
        });
        Ok(vec![Response::Query(QueryResponse::new(
            schema,
            data_row_stream,
        ))])
    }
}

忽略掉泛型以及相应的若干特征，函数体内做的逻辑还是比较清晰的：

创建一列，字段名input，类型字符串，用Arc包装，把用户的查询用Option包一下，vec!处理一下发给encoder得到字节流，这样有表头有数据，就可以发给客户端了

运行参考：

❯ psql -h 127.0.0.1 -p 5432 -U postgres
psql (18.1, server 16.6-pgwire-0.38.2)
Type "help" for help.

postgres=# hello world;
  input
--------------
 hello world;
(1 row)

嗯，不是SQL也可以

pg_catalog支持

更深的兼容需要元数据支持，这就是需要pg_catalog的地方了

pg_catalog需要一系列系统表和视图，而这需要建立在查询引擎之上

GreptimeDB使用DataFusion查询引擎，于是作者维护了一个datafusion-postgres项目作为pgwire DataFusion和 Arrow数据格式的适配器

然而作者表示这个工作实在是太复杂了，他们只专注于了五张主要的表（pg_database pg_namespace pg_tables pg_class pg_attributes），把其余的留给了社区

结语

作者把pg的兼容做成了可复用的库，这意味这不仅GreptimeDB可以用，其他数据基础设施也可以用它构建一个"Postgres-like"生态

# rust# postgreSQL