Rustのfactorybot的なライブラリ「factori」とsqlxでモックデータを生成する【1】

にあえん

October 7, 2023

pythonの単体テストでよく使っていたfactoryboyというツールがあるのですが、これが本当に使いやすい!

factory_boy — Factory Boy stable documentation

これとfixtureを使えば、大抵のテストケースに必要なデータを引数から取得することができます。

で、これと同じようなことをRustでもやりたいな〜と思っていたんですが、「factori」というライブラリがそれに相当するということで、備忘録を兼ねて触った感想などをここに書いていこうと思います。

使い方

とりあえずCargoパッケージに追加します。

[dependencies]
factori = { version = "1.1.0" }
fake = "2.8.0"
uuid = { version = "1.4.1", features = [ "serde", "v4" ] }

使い方はメチャクチャ簡単で、定義された構造体にfactori!マクロを定義してあげれば、create!マクロから呼び出すことができます。

fakeというライブラリと組み合わせてテストデータを作る実装は以下のようになります。

use factori::{factori, create};
use fake::Fake;
use fake::faker::name::raw::*;
use fake::locales::JA_JP;
use uuid::Uuid;

#[derive(Debug)]
pub struct User {
    id: Uuid,
    name: String,
}

factori!(User, {
    default {
        id = Uuid::new_v4(),
        name = Name(JA_JP).fake::<String>(),
    }
});



fn main() {
    // Userを10人作成する
    let user: Vec<User> = [0; 10].iter().map(|_| create!(User)).collect();
    println!("{:?}", user);
}

これだけでUserをちゃちゃっと作る実装ができました。

ではここにOrganizationという関連性を追加します。

use factori::{factori, create};
use fake::Fake;
use fake::faker::company::raw::*;
use fake::faker::name::raw::*;
use fake::locales::JA_JP;
use uuid::Uuid;

#[derive(Debug)]
struct Organization {
    id: Uuid,
    name: String,
}

// ...中略...

factori!(Organization, {
    default {
        id = Uuid::new_v4(),
        name = CompanyName(JA_JP).fake::<String>(),
    }
});

factori!(User, {
    default {
        id = Uuid::new_v4(),
        name = Name(JA_JP).fake::<String>(),
        org = create!(Organization),
    }
});

こんな感じで再度実行すると、これでUserと関連するOrganizationが10個できました!

実行するとこんな感じです。

    Finished dev [unoptimized + debuginfo] target(s) in 3.54s
     Running `target/debug/factori_1`
[User { id: 47d364ab-e24c-4af9-bb08-05e3a0d0cf68, name: "田中 乃愛", org: Organization { id: d314ef7f-ead3-4366-80cb-223e230ea870, name: "(有) 阿部" } }, User { id: ec17304e-5268-4108-8dc2-b8c7d87d648e, name: "和田 華", org: Organization { id: c082704a-8cd1-4d04-b827-d8373f8313ac, name: "藤田 (有)" } }, User { id: cf77936e-2d6b-4ff9-9153-637f0b3ac025, name: "加藤 莉央", org: Organization { id: 502b0ffb-cfc1-422b-baba-ddc420d224be, name: "有限会社 長谷川" } }, User { id: e6b7f6b7-447e-4e87-b777-1b7f45ad6d43, name: "柴田 湊翔", org: Organization { id: 5fc0d3f7-7e63-4963-a9c1-bcdb181cd8d7, name: "佐藤 (株)" } }, User { id: 412ab765-407f-4f16-961c-0fec804f4837, name: "久保 優月", org: Organization { id: 67344bef-8160-47e2-994d-e451d0cb5cbd, name: "西村 有限会社" } }, User { id: a0804301-2500-43eb-9a96-8cd6b217d682, name: "柴田 琴葉", org: Organization { id: e4cf4825-d9d0-426e-a4a1-a9a1d5b5c2ac, name: "高橋 株式会社" } }, User { id: 072aa05a-5229-488f-9cd7-de44d6bdc3df, name: "渡辺 蒼太", org: Organization { id: 74be9444-640e-4040-acb6-d6abe0475af3, name: "有限会社 岡田" } }, User { id: 3c87567e-17ee-416a-a9ed-969c033031d3, name: "近藤 心陽", org: Organization { id: 009efa86-2307-4d79-9a80-7981135d7160, name: "中島 (株)" } }, User { id: c96f7080-c961-4af4-9aff-54802b2c8c78, name: "上田 大知", org: Organization { id: 6a779d1b-c5af-4bde-b75e-9566d1212f7e, name: "松井 (株)" } }, User { id: 68800ba6-fc45-4cd7-a86d-10a6d663bdc0, name: "原田 新", org: Organization { id: 44f46803-8b9a-4f71-bfda-8ca516346898, name: "阿部 (有)" } }]

