IT DASHBOARD Data Importer

Fork

#!/usr/bin/python
# -*- coding: utf-8 -*-

import json
import getpass
import os
import codecs
import re

import csv_parser
import database

# IT DASHBOARD http://www.itdashboard.go.jp/ の
# データをDBにインポートする

CSV_DIR = './csv_data'

def main():
    for f in os.listdir(CSV_DIR):
        (table_name, ext) = f.rsplit('.', 2)
        # table_nameから末尾の_[0-9]*を取り除く
        m = re.search(u'^(.*)_[0-9]*$', table_name)
        if m:
            table_name = m.group(1)
        print 'read %s...' % os.path.join(CSV_DIR, f)
        # IT DASHBOARDのデータはSJIS（cp932）
        fp = codecs.open(os.path.join(CSV_DIR, f), 'r', 'cp932')
        reader = csv_parser.reader(fp)
        header = reader.next()
        print 'put to %s' % table_name
        put_to_database(header, table_name, reader)
        fp.close()

def put_to_database(header, table_name, reader):
    if not table_exists(table_name):
        create_table(table_name, header)

    db = database.Database()
    print 'connect to database...'
    db.connect()

    print 'insert data to %s' % table_name
    # 1行目はカラム名なので読み飛ばす
    col_num = 0
    for row in reader:
        col_num += 1;
        if col_num == 1:
            continue
        quoted_header = ['`%s`' % x for x in header]
        columns = '( %s )' % ', '.join(quoted_header)
        placeholders = '( %s )' % ', '.join(['%s' for x in range(len(row))])
        sql = '''INSERT INTO %s %s VALUES %s;''' % (table_name, columns, placeholders)
        db.execute(sql, row)

    print 'close connection...'
    db.commit()
    db.close()
    print 'done.'

def table_exists(table_name):
    # table_nameテーブルが存在するならTrueを返す
    db = database.Database()
    db.connect()

    print 'check table %s exists...' % table_name
    sql = '''SHOW TABLES LIKE %s;'''
    result = db.query(sql, (table_name,))
    db.close()

    if result == None:
        return False
    if(len(result) != 0):
        return True
    return False

def create_table(table_name, items):
    # 項目名の末尾が「_budget」「_count」「_amount」であれば
    # INT型、それ以外はTEXT型として扱う
    columns = []
    for item in items:
        if re.search('_budget$', item):
            columns.append('''`%s` INT''' % item)
        elif re.search('_count$', item):
            columns.append('''`%s` INT''' % item)
        elif re.search('_amount$', item):
            columns.append('''`%s` INT''' % item)
        else:
            columns.append('''`%s` TEXT''' % item)

    sql = '''CREATE TABLE %s ( %s );''' % (table_name, ', \n'.join(columns))

    db = database.Database()
    print 'connect to database...'
    db.connect()
    
    print 'create table %s...' % table_name
    print sql
    db.execute(sql)
    
    print 'close connection...'
    db.close()
    print 'done.'

if __name__ == '__main__':
    main()