node.jsでテキストファイルを読み込んでMySQLにデータを突っ込む

はじめに

MySQLで吐き出したテキストデータを別のMySQLに突っ込む必要があったので作成。
node.jsで何も考えずにI/O関連のコード書くと酷いことになるといういい例w

• insert.js

var lazy = require('lazy');
var fs = require('fs');
var opts = require('opts');
var mysql = require('mysql');
var _ = require('underscore');
var async = require('async');

//全体で使う変数
var t = {
    def: {}, //テーブル定義
    name: ''
};

//コマンドライン引数の設定
opts.parse([
    {
		short: 't',
		long: 'table',
		description: 'table name',
		value: true,
		require: true
    },{
		short: 'f',
		long: 'file',
		description: 'data file',
		value: true,
		require: true
    },{
		short: 'd',
		long: 'def',
		description: 'table definision (json)',
		value: true,
		require: true
    }
]);

//コマンドライン引数の取得
t.name = opts.get('table');
var dataFile = opts.get('file');
var tDef = opts.get('def');
var tableInfo = null;

console.log('talbe: ' + t.name + ', data: ' + dataFile + ', table def: ' + tDef);

//MySQLの準備
var client = mysql.createClient({
    user: 'user_name',
    password: 'password'
});
client.query('use sample_db');

//前処理 テーブル作成など
var prep = function(table, importDataFunc){
    fs.readFile('./' + tDef, function(err, data){
	if(err){
	    console.log(err);
	}
	tableInfo = JSON.parse(data.toString('utf8'));
	var tableDef = tableInfo.def;

	var defs = [];
	_.each(tableDef, function(type, name){
	    defs.push(name + ' ' + type);
	});
	t.def = tableDef;

	var query = 'CREATE TABLE ' + table + '(' + defs.join(', ') + ');';	
	console.log(query);

	var createTable = function(datafile){
	    //テーブル作成
	    client.query(query, function(){
			console.log('table: ' + table + ' created!');
			importDataFunc(dataFile, createIndexes);
	    });
	};

	//テーブルが既に存在していたら削除する
	try{
	    client.query('DROP TABLE ' + table, function(err){
	    	createTable(dataFile);
	    });
	}catch(e){
	    console.log('already, ' + table +' is exist.');
	    createTable(dataFile);
	}
    });
};

//インデックス作成
var createIndexes = function(tableName, indexes){
    var queries = [];
    _.each(indexes, function(val, key){
    	var query = 'ALTER TABLE ' + tableName + ' ADD INDEX(' + val + ')';
		console.log(query);
		queries.push(query);
    });
    async.forEach(queries, function(query, cb){
		client.query(query, function(){
		    cb();
		});
    }, function(err){
		if(err){
		    console.log(err);
		}else{
		    console.log('create indexes');
		    client.query('COMMIT', function(){
			    process.exit();
		    });
		}
    });
};

//ファイルを読み込んでパースしてDBに書き込む
var queries = [];
var importData = function(table, createIndexesFunc){
    console.log('start import');
    var cols = [];
    _.each(t.def, function(type, name){
    	cols.push(name);
    });
    
    var queryBase = 'INSERT INTO ' + t.name + '(' + cols.join(', ') + ') VALUES (';

    var rs = fs.createReadStream(dataFile, {bufferSize: 256 * 1024});
    //forEachは非同期。rsのendイベントで完了を検知する
    rs.on('end', function(){
		//最初の一行はカラム定義とするので除去
		queries.shift();
	
		console.log('create ' + queries.length + ' queries');

		async.forEach(queries, function(query, cb){
		    client.query(query, function(){
				cb();
		    });
		}, function(err){
		    if(err){
		    	console.log(err);
		    } else {
				console.log('fin: ' + queries.length + ' queries.');
				createIndexesFunc(t.name, tableInfo.indexes);
		    }
		});
    });

    //ファイル読み込み開始
    new lazy(rs).lines.forEach(function(line){
		var items = line.toString().split('\t');
		_.each(items, function(val, i){
		    if(!val.match(/^[0-9]+/ && val !== 'NULL')){
			items[i] = '"' + val + '"';
		    }
		});
		var query = queryBase + items.join(', ') + ')';
		queries.push(query);
    });
};


client.query('SET autocommit=0', function(){
	client.query("START TRANSACTION", function(){
		prep(t.name, importData);
	});
});

以下みたいな感じでテーブル定義を指定する。

• def.json

{
	"def": {
	       "user_id": "int unsigned",
	       "age": "int unsigned",
	       "name": "varchar(255)"
	},
	"indexes": [
	       "user_id"
	]
}

対象としているデータの形式は以下。

• data.txt

user_id	age	name
10001	26	Taro
10002	18	二郎

実行方法

node insert.js -t 【テーブル名】 -d def.json -f data.txt

実践Node.js プログラミング (Programmer's SELECTION)

サーバサイドJavaScript　Node.js入門 (アスキー書籍)