RustのDieselでPostgreSQLサーバにクエリを投げる

概要

ポスグレにクエリを投げるツールを、DieselというORM(Object-relational mapping)を使ってRustで実装してみる。
より素朴なライブラリとしては、下記のpostgresクレートがあり、素直にSQLを書いて実装できるので、最初はDieselよりはこっちが取っつき易いのかもしれない。

• [参考] postgresクレート（今回これは使わない）
  • https://docs.rs/postgres/0.19.1/postgres/

が、Dieselを使うとクエリをRustの思想にしたがって「メソッドチェーン（＝関数をつないで処理する書き方）」で表現できたり、型安全が保証されたり、それ以外でも多くの恩恵がある模様。
今回はRustのDieselクレートを使って、予め作成したポスグレのDBに接続してクエリを投げて結果を取得してみる。
ちなみにDieselでDBそのものを作成することも下記サイトの方法等でできるが、今回の記事では対象外とする。

• dieselのサンプル(sqlite)を動かしてみた
  • https://qiita.com/techno-tanoC/items/4ca564bd1ea8db409496

準備

• サーバ（PostgreSQL）
  • PostgreSQLのサーバを準備。（セットアップ手順は割愛）
  • バージョンは11.5
• クライアント（Rust）
  • RustのDockerイメージを利用
  • バージョン
    • rustc : 1.51.0
    • cargo : 1.51.0

サーバ準備詳細（データベース作成）

今回接続するDBとテーブルを作成する

まずはDBの作成。sampleという名前にする。

$ psql
postgres=# CREATE DATABASE sample;
CREATE DATABASE

一旦psqlからログアウトして、DB名sampleを指定して入り直し、customerというテーブルを作成。
nameが主キーで、ageはNULLを許容するようにする。

$ psql -d sample
sample=# CREATE TABLE customer (
    name TEXT PRIMARY KEY NOT NULL,
    age INTEGER,
    reg_date DATE NOT NULL 
);
CREATE TABLE

テーブルに３つの行を追加。

sample=# INSERT INTO customer
     (name, age, reg_date)
 VALUES
     ('田中 一郎', 34, DATE '2021-03-13'),
     ('佐藤 二郎', NULL, DATE '2021-03-15'),
     ('鈴木 三郎', 43, DATE '2021-03-14');
INSERT 0 3

行が追加されていることを念のため確認。

sample=# SELECT * FROM customer;
   name    | age |  reg_date  
-----------+-----+------------
 田中 一郎 |  34 | 2021-03-13
 佐藤 二郎 |     | 2021-03-15
 鈴木 三郎 |  43 | 2021-03-14

クライアント準備詳細（RustのDockerイメージ作成＆コンテナ起動）

RustのDockerイメージを以下のDockerfileを使って作成。Rustを使うユーザを追加し、コンテナ実行時に追加したユーザでログインできる状態にしておく。

Dockerfile

FROM rust

# ユーザ追加
ARG USERNAME=username
RUN adduser ${USERNAME} && \
    apt -y update
USER ${USERNAME}

上記Dockerfileをもとに、Dockerイメージを作成。

$ docker build -t myrust .

コンテナを起動してログイン。

$ docker run -dit --name myrust myrust
$ docker exec -it myrust /bin/bash
username@2b473184441c:/$

以後は全て、「myrustコンテナ」にログインした状態で作業をする。

Dieselのインストール、設定

• 参考サイト
  • https://diesel.rs/guides/getting-started.html

Diesel CLIのインストール

以下コマンドでインストール。今回はポスグレの機能しか使わないので、オプションでその旨を指定。
※環境によってはlpqがなくてエラーになることがあるので、Ubuntu系OSの場合はsudo apt -y install libpq-devで事前にインストールが必要かも。

$ cargo install diesel_cli --no-default-features --features postgres

Cargoプロジェクトの作成と設定

適当なディレクトリでプロジェクトを作成。名前はpsql_clientとする。
以後、作成されたpsql_clientディレクトリに移動して各種設定を行う。

$ cargo new psql_client
     Created binary (application) `psql_client` package
$ cd psql_client

まず、Cargo.tomlの[dependencies]セクションに必要なパッケージを記載。今回はポスグレの機能と、Date型のデータを扱うため、featuresに["postgres", "chrono"]の２つを指定している。

Cargo.toml

[package]
name = "psql_client"
version = "0.1.0"
authors = ["auther"]
edition = "2018"

[dependencies]
diesel = { version = "1.4.6", features = ["postgres", "chrono"] }
dotenv = "0.15.0"
chrono = "0.4.19"

続いて、ポスグレサーバのURIを記載した.envファイルを作成する。参考サイトにも記載があるが、例えば次のように設定すればOK。URIの最後はDB名を指定するため、今回はsampleと書く。

