PostgreSQLのカスタム型をRustのstructで受け取る

目的

PostgreSQLの検索結果をRustで受け取るには、単純にやると出現するカラムの位置や名前に対して適切な型で受け取る必要がありコード量が多くなり面倒です。SELECTした結果をstructに変換して受け取れれば便利に扱うことができます。

プログラム

SQL

ここではPostgreSQLのカスタム型とストアードプロシージャを使います。ストアードプロシージャごとに型をつけると型を消すことでストアードプロシージャを削除できるのが便利です。

test_get_list.sql

DROP TYPE IF EXISTS type_test_get_list CASCADE;
CREATE TYPE type_test_get_list AS (
  id BIGINT
  ,name TEXT
);

CREATE OR REPLACE FUNCTION test_get_list(
) RETURNS SETOF type_test_get_list AS $FUNCTION$
DECLARE
BEGIN
  RETURN QUERY SELECT
    1::BIGINT
    ,'予定表～①💖ﾊﾝｶｸだ'
  ;
  RETURN QUERY SELECT
    2::BIGINT
    ,'aaa'
  ;
END;
$FUNCTION$ LANGUAGE plpgsql;

Cargo.toml

以下のCargo.tomlを共通して使います。

Cargo.toml

[package]
name = "pg"
version = "0.1.0"
edition = "2018"

[dependencies]
postgres-types = {version="~0.2.0", features=["derive", "with-serde_json-1"]}
serde = {version="~1", features=["derive"]}
serde_json = "~1"
tokio = {version="~1.2", features=["macros", "rt-multi-thread"]}
tokio-postgres = "~0.7.0"

単純に変換する方法

以下では1カラムづつstructの要素に展開しています。一々名前を文字列で与えないといけないあたりタイポしやすく良くないコードです。数が多くなると書くのも辛くなります。

main.rs

use postgres_types::{FromSql};
use serde::{Serialize, Deserialize};

# [derive(Debug, Serialize)]
struct TypeTestGetList {
    id: i64,
    name: String,
}

async fn execute(client: &tokio_postgres::Client
) -> Result<Vec<TypeTestGetList>, Box<dyn std::error::Error>> {
    Ok(client
        .query("SELECT t1.* FROM test_get_list() AS t1", &[])
        .await?
        .iter()
        .map(|row| TypeTestGetList{
            id: row.get("id"),
            name: row.get("name"),
        }).collect())
}

# [tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Connect to the database.
    let (client, connection) =
        tokio_postgres::connect("postgres://user:pass@localhost:5432/test", tokio_postgres::NoTls).await?;

    // The connection object performs the actual communication with the database,
    // so spawn it off to run on its own.
    tokio::spawn(async move {
        if let Err(e) = connection.await {
            eprintln!("connection error: {}", e);
        }
    });

    let value = execute(&client).await?;
    println!("{}", serde_json::to_string(&value)?);
    Ok(())
}

JSONを利用する

SQLの返す値をJSONにしてstructのDeserilizeを使って受け取る方法です。SQL文の中でto_jsonしているのが肝です。カラムごとに展開するコードは無くなりました。しかしPostgreSQLでJSONの作りRust側でそれを戻す処理が行われるため余計な負荷がかかります。

main.rs

# [derive(Debug, Serialize, Deserialize)]
struct TypeTestGetList {
    id: i64,
    name: String,
}

async fn execute(client: &tokio_postgres::Client
) -> Result<Vec<TypeTestGetList>, Box<dyn std::error::Error>> {
    Ok(client
        .query("SELECT to_json(t1.*) FROM test_get_list() AS t1", &[])
        .await?
        .iter()
        .map(|row| {
            let json: serde_json::Value = row.get(0);
            serde_json::from_value(json).unwrap()
        })
        .collect())
}

structでカスタム型を指定する

FromSqlを使うと直接structに型を指定できます。
記述も簡単になって余計な負荷もありません。
コメントアウトしているクエリーを使うとSQLを簡単に書けますが、WHERE句をつけることが出来なくなるので、適切な方を使ってください。

main.rs

# [derive(Debug, FromSql, Serialize)]
# [postgres(name = "type_test_get_list")]
struct TypeTestGetList {
    id: i64,
    name: String,
}

async fn execute(client: &tokio_postgres::Client
) -> Result<Vec<TypeTestGetList>, Box<dyn std::error::Error>> {
    Ok(client
        .query("SELECT ROW(t1.*)::type_test_get_list FROM test_get_list() AS t1", &[])
        //.query("SELECT test_get_list()", &[])
        .await?
        .iter()
        .map(|row| row.get(0))
        .collect())
}

まとめ

カスタム型を使ってPostgreSQLの結果を簡単にRustのstructに変換することができました。
ちなみにカスタム型を使わないでstructに受け取るにはJSONの方法を使ってください。こちらの方が自由度は高くなります。

おまけ

せっかくなのでToSqlの方も試してみました。入力パラメーターにも型が付くのでRust的にはうれしい。あとNULLと配列も試してみました。

stored.sql

DROP TYPE IF EXISTS type_test_get_list2 CASCADE;
CREATE TYPE type_test_get_list2 AS (
  id BIGINT
  ,name TEXT
  ,nullable TEXT
  ,ary TEXT[]
);

DROP TYPE IF EXISTS type_in_test_get_list2 CASCADE;
CREATE TYPE type_in_test_get_list2 AS (
  id BIGINT
  ,name TEXT
);


CREATE OR REPLACE FUNCTION test_get_list2(
  p_parameter type_in_test_get_list2
) RETURNS SETOF type_test_get_list2 AS $FUNCTION$
DECLARE
BEGIN
  RETURN QUERY SELECT
    1::BIGINT
    ,'予定表～①💖ﾊﾝｶｸだ'
    ,'a'
    ,ARRAY['x', 'y', 'z']
  ;
  RETURN QUERY SELECT
    2::BIGINT
    ,'aaa'
    ,NULL::TEXT
    ,ARRAY[]::TEXT[]
  ;
  RETURN QUERY SELECT
    p_parameter.id
    ,p_parameter.name
    ,NULL::TEXT
    ,ARRAY[]::TEXT[]
  ;
END;
$FUNCTION$ LANGUAGE plpgsql;

main.rs

use postgres_types::{FromSql, ToSql};
use serde::{Serialize, Deserialize};

# [derive(Debug, FromSql, Serialize, Deserialize)]
# [postgres(name = "type_test_get_list2")]
struct TypeTestGetList2 {
    id: i64,
    name: String,
    nullable: Option<String>,
    ary: Vec<String>
}

# [derive(Debug, ToSql, FromSql)]
# [postgres(name = "type_in_test_get_list2")]
struct TypeInTestGetList2 {
    id: i64,
    name: String,
}

async fn execute(client: &tokio_postgres::Client
) -> Result<Vec<TypeTestGetList2>, Box<dyn std::error::Error>> {
    let parameter = TypeInTestGetList2 {
        id: 999,
        name: "zzz".to_owned()
    };
    let sql = r#"
        SELECT 
            ROW(t1.*)::type_test_get_list2 
        FROM 
            test_get_list2(
                p_parameter := $1
            ) AS t1
    "#;
    Ok(client
        .query(sql, &[&parameter])
        .await?
        .iter()
        .map(|row| row.get(0))
        .collect())
}