ただただ、Rustを書いてみたい

Tech系のブログや記事を見てると、Rustを目にする機会が多くなってきました。書いてみたくなったので、CSVをインポートしてDBに挿入する処理をRustで書くというお題目を設定しました。

前提

Railsで、Articleというテーブルに対して、約2万行のCSVファイルを読み込んで挿入するというタスクがあります。
すでにPostgreslでDBがあり、Articleというテーブルも存在しています。
Article は以下のschemaです。


create_table "articles", force: :cascade do |t|
    t.string "title" # タイトル
    t.string "body" # 本文
end

今回はその処理をRustで書いてみます。

CSVインポートして、ただ表示する

まずはCSVインポートですが、Cargoのパッケージがあるのでそれを使います。

Cargo.toml

[dependencies]
csv = "1.1"

次に、CSVを読み込んで表示する処理です。（https://docs.rs/csv/ のチュートリアルの通り）

src/bin/show_csv.rs

extern crate rust;

use self::rust::*;
use std::io;
use std::process;

fn read() -> Result<(), Box<dyn Error>> {
    let mut rdr = csv::Reader::from_reader(io::stdin());
    for result in rdr.records() {
        let record = result?;
        let title = &record[0];
        let body = &record[1];
        println!("title: {}", &title);
        println!("body: {}", &body);
    }
    Ok(())
}

fn main() {
    if let Err(err) = read() {
        println!("error running read: {}", err);
        process::exit(1);
    }
}

実行

$ cargo run --bin show_csv < sample.csv

title: AAAAAA
body: BBBBB

Dieselを使ってDBへアクセスする

次に、DBへアクセスします。
RustからDBへアクセスする方法はいろいろありますが、今回は少しでもRailsに近い感じにするため、ORMの Diesel を使いました。

以下はしばらく、 https://diesel.rs/guides/getting-started/ に書かれている通りの設定が続きます。

Cargo.toml

[dependencies]
diesel = { version = "1.0.0", features = ["postgres"] }

データベースへの接続設定を環境変数にセットし、セットアップします。

$ echo DATABASE_URL=postgres://username:password@localhost/yama_sample > .env
$ diesel setup

使用するテーブルのmigrateを行います。

$ diesel migration generate create_articles

するとmigrationsフォルダ以下に、テーブル作成と削除用のSQLファイルが生成されるので、適宜設定します。

migrations/2019-12-05-213319_create_articles/up.sql

CREATE TABLE articles (
  id SERIAL PRIMARY KEY,
  title TEXT,
  body TEXT
)

Modelを定義します。

src/models.rs

use super::schema::articles;

#[derive(Queryable)]
pub struct Article {
    pub id: i64,
    pub title: String,
    pub body: String,
}

チュートリアルでは、 id: i32 になっていましたが、Railsで作ったテーブルはBigIntだったため、 i64 へ変更しています。

Schemaを定義します。

src/schema.rs

table! {
    articles (id) {
        id -> BigInt,
        title -> Text,
        body -> Text,
    }
}

こちらもidを IntegerからBigIntへ変更しています。

これで最低限の設定は完了です。

CSVで読み込んだデータをテーブルに挿入

いよいよ本丸です。
まず、Insert用の構造体を定義します。

src/models.rs

#[derive(Insertable)]
#[table_name="articles"]
pub struct NewArticle<'a> {
    pub title: &'a str,
    pub body: &'a str,
}

次に、Insert用の関数を定義します。

src/lib.rs

pub fn create_article<'a>(conn: &PgConnection, title: &'a str, body: &'a str) -> Article {
    use schema::articles;

    let new_article = NewArticle {
        title: title,
        body: body,
    };

    diesel::insert_into(articles::table)
        .values(&new_article)
        .get_result(conn)
        .expect("Error saving new article")
}

最後に、CSVを読み込み、そのデータをDBへInsertするロジックを追加します。

src/bin/update_article.rs


extern crate rust;
extern crate diesel;

use self::rust::*;
use std::error::Error;
use std::io;
use std::process;

fn create() -> Result<(), Box<dyn Error>> {
    let connection = establish_connection();
    let mut rdr = csv::Reader::from_reader(io::stdin());
    for result in rdr.records() {
        let record = result?;
        let title = &record[0];
        let body = &record[1];

        let article = create_article(&connection, &title, &body);
    }
    Ok(())
}

fn main() {
    if let Err(err) = create() {
        println!("error running create: {}", err);
        process::exit(1);
    }
}

実行すると、CSVのデータが1行ずつテーブルにInsertされます。

Rustに任せるのも一手

詳細に検証すると数値がずれるとは思いますが、速報値として、

同じ処理をRailsのままでやる: 2分7秒
Rust: 30秒

という結果でした。もちろん、Railsの方でbulk insertを使うなどしてチューニングすることはできますし、単純に比較するものでもないと思います。
今回は、Rustでやったらどうなるかという好奇心から実装しましたが、Dieselという強力なORMもあることですし、この処理をRustに任せるというのは選択肢としてありうるのかなと思います。

次は @kawano-fusic の「S3×Lambda×Cloudwatch Eventsで、バッチ処理の監視機構を簡単に導入する」です。
はりきってどうぞ！