csv2mysql

<span style="color: #0000ff">import</span><span style="color: #000000"> os
</span><span style="color: #0000ff">import</span><span style="color: #000000"> re
</span><span style="color: #0000ff">import</span><span style="color: #000000"> sys
</span><span style="color: #0000ff">import</span><span style="color: #000000"> csv
</span><span style="color: #0000ff">import</span><span style="color: #000000"> time
</span><span style="color: #0000ff">import</span><span style="color: #000000"> argparse
</span><span style="color: #0000ff">import</span><span style="color: #000000"> collections
</span><span style="color: #0000ff">import</span><span style="color: #000000"> MySQLdb
</span><span style="color: #0000ff">import</span><span style="color: #000000"> warnings 
</span><span style="color: #008000">#</span><span style="color: #008000"> suppress annoying mysql warnings</span>
warnings.filterwarnings(action=<span style="color: #800000">'</span><span style="color: #800000">ignore</span><span style="color: #800000">'</span>, category=<span style="color: #000000">MySQLdb.Warning) 



</span><span style="color: #0000ff">def</span><span style="color: #000000"> get_type(s):
    </span><span style="color: #800000">"""</span><span style="color: #800000">Find type for this string
    </span><span style="color: #800000">"""</span>
    <span style="color: #008000">#</span><span style="color: #008000"> try integer type</span>
    <span style="color: #0000ff">try</span><span style="color: #000000">:
        v </span>=<span style="color: #000000"> int(s)
    </span><span style="color: #0000ff">except</span><span style="color: #000000"> ValueError:
        </span><span style="color: #0000ff">pass</span>
    <span style="color: #0000ff">else</span><span style="color: #000000">:
        </span><span style="color: #0000ff">if</span> abs(v) > 2147483647<span style="color: #000000">:
            </span><span style="color: #0000ff">return</span> <span style="color: #800000">'</span><span style="color: #800000">bigint</span><span style="color: #800000">'</span>
        <span style="color: #0000ff">else</span><span style="color: #000000">:
            </span><span style="color: #0000ff">return</span> <span style="color: #800000">'</span><span style="color: #800000">int</span><span style="color: #800000">'</span>
    <span style="color: #008000">#</span><span style="color: #008000"> try float type</span>
    <span style="color: #0000ff">try</span><span style="color: #000000">:
        float(s)
    </span><span style="color: #0000ff">except</span><span style="color: #000000"> ValueError:
        </span><span style="color: #0000ff">pass</span>
    <span style="color: #0000ff">else</span><span style="color: #000000">:
        </span><span style="color: #0000ff">return</span> <span style="color: #800000">'</span><span style="color: #800000">double</span><span style="color: #800000">'</span>

    <span style="color: #008000">#</span><span style="color: #008000"> check for timestamp</span>
    dt_formats =<span style="color: #000000"> (
        (</span><span style="color: #800000">'</span><span style="color: #800000">%Y-%m-%d %H:%M:%S</span><span style="color: #800000">'</span>, <span style="color: #800000">'</span><span style="color: #800000">datetime</span><span style="color: #800000">'</span><span style="color: #000000">),
        (</span><span style="color: #800000">'</span><span style="color: #800000">%Y-%m-%d %H:%M:%S.%f</span><span style="color: #800000">'</span>, <span style="color: #800000">'</span><span style="color: #800000">datetime</span><span style="color: #800000">'</span><span style="color: #000000">),
        (</span><span style="color: #800000">'</span><span style="color: #800000">%Y-%m-%d</span><span style="color: #800000">'</span>, <span style="color: #800000">'</span><span style="color: #800000">date</span><span style="color: #800000">'</span><span style="color: #000000">),
        (</span><span style="color: #800000">'</span><span style="color: #800000">%H:%M:%S</span><span style="color: #800000">'</span>, <span style="color: #800000">'</span><span style="color: #800000">time</span><span style="color: #800000">'</span><span style="color: #000000">),
    )
    </span><span style="color: #0000ff">for</span> dt_format, dt_type <span style="color: #0000ff">in</span><span style="color: #000000"> dt_formats:
        </span><span style="color: #0000ff">try</span><span style="color: #000000">:
            time.strptime(s, dt_format)
        </span><span style="color: #0000ff">except</span><span style="color: #000000"> ValueError:
            </span><span style="color: #0000ff">pass</span>
        <span style="color: #0000ff">else</span><span style="color: #000000">:
            </span><span style="color: #0000ff">return</span><span style="color: #000000"> dt_type
   
    </span><span style="color: #008000">#</span><span style="color: #008000"> doesn't match any other types so assume text</span>
    <span style="color: #0000ff">if</span> len(s) > 255<span style="color: #000000">:
        </span><span style="color: #0000ff">return</span> <span style="color: #800000">'</span><span style="color: #800000">text</span><span style="color: #800000">'</span>
    <span style="color: #0000ff">else</span><span style="color: #000000">:
        </span><span style="color: #0000ff">return</span> <span style="color: #800000">'</span><span style="color: #800000">varchar(255)</span><span style="color: #800000">'</span>


