RustでPostgreSQLで自分で定義した型の利用

目的

前回PostgreSQLが持っている型とRustの型をマッピングしてみました。

今回は自分て定義した型でPostgreSQLとRustのstructやenumにマッピングしたいと思います。
リファレンスには定義の仕方が詳細に書いてありますが、使い方は書いてないのでこの記事が参考になると思います。

リポジトリはこちら。

コード

PostgreSQL

ここではOSの種類を表すenumとしてtype_enum_osを、整数を意味するドメインとしてdomain_integerを定義します。

それらを使ってストアードプロシージャーで値をやり取りする型としてtype_pg_get_select_compositeを定義します。

ストアードプロシージャーでは受け取った値を2回返しています。

pg_get_select_composite.sql

DROP TYPE IF EXISTS type_enum_os CASCADE;
CREATE TYPE type_enum_os AS ENUM (
  'linux',
  'mac_os',
  'windows'
);

DROP DOMAIN IF EXISTS domain_integer CASCADE;
CREATE DOMAIN domain_integer AS BIGINT CHECK(VALUE >= 0);

DROP TYPE IF EXISTS type_pg_get_select_composite CASCADE;
CREATE TYPE type_pg_get_select_composite AS (
  os type_enum_os,
  cpu_count domain_integer,
  memory_size BIGINT
);

CREATE OR REPLACE FUNCTION pg_get_select_composite(
  p_value type_pg_get_select_composite DEFAULT NULL
) RETURNS SETOF type_pg_get_select_composite AS $FUNCTION$
DECLARE
BEGIN
  RETURN NEXT p_value;
  RETURN NEXT p_value;
END;
$FUNCTION$ LANGUAGE plpgsql;

Rust

設定した値が2回になって返ってくるのを確認するRustのコードは以下のようななります。
Rust側ではmemory_sizeをNewTypeとして定義しています。
enumとdomainと型はリファレンスにある通りに書いています。

一旦全部のコードを紹介した後で利用の仕方について説明します。

composite.rs

use postgres_types::{FromSql, ToSql};

use crate::pool::PgClient;

const SQL: &str = r#"
SELECT
    ROW(t1.*)::type_pg_get_select_composite
FROM
    pg_get_select_composite(
        p_value := $1
    ) AS t1 
"#;

#[derive(Debug, ToSql, FromSql, PartialEq)]
#[postgres(name = "type_enum_os", rename_all = "snake_case")]
enum TypeEnumOs {
    Linux,
    MacOs,
    Windows
}

#[derive(Debug, ToSql, FromSql, PartialEq)]
#[postgres(name = "domain_integer", rename_all = "snake_case")]
struct DomainInteger(i64);

#[derive(Debug, ToSql, FromSql, PartialEq)]
#[postgres(transparent)]
struct MemorySizeUnit(i64);

#[derive(Debug, ToSql, FromSql, PartialEq)]
#[postgres(name = "type_pg_get_select_composite", rename_all = "snake_case")]
struct TypePgGetSelectComposite {
    os: TypeEnumOs,
    cpu_count: DomainInteger,
    memory_size: MemorySizeUnit,
}

pub(crate) async fn execute(pg_client: &PgClient) -> anyhow::Result<()> {
    let data = TypePgGetSelectComposite {
        os: TypeEnumOs::Linux,
        cpu_count: DomainInteger(2),
        memory_size: MemorySizeUnit(1024),
    };

    let rows: Vec<TypePgGetSelectComposite> = pg_client.query(SQL, &[&data]).await?
        .iter()
        .map(|row| row.get(0))
        .collect();

    assert_eq!(rows[0], data);
    assert_eq!(rows[1], data);
    println!("{:?}", rows[0]);

    Ok(())
}

SQLの呼び出し方に特徴があります。ROW()を使って行を1つのカラムに変換しています。これはRustのRow型から直接定義したStructに変換できないからです。変換できるのはカラムだけです。

const SQL: &str = r#"
SELECT
    ROW(t1.*)::type_pg_get_select_composite
FROM
    pg_get_select_composite(
        p_value := $1
    ) AS t1 
"#;

Rowから目的のstructに書き換える部分は以下になります。上記で説明した通り最初のカラムを目的のstructに変換しています。

let rows: Vec<TypePgGetSelectComposite> = pg_client.query(SQL, &[&data]).await?
        .iter()
        .map(|row| row.get(0))
        .collect();

まとめ

PostgreSQLで自分で定義した型とRustのstructやenumと相互に変換できるようになりました。
ストアードプロシージャーを作ると自分で戻る型を作ることが多いので、今回の内容が役に立つと思います。