Rust 的 SQL Optimistic Lock

上一篇聊到怎麼用 PostgreSQL 的原生型別 xmin 來做 optimistic lock，身為一個 Rust 愛好者，當然想在 Rust 上面也照做一次，結果沒想到中間卡了一堆意料之外的東西。這篇就來分享一下遇到的狀況，還有最後是怎麼解決的。

挑一個 SQL Client

Rust 裡面能連 SQL 資料庫的 crate 其實不少。我自己是 .NET 出身的，想找的當然是像 EF Core 那樣強大的東西，而比較有名的 ORM crate 有兩個：Diesel 跟 SeaORM。因為 Diesel 不支援 async，一開始我是打算用 SeaORM 的，不過這兩套用起來都沒有 EF Core 那麼方便，可能是 Rust 跟 C# 的機制本來就差蠻多的關係。最後選了 sqlx，一方面 SeaORM 本身就是建立在 sqlx 上面，另一方面直接用 SQL 操作資料反而比較直觀。

第一次嘗試

跟上次在 Dotnet 做的一樣，先準備一張測試用的資料表：

CREATE TABLE IF NOT EXISTS test_table (
    id      UUID        PRIMARY KEY,
    value   VARCHAR
)

反正只是簡單測試，connection pool 就先不管了，直接用 PgConnection 塞資料進去：

let id = Uuid::new_v4();

let xmin: u32 = sqlx::query_scalar(
    r#"
    INSERT INTO test_table (id, value) VALUES ($1, $2) RETURNING xmin
    "#,
)
.bind(id)
.bind("test")
.fetch_one(&mut conn)
.await
.expect("Insert row fail.");

結果就出事了，有點意外：

error[E0277]: the trait bound `u32: sqlx::Decode<'_, sqlx::Postgres>` is not satisfied
   --> lab_optimistic_lock/src/lib.rs:63:10
    |
63  |         .fetch_one(&mut conn)
    |          ^^^^^^^^^ the trait `sqlx::Decode<'_, sqlx::Postgres>` is not implemented for `u32`, which is required by `for<'r> (_,): FromRow<'r, PgRow>`

原因是 sqlx-postgres 沒有支援 Rust 的 u32 跟 Postgres 型別之間的轉換，Rust 與 Postgres 型別的對應可以在這裡查到。那換成 i64 應該可以吧。

let xmin: i64 = sqlx::query_scalar(
    r#"
    INSERT INTO test_table (id, value) VALUES ($1, $2) RETURNING xmin
    "#,
)
.bind(id)
.bind("test")
.fetch_one(&mut conn)
.await
.expect("Insert row fail.");

還是噴錯：

Insert row fail.: ColumnDecode { index: "0", source: "mismatched types; Rust type `i32` (as SQL type `INT4`) is not compatible with SQL type `xid`" }

這次是因為 PostgreSQL 的 driver 根本沒支援 Rust 型別跟 PostgreSQL Xid 之間的轉換。我後來也翻了 Diesel、SeaORM 跟 rust-postgres，通通都沒有，看來只能自己動手實作了。

實作 Rust 型別與 Postgres Xid 之間的轉換

sqlx 的 PostgreSQL driver 其實有幫 PostgreSQL 的 oid 做另一組轉換，對照 Postgres 的文件來看，拿來當範本應該不錯。

先在 Rust 這邊建一個新的 struct 來代表 Xid：

pub struct Xid(pub u32);

Xid 需要實作三個 trait，這邊用的是包一層 u32 的 wrapper。

• Type trait

用來標示 xid 這個型別的 oid。上次提過，oid 是 PostgreSQL 用來代表物件的值。

impl Type<Postgres> for Xid {
    fn type_info() -> PgTypeInfo {
        PgTypeInfo::with_oid(Oid(28))
    }
}

PgTypeInfo 這個 enum 裡面沒有 xid，不過自己定義就好，xid 的值是 28。

• Encode

impl Encode<'_, Postgres> for Xid {
    fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> IsNull {
        buf.extend(&self.0.to_be_bytes());

        IsNull::No
    }
}

這段是告訴 sqlx 怎麼把 PostgreSQL 型別轉成 Rust 的 struct。

• Decode

impl Decode<'_, Postgres> for Xid {
    fn decode(value: PgValueRef<'_>) -> Result<Self, BoxDynError> {
        Ok(Self(match value.format() {
            PgValueFormat::Binary => BigEndian::read_u32(value.as_bytes()?),
            PgValueFormat::Text => value.as_str()?.parse()?,
        }))
    }
}

這段就是 Encode 的反向。

三個 trait 都實作完之後，就能順利拿到剛塞進去那筆資料的 XID 了。

其實 sqlx-postgres 有提供 macro 來簡化這整段：

#[derive(sqlx::Type)]
#[sqlx(type_name = "xid")]
pub struct Xid(pub i32);

不過沒辦法直接吃 u32，只能用 i32 之類的其他型別繞過去。

End

完整的 optimistic lock 範例放在這裡，有任何問題或想法都歡迎跟我說！