$ echo DATABASE_URL=postgres://[ユーザ名]:[パスワード]@[ホスト名]:[ポート番号]/sample > .env

DBのURIを設定後、以下コマンドを実行。

$ diesel setup
Creating migrations directory at: /path/psql_client/migrations

setupに成功すると、migrationsディレクトリができる。テーブルを新たに作る場合は、ここにSQLのCREATE文を書けばいいようだが、今回は既にテーブルを作成済なのでスキップする。
もう一つ、diesel.tomlというファイルが生成されており、以下のようにデータベースのスキーマの定義をsrc/schema.rsに記載せよ、という記述がある。ここは次の「実装」で書いていく。

diesel.toml

[print_schema]
file = "src/schema.rs"

実装

今回のプロジェクトのディレクトリ構成は下記。
まず、src/schema.rsにテーブル情報のマクロを記載し、続いて本体のsrc/main.rsを実装していく。

.
├── Cargo.toml
├── diesel.toml
└── src
    ├── main.rs
    └── schema.rs

src/schema.rsの作成

dieselのprint-schemaというコマンドでスキーマ内にあるテーブル情報のマクロを作ってくれる。これをそのままsrc/schema.rsへリダイレクトすればOK。

$ diesel print-schema > src/schema.rs

src/schema.rsの中身は以下

src/schema.rs

table! {
  customer (name) {
    name -> Text,
    age -> Nullable<Int4>,
    reg_date -> Date,
  }
}

テーブル名customerの次の括弧部分に記載される変数（本例ではname）は、テーブルの主キー。またageは値がNULLになる可能性があるので、Nullable<T>型の構造体で定義されていることに注意。
今回は半分手動でsrc/schema.rsを作成したが、Rust + MySQL + Diesel（1）にも記載があるように、diesel migration runでテーブルを作成した場合はsrc/schema.rsは自動で作成される模様。

src/main.rsの作成

とりあえず以下のように書くと、customerテーブルの各行の情報を取得できる。

src/main.rs

#[macro_use]
extern crate diesel;
extern crate dotenv;

use std::env;
use diesel::prelude::*;
use diesel::pg::PgConnection;
use dotenv::dotenv;

// customerテーブルの日付型カラム(reg_date)を扱うのに必要
use chrono::NaiveDate;

// "src/schema.rs"で定義したマクロを使えるようにする
mod schema;
use schema::customer::dsl::*;

// テーブルcustomerの各行の情報を格納する構造体
// （Web上のサンプルプロブラムでは別ソースmodels.rsに書かれていることが多い）
// ageはNULLを取ることがあり、Option型で受ける必要あり
#[derive(Queryable, Debug)]
struct Customer {
  name: String,
  age: Option<i32>,
  reg_date: NaiveDate,
}

/// DBの接続を確立
fn establish_connection() -> PgConnection {
  dotenv().ok();

  // .envファイルに定義された環境変数DATABASE_URLを取得してDBに接続する
  let database_url = env::var("DATABASE_URL")
    .expect("DATABASE_URL must be set");
  PgConnection::establish(&database_url)
    .expect(&format!("Error connecting to {}", database_url))
}

fn main() {
  // DBの接続を確立
  let conn = establish_connection();

  // customerテーブルの各情報を取得
  // SQLで「SELECT * FROM customer;」をやっているのと同じ
  let results = customer
    .load::<Customer>(&conn)
    .expect("Error loading customer");

  // 結果を表示
  for r in results {
    println!("{:?}", r);
  }
}

ビルドして実行

customerテーブルの各行の値を出力

$ cargo run
   Compiling psql_client v0.1.0 (/path/psql_client)
    Finished dev [unoptimized + debuginfo] target(s) in 0.97s
     Running `target/debug/psql_client`
Customer { name: "田中 一郎", age: Some(34), reg_date: 2021-03-13 }
Customer { name: "佐藤 二郎", age: None, reg_date: 2021-03-15 }
Customer { name: "鈴木 三郎", age: Some(43), reg_date: 2021-03-14 }

ageはNULLを取り得るので、Option型になっている。佐藤さんのageはNULLで、Customer構造体にはNoneが格納されていることがわかる。

クエリのバリエーション

main関数内でresultsを取得する部分はメソッドチェーンにより様々なクエリを指定可能。

昇順ソート（ORDER BY句相当）

  let results = customer
    .order(reg_date.asc()) // 追加
    .load::<Customer>(&conn)
    .expect("Error loading customer");

出力結果
Customer { name: "田中 一郎", age: Some(34), reg_date: 2021-03-13 }
Customer { name: "鈴木 三郎", age: Some(43), reg_date: 2021-03-14 }
Customer { name: "佐藤 二郎", age: None, reg_date: 2021-03-15 }

絞り込み（WHERE句相当）

  let results = customer
    .filter(name.eq("田中 一郎")) // 追加
    .load::<Customer>(&conn)
    .expect("Error loading customer");

