build_tables.py

#! /usr/bin/env python

import MySQLdb
from argparse import ArgumentParser
import gzip


def add_table(options):
    link = MySQLdb.connect(options.host, options.user, options.password)
    cursor = link.cursor()
    cursor.execute("use {};".format(options.database))
    if options.file.endswith(".gz"):
        infile = gzip.open(options.file, "rb")
    else:
        infile = open(options.file, "r")
    build = """create table {name} (assembly char(4), position int(9), context char(3), mc int(3), h int(3));""".format(name=options.name)
    cursor.execute(build)
    for line in infile:
        line = line.rsplit()
        chrom = line[0]
        if chrom == 'chrom':
            pass  # header
        else:
            pos = int(line[2])
            context = line[3]
            mc = int(line[6])
            h = int(line[7])
            add_data = """insert into {name} (assembly, position, context, mc, h) values ('{chrom}', {pos}, '{context}', {mc}, {h});""".format(
                   name=options.name,
                   chrom=chrom,
                   pos=pos,
                   context=context,
                   mc=mc,
                   h=h)
            cursor.execute(add_data)
    cursor.close()
    link.close()
    infile.close()


def indexTables(options):
    link = MySQLdb.connect(options.host, options.user, options.password)
    cursor = link.cursor()
    cursor.execute("use {};".format(options.database))
    query = "alter table {} add index (position)".format(options.name)
    cursor.execute(query)
    cursor.close()
    link.close()

if __name__ == "__main__":
    parser = ArgumentParser(description='Build MySQL tables from file')
    parser.add_argument('--host', help='MySQL host address', required=True)
    parser.add_argument('-u', '--user', help='MySQL user name', required=True)
    parser.add_argument('-p', '--password', help='MySQL password', required=True)
    parser.add_argument('-d', '--database', help='MySQL database name', required=True)
    parser.add_argument("-f", "--file", help="File name", required=True)
    parser.add_argument("-n", "--name", help="Sample name", required=True)
    options = parser.parse_args()
    add_table(options)
    indexTables(options)