<span style="color: #0000ff">def</span> most_common(l, default=<span style="color: #800000">'</span><span style="color: #800000">varchar(255)</span><span style="color: #800000">'</span><span style="color: #000000">):
    </span><span style="color: #800000">"""</span><span style="color: #800000">Return most common value from list
    </span><span style="color: #800000">"""</span>
    <span style="color: #008000">#</span><span style="color: #008000"> some formats trump others</span>
    <span style="color: #0000ff">if</span><span style="color: #000000"> l:
        </span><span style="color: #0000ff">for</span> dt_type <span style="color: #0000ff">in</span> (<span style="color: #800000">'</span><span style="color: #800000">text</span><span style="color: #800000">'</span>, <span style="color: #800000">'</span><span style="color: #800000">bigint</span><span style="color: #800000">'</span><span style="color: #000000">):
            </span><span style="color: #0000ff">if</span> dt_type <span style="color: #0000ff">in</span><span style="color: #000000"> l:
                </span><span style="color: #0000ff">return</span><span style="color: #000000"> dt_type
        </span><span style="color: #0000ff">return</span> max(l, key=<span style="color: #000000">l.count)
    </span><span style="color: #0000ff">return</span><span style="color: #000000"> default


</span><span style="color: #0000ff">def</span> get_col_types(input_file, max_rows=1000<span style="color: #000000">):
    </span><span style="color: #800000">"""</span><span style="color: #800000">Find the type for each CSV column
    </span><span style="color: #800000">"""</span><span style="color: #000000">
    csv_types </span>=<span style="color: #000000"> collections.defaultdict(list)
    reader </span>=<span style="color: #000000"> csv.reader(open(input_file))
    </span><span style="color: #008000">#</span><span style="color: #008000"> test the first few rows for their data types</span>
    <span style="color: #0000ff">for</span> row_i, row <span style="color: #0000ff">in</span><span style="color: #000000"> enumerate(reader):
        </span><span style="color: #0000ff">if</span> row_i ==<span style="color: #000000"> 0:
            header </span>=<span style="color: #000000"> row
        </span><span style="color: #0000ff">else</span><span style="color: #000000">:
            </span><span style="color: #0000ff">for</span> col_i, s <span style="color: #0000ff">in</span><span style="color: #000000"> enumerate(row):
                data_type </span>=<span style="color: #000000"> get_type(s)
                csv_types[header[col_i]].append(data_type)
 
        </span><span style="color: #0000ff">if</span> row_i ==<span style="color: #000000"> max_rows:
            </span><span style="color: #0000ff">break</span>

    <span style="color: #008000">#</span><span style="color: #008000"> take the most common data type for each row</span>
    <span style="color: #0000ff">return</span> [most_common(csv_types[col]) <span style="color: #0000ff">for</span> col <span style="color: #0000ff">in</span><span style="color: #000000"> header]