sqlxと組み合わせる

さて、ここからが本題です。

create!でのインスタンス化に紐付けてsqlxを使用したインサート処理を行えるようにしたいです。

まぁとりあえずテーブルを作っておきます。

以下のようなdocker-composeのyamlを用意しておきます。

version: '3'

services:
  db:
    image: postgres:15
    container_name: postgres
    ports:
        - 5432:5432
    environment:
        - POSTGRES_PASSWORD=example

つづけてテーブル作成をしていきます。

最近知ったpsqldefを使っていきましょう!

(紹介記事はこちら)

CREATE EXTENSION "uuid-ossp";

CREATE TABLE public.organization (
    id uuid PRIMARY KEY DEFAULT uuid_generate_v4(),
    name varchar(255)
);

CREATE TABLE public.user (
    id uuid PRIMARY KEY DEFAULT uuid_generate_v4(),
    name varchar(255),
    org_id uuid NOT NULL,
    CONSTRAINT fk_user_org_id FOREIGN KEY (org_id) REFERENCES organization (id)
);

こんな感じで定義して

psqldef -h localhost -U postgres --password=example [DB名] < schema.sql

これで適用できました。

あとはsqlxの実行のために依存関係も追加しておきましょう。

[dependencies]
# ...中略...
sqlx = { version = "0.7", features = [ "runtime-tokio", "postgres", "uuid" ] }
tokio = { version = "1.32.0", features = [ "full" ] }

あとは作成したUserをDBにインサートする関数を書きます。

use factori::{factori, create};
use fake::Fake;
use fake::faker::company::raw::*;
use fake::faker::name::raw::*;
use fake::locales::JA_JP;
use sqlx::postgres::{PgPool, PgPoolOptions};
use sqlx::{Postgres, QueryBuilder};
use uuid::Uuid;

#[derive(Debug)]
pub struct Organization {
    id: Uuid,
    name: String,
}

#[derive(Debug)]
pub struct User {
    id: Uuid,
    name: String,
    org: Organization,
}

factori!(Organization, {
    default {
        id = Uuid::new_v4(),
        name = CompanyName(JA_JP).fake::<String>(),
    }
});

factori!(User, {
    default {
        id = Uuid::new_v4(),
        name = Name(JA_JP).fake::<String>(),
        org = create!(Organization),
    }
});

async fn insert_user(pool: &PgPool, user: &User) {
    let cnt: i64 = sqlx::query_scalar("SELECT count(*) FROM organization WHERE id = $1")
        .bind(user.org.id)
        .fetch_one(pool)
        .await.unwrap();

    if cnt == 0 {
        insert_organization(pool, &user.org).await;
    }

    let mut builder: QueryBuilder<Postgres> = QueryBuilder::new(r#"
    INSERT INTO public.user (
        id,
        name,
        org_id
    ) VALUES (
    "#);
    let mut separated = builder.separated(", ");
    separated.push_bind(user.id);
    separated.push_bind(user.name.clone());
    separated.push_bind(user.org.id);
    separated.push_unseparated(")");

    builder.build().execute(pool).await.unwrap();
}

async fn insert_organization(pool: &PgPool, organization: &Organization) {
    let mut builder: QueryBuilder<Postgres> = QueryBuilder::new(r#"
    INSERT INTO organization (
        id,
        name
    )
    VALUES (
    "#);
    let mut separated = builder.separated(", ");
    separated.push_bind(organization.id);
    separated.push_bind(organization.name.clone());
    separated.push_unseparated(")");
    builder.build().execute(pool)
        .await.unwrap();
}

#[tokio::main]
async fn main() {
    let pool = PgPoolOptions::new()
        .connect("postgres://postgres:example@localhost:5432/psqldef?schema=public")
        .await.unwrap();

    // Userを10人作成する
    let users: Vec<User> = [0; 10].iter().map(|_| create!(User)).collect();
    for user in users {
        insert_user(&pool, &user).await;
    }
}

これでできました。

ただこのような定型的な処理はできればもっと楽に実装しちゃいたいですね。

まとめ

factoriとSqlXを組み合わせた処理を実装しました。

次回はマクロを使ってcreate!のデコレータを実装し、より汎用的な形に直してみようと思います!

ということでマクロを1から学んでから次回書きます。