PythonでCSVをDBにインポート(with Flask-Script and SQLAlchemy)

Flask + SQLAlchemyを使っていて、フィクスチャ、サンプルデータ、マスタデータ等をDBに入れたいが、何かしらのファイル形式（今回はCSV）で書いたものをDBへとインポートする仕組みを作りたい

LOAD DATA INFILE 使えばええやんって話もあるけど、インポートの時点である程度のバリデートもできる(実際はデータを入れてから流すケースが多そうではあるが)ので、一考の価値はありかと

例えば、下記のようなテーブル(Message)に対して

mysql> desc message;
+------------+-----------+------+-----+---------+----------------+
| Field      | Type      | Null | Key | Default | Extra          |
+------------+-----------+------+-----+---------+----------------+
| message_id | int(11)   | NO   | PRI | NULL    | auto_increment |
| author     | text      | YES  |     | NULL    |                |
| text       | text      | YES  |     | NULL    |                |
+------------+-----------+------+-----+---------+----------------+
3 rows in set (0.00 sec)

以下のようなCSV(message.csv)を用意して、データを取り込むようにする

message_id,author,text
1,Tom,This is a sample message
2,Ken,Oops!

Flask-Scriptとは

Flaskで外部スクリプトを簡単に書くためのextention。

Djangoでいうmanage.pyのようなものが実現できる

Flask-Script — Flask-Script 0.4.0 documentation

@manager.command
def hello():
    "Just say hello"
    print "hello"

python manage.py hello
> hello

↑は公式の引用だが、こういった形で使用できる

CSVをDBにインポートする

CSVに対してmodel(SQLAlchemy)を動的にロードしDBに突っ込む

将来のmodelの追加・変更を考えるとインポートのロジック自体にmodelの情報を持たせたくない

ここではCSVのファイル名をクラス名、CSVの1行目をクラスの各属性の情報として役立てるようにしている

from flask_script import Manager
from skeleton import app, db
import csv
import importlib
import os


manager = Manager(app)


@manager.command
def import_csv():
    fixtures_dir = app.config['FIXTURES_DIR']
    models = importlib.import_module('skeleton.models')

    for file_name in os.listdir(fixtures_dir):
        class_name = file_name.replace(".csv", "").capitalize()
        Klass = getattr(models, class_name)
        with open(fixtures_dir + '/' + file_name, encoding='utf-8') as csv_file:
            reader = csv.reader(csv_file, delimiter=',')
            header = next(reader)
            for row in reader:
                klass = Klass()
                for i in range(len(header)):
                    setattr(klass, header[i], row[i])
                db.session.add(klass)
            db.session.commit()


if __name__ == "__main__":
    manager.run()

件数多ければbulk insertにする必要がある

Pythonメモ

• importlib : モジュールの動的ロードしてくれるライブラリ
• getattr : object の指名された属性の値を返す組み込み関数。getattr(x, ‘foobar’) は x.foobar と等価
• setattr : 値を属性に関連付ける組み込み関数。setattr(x, ‘foobar’, 123) は x.foobar = 123 と等価

31.5. importlib — import の実装 — Python 3.6.0 ドキュメント
2. 組み込み関数 — Python 3.6.0 ドキュメント

modelやdb周りの簡素な設定も含んだコードはこれ