</span><span style="color: #0000ff">def</span><span style="color: #000000"> get_schema(table, header, col_types):
    </span><span style="color: #800000">"""</span><span style="color: #800000">Generate the schema for this table from given types and columns
    </span><span style="color: #800000">"""</span><span style="color: #000000">
    schema_sql </span>= <span style="color: #800000">"""</span><span style="color: #800000">CREATE TABLE IF NOT EXISTS %s ( 
        id int NOT NULL AUTO_INCREMENT,</span><span style="color: #800000">"""</span> %<span style="color: #000000"> table 

    </span><span style="color: #0000ff">for</span> col_name, col_type <span style="color: #0000ff">in</span><span style="color: #000000"> zip(header, col_types):
        schema_sql </span>+= <span style="color: #800000">'</span><span style="color: #800000">\n%s %s,</span><span style="color: #800000">'</span> %<span style="color: #000000"> (col_name, col_type)

    schema_sql </span>+= <span style="color: #800000">"""</span><span style="color: #800000">\nPRIMARY KEY (id)
        ) DEFAULT CHARSET=utf8;</span><span style="color: #800000">"""</span>
    <span style="color: #0000ff">return</span><span style="color: #000000"> schema_sql


</span><span style="color: #0000ff">def</span><span style="color: #000000"> get_insert(table, header):
    </span><span style="color: #800000">"""</span><span style="color: #800000">Generate the SQL for inserting rows
    </span><span style="color: #800000">"""</span><span style="color: #000000">
    field_names </span>= <span style="color: #800000">'</span><span style="color: #800000">, </span><span style="color: #800000">'</span><span style="color: #000000">.join(header)
    field_markers </span>= <span style="color: #800000">'</span><span style="color: #800000">, </span><span style="color: #800000">'</span>.join(<span style="color: #800000">'</span><span style="color: #800000">%s</span><span style="color: #800000">'</span> <span style="color: #0000ff">for</span> col <span style="color: #0000ff">in</span><span style="color: #000000"> header)
    </span><span style="color: #0000ff">return</span> <span style="color: #800000">'</span><span style="color: #800000">INSERT INTO %s (%s) VALUES (%s);</span><span style="color: #800000">'</span> %<span style="color: #000000"> \
        (table, field_names, field_markers)


</span><span style="color: #0000ff">def</span><span style="color: #000000"> format_header(row):
    </span><span style="color: #800000">"""</span><span style="color: #800000">Format column names to remove illegal characters and duplicates
    </span><span style="color: #800000">"""</span><span style="color: #000000">
    safe_col </span>= <span style="color: #0000ff">lambda</span> s: re.sub(<span style="color: #800000">'</span><span style="color: #800000">\W+</span><span style="color: #800000">'</span>, <span style="color: #800000">'</span><span style="color: #800000">_</span><span style="color: #800000">'</span>, s.lower()).strip(<span style="color: #800000">'</span><span style="color: #800000">_</span><span style="color: #800000">'</span><span style="color: #000000">)
    header </span>=<span style="color: #000000"> []
    counts </span>=<span style="color: #000000"> collections.defaultdict(int)
    </span><span style="color: #0000ff">for</span> col <span style="color: #0000ff">in</span><span style="color: #000000"> row:
        col </span>=<span style="color: #000000"> safe_col(col)
        counts[col] </span>+= 1
        <span style="color: #0000ff">if</span> counts[col] > 1<span style="color: #000000">:
            col </span>= <span style="color: #800000">'</span><span style="color: #800000">{}{}</span><span style="color: #800000">'</span><span style="color: #000000">.format(col, counts[col])
        header.append(col)
    </span><span style="color: #0000ff">return</span><span style="color: #000000"> header


