Rustでcsvを読み込みDBにインポートする

こちらは [Fusic Advent Calendar 2019 - Qiita](https://qiita.com/advent-calendar/2019/fusic) 6日目の記事です。
昨日は @tsukabo による [社内基幹システムのバージョンアップに挑んだ話](https://qiita.com/tsukabo/items/29a66ff49638e5b219a3) でした。
バージョンアップはエンジニアの宿命でもありますよね。私も直近、バージョンアップ案件が控えているので参考になります。

今回は、Railsの大容量のCSVインポート処理をRustで書いてみた件です。

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

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

```ruby

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/ のチュートリアルの通り）

```rust: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を定義します。

```rust: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を定義します。

```rust:src/schema.rs
table! {
    articles (id) {
        id -> BigInt,
        title -> Text,
        body -> Text,
    }
}
```
こちらもidを `Integer`から`BigInt`へ変更しています。


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

## CSVで読み込んだデータをテーブルに挿入
いよいよ本丸です。
まず、Insert用の構造体を定義します。

```rust:src/models.rs
#[derive(Insertable)]
#[table_name="articles"]
pub struct NewArticle<'a> {
    pub title: &'a str,
    pub body: &'a str,
}
```

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

```rust: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するロジックを追加します。

```rust: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で、バッチ処理の監視機構を簡単に導入する」です。
はりきってどうぞ！