PostgreSQLのネストした複合型をRustで取得してみる

目的

ストアードプロシージャの戻り値で複合型を受け取ることがよくあります。
ヘッダー明細のような親子のデータを取得すると今までは子のデータははJSON配列として1つのカラムとして受け取っていました。

Rustでデータを扱う時にJSONでは型が定義できないので、自前で変換するなど面倒なことになります。
ここでは親子共に複合型を定義して、それがRustのstructとしてきちんと型のあるデータに展開していきます。

また出力だけでなく、入力も複合型を使ってデータを渡します。

以前ネストしていない複合型をRustで扱っている記事も書いているので参考にしてください。
PostgreSQLのカスタム型をRustのstructで受け取る

コード

ストアードプロシージャ

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


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

DROP TYPE IF EXISTS type_user_list_output CASCADE;
CREATE TYPE type_user_list_output AS (
  user_id BIGINT
  ,name TEXT
  ,hobby type_hobby[]
);


CREATE OR REPLACE FUNCTION test_list_user(
  p_parameter type_user_list_input
) RETURNS SETOF type_user_list_output AS $FUNCTION$
DECLARE
  w_hobbies type_hobby[];
BEGIN
  SELECT
    ARRAY_AGG(t1.*)
  INTO
    w_hobbies
  FROM
    (VALUES (1, '読書', '00101'), (2, '映画', '00102')) AS t1(id, name, hobby_kbn)
  ;

  RETURN QUERY SELECT
    p_parameter.id
    ,p_parameter.name
    ,w_hobbies
  ;
END;
$FUNCTION$ LANGUAGE plpgsql;

Rust

Cargo.toml

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

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

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

models.rs

pub mod test_user_list_input;
pub mod test_user_list_output;
pub mod test_hobby;

models/test_hobby.rs

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


#[derive(Debug, FromSql, ToSql, Serialize, Deserialize)]
#[postgres(name = "type_hobby")]
pub struct Hobby{
    pub id: i64,
    pub name: String,
    pub hobby_kbn: String,
}

models/test_user_input.rs

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


#[derive(Debug, FromSql, ToSql, Serialize, Deserialize)]
#[postgres(name = "type_user_list_input")]
pub struct UserListInput{
    pub id: i64,
    pub name: String,
}

models/test_user_output.rs

use postgres_types::{FromSql, ToSql};
use serde::{Deserialize, Serialize};
use crate::models::test_hobby::Hobby;

#[derive(Debug, FromSql, ToSql, Serialize, Deserialize)]
#[postgres(name = "type_user_list_output")]
pub struct UserListOutput{
    pub user_id: i64,
    pub name: String,
    pub hobby: Vec<Hobby>,
}

main.rs

pub mod models;
use crate::models::{
    test_user_list_input::UserListInput,
    test_user_list_output::UserListOutput,
};

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

結果

[
  {
    "user_id":999,
    "name":"zzz",
    "hobby":[
      {"id":1,"name":"読書","hobby_kbn":"00101"},
      {"id":2,"name":"映画","hobby_kbn":"00102"}
    ]
  }
]

おまけ

以前ERD Toolの紹介という記事を書きました。
今回のコードはERD Toolの新機能であるinterfaceとstructを利用しています。

ストアードプロシージャはDECLAREからEND;までは手で書いてますが、それ以外は全て自動生成です。
またRustはmodels.rsとmodels以下の3つのファイルは自動生成です。
今回時間が無くてできませんでしたが、main.rsのasync fn executeは自動生成できるはずです。

{
  "meta": {
    "version": "1.0.0",
    "lname": "テスト",
    "pname": "test"
  },
  "includes": ["domain.erd.json"],
  "templates": {
    "interface": [
      {
        "template": "erd/stored_sql.ejs",
        "file": "sql/${meta.pname}_${pname}.sql"
      }
    ],
    "struct": [
      {
        "template": "erd/model_rs.ejs",
        "file": "src/models/${meta.pname}_${pname}.rs"
      },
      {
        "template": "erd/model_mod_rs.ejs",
        "file": "src/models.rs"
      }
    ]
  },
  "interfaces": [
    {
      "lname": "ユーザー一覧",
      "pname": "list_user",
      "input": "ユーザー一覧インプット",
      "output": "ユーザー一覧アウトプット"
    }
  ],
  "structs": [
    {
      "lname": "ユーザー一覧インプット",
      "pname": "user_list_input",
      "parameters": [
        {
          "domain": "ID"
        },
        {
          "domain": "名前"
        }
      ]
    },
    {
      "lname": "ユーザー一覧アウトプット",
      "pname": "user_list_output",
      "parameters": [
        {
          "domain": "ID",
          "pname": "user"
        },
        {
          "domain": "名前"
        },
        {
          "type": "ホビー[]",
          "notNull": true
        }
      ]
    },
    {
      "lname": "ホビー",
      "pname": "hobby",
      "parameters": [
        {
          "domain": "ID"
        },
        {
          "domain": "名前"
        },
        {
          "domain": "区分",
          "lname": "ホビー",
          "pname": "hobby"
        }
      ]
    }
  ]
}