PythonでCSV,TSVファイルをSQLiteにインポートする方法

はじめに

PythonでSQLiteのDBにCSVやTSVデータをインポートしたい！
そんな方のための、スクリプト一発でインポートする方法の解説です。

SQLiteとはなにか、どうやってセットアップするのか等は以下の記事を参考にしてください。
PythonとSQLiteで実現するローカルでお手軽SQL環境作成術

インポート方法

使い方はこの下の「解説」にて。

スクリプト

import_data_to_sqlite.py

import sqlite3
import csv
import os

#######↓ここのパラメータを変える↓#######
dbname = ''
target_table_name = ''
import_table_name = ''
is_create_table = 
is_header_skip = 
#####################################


#######↓インポート先のテーブルDDL↓#######
sql_script = """

"""
#######################################

class ImportSQLite():
    def __init__(self, dbname, target_table_name, import_data_name, is_create_table, is_header_skip=False, sql_create_table=None):
        """
        csvまたはtsvファイルをSQLiteへインポートする
        :param dbname: text 接続先DB名
        :param target_table_name: text インポート先となるDB上のテーブル名
        :param import_data_name: text インポートしたいデータ名
        :param is_create_table: boolean インポート先となるテーブルを作成するか否か
        :param is_header_skip: boolean インポートするデータのヘッダーを読み飛ばすか否か
        :param sql_create_table: text インポート先となるテーブルのDDL
        """
        self.dbname = dbname
        self.target_table_name = target_table_name
        self.import_data_name = import_data_name
        self.is_create_table = is_create_table
        self.is_header_skip = is_header_skip
        _, raw_delimiter = os.path.splitext(import_data_name)
        if raw_delimiter == '.csv':
            self.delimiter = ','
        elif raw_delimiter == '.tsv':
            self.delimiter = '\t'
        else:
            raise ValueError('Import file should be csv or tsv.')

        if is_create_table:
            if not sql_create_table:
                raise ValueError('It\'s necessary of sql to create table')
            else:
                self.sql_create_table = sql_create_table


    def read_import_file(self):
        with open(self.import_data_name, 'r', encoding='utf-8') as f:
            reader = csv.reader(f, delimiter=self.delimiter)
            if self.is_header_skip:
                header = next(reader)

            return [i for i in reader]


    def pick_column_num(self, import_data):
        """
        インポートファイルの列数を算出する
        :param import_data: array(two-dimensional)
        :return: int
        """
        columns = []
        for raw in import_data:
            columns.append(len(raw))
        if len(set(columns)) == 1:
            return columns[0]
        else:
            raise ValueError('this import files has diffrenect column numbers.')


    def insert_csv_file(self):
        input_file = self.read_import_file()
        column = self.pick_column_num(input_file)
        val_questions = ['?' for i in range(column)]
        cur.executemany("insert into {0} values ({1})".format(self.target_table_name, ','.join(val_questions)), input_file)


if __name__ == '__main__':

    sql = ImportSQLite(
        dbname=dbname,
        target_table_name=target_table_name,
        import_data_name=import_table_name,
        is_create_table=is_create_table,
        is_header_skip= is_header_skip,
        sql_create_table=sql_script
    )

    conn = sqlite3.connect(sql.dbname)
    cur = conn.cursor()

    if sql.is_create_table:
        cur.execute('drop table if exists {};'.format(target_table_name))
        cur.execute(sql.sql_create_table)

    sql.insert_csv_file()

    conn.commit()
    conn.close()

解説

設定項目

変数名	型	設定方法
dbname	Text	接続先のDB名。パスの形で指定する。
target_table_name	Text	インポート先となるDB上のテーブルの名前
import_table_name	Text	インポートしたいデータの名前。パスの形で指定する。
is_create_table	Boolean(True or False)	インポート先となるテーブルを作成するか否か
is_header_skip	Boolean(True or False)	インポートするデータのヘッダーを読み飛ばすか否か

諸注意

• インポート先となるテーブルを作成する場合、作成クエリをセットしないとエラーで落ちます。
• csv,tsv以外のファイルをインプットとして指定すると弾かれます。

使用法とサンプル

例えば以下のようなCSVがあって、接続先のDBにはまだテーブルがなかった場合

students.csv

id,name,class,blood_type
1,Mike,Moon,B
2,Bob,Song,A
3,Gonzalez,Star,AB
4,Alex,Moon,

設定項目はこのようになります。

（冒頭）import_data_to_sqlite.py

#######↓ここのパラメータを変える↓#######
dbname = 'test.db' # '/home/user/test.db'といった形で指定も可能
target_table_name = 'students'
import_table_name = 'students.csv' # '/home/user/students.csv'といった形で指定可能
is_create_table = True
is_header_skip = True
#################################


#######↓インポート先のテーブルDDL作成↓#######
sql_script = """
create table students(
    id integer,
    name text,
    class text,
    blood_type text
);
"""
#######################################

あとは普通に起動するだけ

kick.sh

$ python import_data_to_sqlite.py

おしまい。