Rust入門(6) 簡単なツールを作ってみる

Rust 入門の続き。

仕事でテキストデータをRDBにインポートする必要があったので、Rustの勉強がてらテキストファイルを読み込んで、SQLのINSERT文を出力するプログラムを書いてみることにした。

テキストデータは単純なTSV形式のファイルで改行コードはLF。
対象となるテキストデータは複数存在するので、インポート先のテーブルの情報は pg_dump コマンドで出力したもので以下のような形式。今回対象となるシステムでは数値型と文字列型しかデータ型は利用していないので、インポートするさいにクォートが必要なのは文字列型のみとなる。

CREATE TABLE m_master (
    category integer NOT NULL,
    id smallint NOT NULL,
    name character varying(100),
    order_no smallint DEFAULT 0,
    del_flg smallint DEFAULT 0,
    in_time character varying(12) DEFAULT to_char(now(), 'YYYYMMDDHH24MI'::text),
    in_staff integer,
    up_time character varying(12) DEFAULT to_char(now(), 'YYYYMMDDHH24MI'::text),
    up_staff integer
);

とりあえず、いろいろ試行錯誤しながらコーディングして、とりあえず動くとこまで持っていったのが次のコード。
CRERATE TABLE のSQLを読み込んで構造体を返す部分は他にも転用できそうな気がするのでモジュールとして作成した。

// table_def.rs

extern crate regex;
use regex::Regex;

#[derive(Debug)]
pub struct FieldDef {
    pub field_name: String,
    pub field_type: String,
    pub options: String
}

#[derive(Debug)]
pub struct TableDef {
    pub table_name: String,
    pub fields: Vec<FieldDef>
}

fn parse_fields(s: &str) -> Vec<FieldDef> {
    let mut vec:Vec<FieldDef> = Vec::new();
    for line in s.split("\n") {
        let re = Regex::new("(\\S+)\\s+(\\S+)(.*),").unwrap();
        if let Some(m) = re.captures(line) {
            vec.push(FieldDef {
                field_name: (&m[1]).to_string(),
                field_type: (&m[2]).to_string(),
                options: (&m[3]).to_string()
            });
        }
    }
    vec
}

impl TableDef {
    pub fn parse(sql: &str) -> Result<TableDef, &str> {
        let re = Regex::new("(?s)CREATE\\s+TABLE\\s+(\\S+)\\s+\\(\\s*(.+)\\s*\\);").unwrap();
        match re.captures(sql) {
            Some(m) => Ok(TableDef { 
                table_name: (&m[1]).to_string(),
                fields: parse_fields(&m[2])
            }),
            None => Err("Invalid Data")
        }
    }
}

SQLをパースする部分は定形なので正規表現で抜き出している。文字列には改行コードを含むため(?s)を指定して改行コードをく白文字にマッチするようにしてる。フィールド定義については一度まるごと抜き出したのを改行コードで分割したあと、さらに正規表現でフィールド名とデータ型などを抜き出すようにしている。ここら辺、正規表現でもっと頑張れば簡潔に書けそうな気がする。フィールド名以降の属性のパースについては、文字列型かそれ以外の判定さえできれば問題ないので非常に杜撰なものとなってる。

で、メインの処理はこんな感じに。

// main.rs

use std::env;
use std::fs;
use std::io::BufReader;
use std::io::BufRead;
use std::io::BufWriter;
use std::io::stdout;
use std::io::Write;
use std::result::Result;
mod table_def;
use table_def::TableDef;

fn load_table_def(s: &str) -> Result<TableDef, &str> {
    if let Ok(content) = fs::read_to_string(s) {
        if let Ok(t) = TableDef::parse(&content) {
            return Ok(t);
        }    
    }
    Err("Read Error")
}

fn output(table_def: &TableDef, data_path: &str) {
    if let Ok(file) = std::fs::File::open(data_path) {
        let mut fb = BufReader::new(file);
        let mut buff = String::new();
        let out = stdout();
        let mut out = BufWriter::new(out.lock());

        let mut fields = table_def.fields.iter().map(|f| &f.field_name).fold(String::new(), |s, n| s + n + ",");
        fields.pop();
        writeln!(out, "INSERT INTO {} ({}) VALUES ", &table_def.table_name, &fields).unwrap();
        let mut sw = false;
        loop {
            match fb.read_line(&mut buff) {
                Ok(n) => {
                    if n > 0 {
                        buff.pop();
                        if sw {
                            write!(out, ",\n(").unwrap();
                        } else {
                            write!(out, "(").unwrap();
                            sw = true;
                        }
                        for (i, v) in (&buff).split("\t").enumerate() {                            
                            let value = if table_def.fields.len() > i && table_def.fields[i].field_type == "character" {
                                format!("'{}'", v)
                            } else {
                                if v != "" { v.to_string() } else { "NULL".to_string() }
                            };
                            if i == 0 {
                                write!(out, "{}", value).unwrap();
                            } else {
                                write!(out, ",{}", value).unwrap();
                            }
                        }
                        write!(out, ")").unwrap();
                        buff.clear();
                    } else {
                        break;
                    }
                },
                _ => break
            }
        }
        writeln!(out, ";").unwrap();
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let table_def = load_table_def(&args[1]);
    match table_def {
        Ok(t) => {
            output(&t, &args[2]);
        },
        Err(err) => {
            println!("{:?}", err);
        }
    }
}

一番、苦労したのは項目定義の構造体のVecから項目名を取り出して連結した一つの文字列を取得するとこで、いろいろ試行錯誤た結果、次のような感じになった。

let mut fields = table_def.fields.iter().map(|f| &f.field_name).fold(String::new(), |s, n| s + n + ",");
fields.pop();

配列であれば connect() というメソッドがあって&strの連結ができるようだが、Vec<String> を簡単に結合するメソッドのようなものはなさそう。仕方がないので、map() で構造体からフィールド名を保持する String の参照を抜き出して、fold() で Stringに連結するようにした。末尾に余計な","が付くので最後に fields.pop() で取り除いている。

配列、スライス、Vecの違い、iter()とinto_iter()の違い、&strとStringの違いなど、いろいろまだキチンと理解していないのでなかなかコンパイルできるコードにたどり着けず苦労した。もっとエレガントというか、良くあるコードなのでイディオムが存在してそうな気がする。

あと、これはいまだに訳が判っていないトコがあって、テキストファイルから読み込んだテーブル定義のStringをパース処理するメソッドに喰わせた結果の戻り地の処理で、以下のような記述でコンパイルが通るようなったのだが、

fn load_table_def(s: &str) -> Result<TableDef, &str> {
    if let Ok(content) = fs::read_to_string(s) {
        if let Ok(t) = TableDef::parse(&content) {
            return Ok(t);
        }    
    }
    Err("Read Error")
}

なぜ、↓だとコンパイルが通らないのかが判らない。

fn load_table_def(s: &str) -> Result<TableDef, &str> {
    if let Ok(content) = fs::read_to_string(s) {
        TableDef::parse(&content)
    } else {
      Err("Read Error")
    }
}

returns a value referencing data owned by the current functionrustc(E0515)

というエラーが発生するのだが、なぜ戻り値をそのまま返すのが駄目で、一度、Resutから値を取り出して再びResultで包むとエラーにならないのか、その理屈が判らない。

うーん、Rust の道はなかなかに険しい。