出力結果
Customer { name: "田中 一郎", age: Some(34), reg_date: 2021-03-13 }

一部カラム（nameとreg_dateのみ）取得

SELECT name, reg_date FROM customer; と同じ動作

  let result = customer
    .select((name, reg_date))
    .get_results::<(String, NaiveDate)>(&conn)
    .unwrap();

出力結果
("田中 一郎", 2021-03-13)
("佐藤 二郎", 2021-03-15)
("鈴木 三郎", 2021-03-14)

※Dieselのドキュメントによると、get_resultsメソッドの処理はloadと同じらしい。

集約関数 COUNT

SELECT COUNT(*) FROM sample; と同じ動作

  let count_star = customer
		.count()
		.get_result::<i64>(&conn)
		.unwrap();

  assert_eq!(count_star, 3); // 結果が3になることの確認

SELECT COUNT(age) FROM sample; と同じ動作

  use diesel::dsl::*; // count関数を使うのに必要

  let count_age = customer
    .select(count(age))
    .get_result::<i64>(&conn)
    .unwrap();

  assert_eq!(count_age, 2);

集約関数 AVG

SELECT AVG(age) FROM sample; と同じ動作

  use diesel::dsl::*;          // avg関数を使うのにインポート必要
  use bigdecimal::BigDecimal;  // クエリ結果がBigDecimal型になるので必要

  let average = customer
    .select(avg(age))
    .get_result::<Option<BigDecimal>>(&conn)
    .unwrap()   // ageがNullable型なので、追加でunwrapが必要
    .unwrap();

  use std::str::FromStr;
  assert_eq!(average, BigDecimal::from_str(&"38.5").unwrap());

平均値だとBigDecimal型を返すようなので、Cargo.tomlの[dependencies]セクションで、bigdecimalクレートの追加が必要。また、dieselクレートのfeaturesで"numeric"も追記する必要あり。（これに気づくのに２時間近くかかった…）

Cargo.toml

(略)

[dependencies]
diesel = { version = "1.4.6", features = ["postgres", "chrono", "numeric"] } # "numeric"追加
dotenv = "0.15.0"
chrono = "0.4.19"
bigdecimal = "0.1.2" # 追加

また動作確認はしていないが、最大値(max)はBigInt型で返るようなので、平均(avg)と同様にクレート追加が必要になるかも。

SQL文を直書きする方法

sql_query関数を使うことで、SQL文を直書きして結果を取得することもできるが、Rust側で値を正しく受け取るための構造体定義など、事前準備が必要となる点に注意。

  let result = 
    diesel::sql_query("SELECT COUNT(name), AVG(age) FROM customer")
    .get_result::<CountAvg>(&conn) // CountAvg構造体は別途定義が必要
    .unwrap();

  println!("{:?}", result);

出力結果
CountAvg { count: 3, avg: Some(BigDecimal("38.5000000000000000")) }

この例の場合、src/schema.rsとsrc/main.rsに以下のように、クエリの結果を受け取るためのマクロと構造体を定義する。

src/schema.rs

// 以下を追記
table! {
  count_avg (count) {
    count -> Int8,
    avg -> Nullable<Numeric>,
  }
}

src/main.rs

// 以下を追記
use self::schema::count_avg;

// QueryableByNameトレイトの実装が必要
#[derive(QueryableByName, Debug)]
#[table_name = "count_avg"]
struct CountAvg {
  count: i64,
  avg: Option<bigdecimal::BigDecimal>,
}

直書きSQL文の中で変数指定

DBサーバがポスグレの場合は、以下のように$n(n=1,2,3...)を設定後、bindメソッドで$nの値を指定すればOK。

  let result = 
    diesel::sql_query("
      SELECT name, age, reg_date FROM customer
      WHERE age >= $1
      ")
    .bind::<diesel::sql_types::Integer, _>(40)
    .load::<CustomerByName>(&conn)
    .unwrap();

  println!("{:?}", result);

出力結果
[CustomerByName { name: "鈴木 三郎", age: Some(43), reg_date: 2021-03-14 }]

上で（何の説明もなしに）出てきたCustomerByName構造体は、事前に以下のように定義しておく。
sql_queryで直書きのクエリを使うため、QueryableByNameトレイトの実装が必要。

src/main.rs

// 以下を追記
#[derive(QueryableByName, Debug)]
#[table_name = "customer"]
struct CustomerByName {
  name: String,
  age: Option<i32>,
  reg_date: NaiveDate,
}

動作確認はしていないが、DBサーバがMySQLの場合は$nではなく、?と書けばいいはず。（というか、Webで見つかるほとんどのサンプルは?で書かれていた。なので筆者は最初、ポスグレも?で書くのかと思いこんで実装するもエラーで動かず、$nで書くこと気づくのにえらい時間を要してしまった…）