</span><span style="color: #0000ff">def</span> main(input_file, user, password, host, table, database, max_inserts=10000<span style="color: #000000">):
    </span><span style="color: #0000ff">print</span> <span style="color: #800000">"</span><span style="color: #800000">Importing `%s' into MySQL database `%s.%s'</span><span style="color: #800000">"</span> %<span style="color: #000000"> (input_file, database, table)
    db </span>= MySQLdb.connect(host=host, user=user, passwd=password, charset=<span style="color: #800000">'</span><span style="color: #800000">utf8</span><span style="color: #800000">'</span><span style="color: #000000">)
    cursor </span>=<span style="color: #000000"> db.cursor()
    </span><span style="color: #008000">#</span><span style="color: #008000"> create database and if doesn't exist</span>
    cursor.execute(<span style="color: #800000">'</span><span style="color: #800000">CREATE DATABASE IF NOT EXISTS %s;</span><span style="color: #800000">'</span> %<span style="color: #000000"> database)
    db.select_db(database)

    </span><span style="color: #008000">#</span><span style="color: #008000"> define table</span>
    <span style="color: #0000ff">print</span> <span style="color: #800000">'</span><span style="color: #800000">Analyzing column types ...</span><span style="color: #800000">'</span><span style="color: #000000">
    col_types </span>=<span style="color: #000000"> get_col_types(input_file)
    </span><span style="color: #0000ff">print</span><span style="color: #000000"> col_types

    header </span>=<span style="color: #000000"> None
    </span><span style="color: #0000ff">for</span> i, row <span style="color: #0000ff">in</span><span style="color: #000000"> enumerate(csv.reader(open(input_file))):
        </span><span style="color: #0000ff">if</span><span style="color: #000000"> header:
            </span><span style="color: #0000ff">while</span> len(row) <<span style="color: #000000"> len(header):
                row.append(</span><span style="color: #800000">''</span>) <span style="color: #008000">#</span><span style="color: #008000"> this row is missing columns so pad blank values</span>
<span style="color: #000000">            cursor.execute(insert_sql, row)
            </span><span style="color: #0000ff">if</span> i % max_inserts ==<span style="color: #000000"> 0:
                db.commit()
                </span><span style="color: #0000ff">print</span> <span style="color: #800000">'</span><span style="color: #800000">commit</span><span style="color: #800000">'</span>
        <span style="color: #0000ff">else</span><span style="color: #000000">:
            header </span>=<span style="color: #000000"> format_header(row)
            schema_sql </span>=<span style="color: #000000"> get_schema(table, header, col_types)
            </span><span style="color: #0000ff">print</span><span style="color: #000000"> schema_sql
            </span><span style="color: #008000">#</span><span style="color: #008000"> create table</span>
            cursor.execute(<span style="color: #800000">'</span><span style="color: #800000">DROP TABLE IF EXISTS %s;</span><span style="color: #800000">'</span> %<span style="color: #000000"> table)
            cursor.execute(schema_sql)
            </span><span style="color: #008000">#</span><span style="color: #008000"> create index for more efficient access</span>
            <span style="color: #0000ff">try</span><span style="color: #000000">:
                cursor.execute(</span><span style="color: #800000">'</span><span style="color: #800000">CREATE INDEX ids ON %s (id);</span><span style="color: #800000">'</span> %<span style="color: #000000"> table)
            </span><span style="color: #0000ff">except</span><span style="color: #000000"> MySQLdb.OperationalError:
                </span><span style="color: #0000ff">pass</span> <span style="color: #008000">#</span><span style="color: #008000"> index already exists</span>

            <span style="color: #0000ff">print</span> <span style="color: #800000">'</span><span style="color: #800000">Inserting rows ...</span><span style="color: #800000">'</span>
            <span style="color: #008000">#</span><span style="color: #008000"> SQL string for inserting data</span>
            insert_sql =<span style="color: #000000"> get_insert(table, header)

    </span><span style="color: #008000">#</span><span style="color: #008000"> commit rows to database</span>
    <span style="color: #0000ff">print</span> <span style="color: #800000">'</span><span style="color: #800000">Committing rows to database ...</span><span style="color: #800000">'</span><span style="color: #000000">
    db.commit()
    </span><span style="color: #0000ff">print</span> <span style="color: #800000">'</span><span style="color: #800000">Done!</span><span style="color: #800000">'</span>



<span style="color: #0000ff">if</span> <span style="color: #800080">__name__</span> == <span style="color: #800000">'</span><span style="color: #800000">__main__</span><span style="color: #800000">'</span><span style="color: #000000">:
    parser </span>= argparse.ArgumentParser(description=<span style="color: #800000">'</span><span style="color: #800000">Automatically insert CSV contents into MySQL</span><span style="color: #800000">'</span><span style="color: #000000">)
    parser.add_argument(</span><span style="color: #800000">'</span><span style="color: #800000">--table</span><span style="color: #800000">'</span>, dest=<span style="color: #800000">'</span><span style="color: #800000">table</span><span style="color: #800000">'</span>, help=<span style="color: #800000">'</span><span style="color: #800000">Set the name of the table. If not set the CSV filename will be used</span><span style="color: #800000">'</span><span style="color: #000000">)
    parser.add_argument(</span><span style="color: #800000">'</span><span style="color: #800000">--database</span><span style="color: #800000">'</span>, dest=<span style="color: #800000">'</span><span style="color: #800000">database</span><span style="color: #800000">'</span>, default=<span style="color: #800000">'</span><span style="color: #800000">test</span><span style="color: #800000">'</span>, help=<span style="color: #800000">'</span><span style="color: #800000">Set the name of the database. If not set the test database will be used</span><span style="color: #800000">'</span><span style="color: #000000">)
    parser.add_argument(</span><span style="color: #800000">'</span><span style="color: #800000">--user</span><span style="color: #800000">'</span>, dest=<span style="color: #800000">'</span><span style="color: #800000">user</span><span style="color: #800000">'</span>, default=<span style="color: #800000">'</span><span style="color: #800000">root</span><span style="color: #800000">'</span>, help=<span style="color: #800000">'</span><span style="color: #800000">The MySQL login username</span><span style="color: #800000">'</span><span style="color: #000000">)
    parser.add_argument(</span><span style="color: #800000">'</span><span style="color: #800000">--password</span><span style="color: #800000">'</span>, dest=<span style="color: #800000">'</span><span style="color: #800000">password</span><span style="color: #800000">'</span>, default=<span style="color: #800000">''</span>, help=<span style="color: #800000">'</span><span style="color: #800000">The MySQL login password</span><span style="color: #800000">'</span><span style="color: #000000">)
    parser.add_argument(</span><span style="color: #800000">'</span><span style="color: #800000">--host</span><span style="color: #800000">'</span>, dest=<span style="color: #800000">'</span><span style="color: #800000">host</span><span style="color: #800000">'</span>, default=<span style="color: #800000">'</span><span style="color: #800000">localhost</span><span style="color: #800000">'</span>, help=<span style="color: #800000">'</span><span style="color: #800000">The MySQL host</span><span style="color: #800000">'</span><span style="color: #000000">)
    parser.add_argument(</span><span style="color: #800000">'</span><span style="color: #800000">input_file</span><span style="color: #800000">'</span>, help=<span style="color: #800000">'</span><span style="color: #800000">The input CSV file</span><span style="color: #800000">'</span><span style="color: #000000">)
    args </span>= parser.parse_args(sys.argv[1<span style="color: #000000">:])
    </span><span style="color: #0000ff">if</span> <span style="color: #0000ff">not</span><span style="color: #000000"> args.table:
        </span><span style="color: #008000">#</span><span style="color: #008000"> use input file name for table</span>
        args.table =<span style="color: #000000"> os.path.splitext(os.path.basename(args.input_file))[0]
    
    main(args.input_file, args.user, args.password, args.host, args.table, args.database)</span>