Python多进程导入CSV至数据库

#!/usr/bin/env python# -*- coding: utf-8 -*-import codecsimport csvimport loggingimport multiprocessingimport osimport warningsimport clickimport MySQLdbimport sqlalchemywarnings.filterwarnings('ignore', category=MySQLdb.Warning)# 批量插入的记录数量BATCH = 5000DB_URI = 'mysql://root@localhost:3306/example?charset=utf8'engine = sqlalchemy.create_engine(DB_URI)def get_table_cols(table):  sql = 'SELECT * FROM `{table}` LIMIT 0'.format(table=table)  res = engine.execute(sql)  return res.keys()def insert_many(table, cols, rows, cursor):  sql = 'INSERT INTO `{table}` ({cols}) VALUES ({marks})'.format(      table=table,      cols=', '.join(cols),      marks=', '.join(['%s'] * len(cols)))  cursor.execute(sql, *rows)  logging.info('process %s inserted %s rows into table %s', os.getpid(), len(rows), table)def insert_worker(table, cols, queue):  rows = []  # 每个子进程创建自己的 engine 对象  cursor = sqlalchemy.create_engine(DB_URI)  while True:    row = queue.get()    if row is None:      if rows:        insert_many(table, cols, rows, cursor)      break    rows.append(row)    if len(rows) == BATCH:      insert_many(table, cols, rows, cursor)      rows = []def insert_parallel(table, reader, w=10):  cols = get_table_cols(table)  # 数据队列，主进程读文件并往里写数据，worker 进程从队列读数据  # 注意一下控制队列的大小，避免消费太慢导致堆积太多数据，占用过多内存  queue = multiprocessing.Queue(maxsize=w*BATCH*2)  workers = []  for i in range(w):    p = multiprocessing.Process(target=insert_worker, args=(table, cols, queue))    p.start()    workers.append(p)    logging.info('starting # %s worker process, pid: %s...', i + 1, p.pid)  dirty_data_file = './{}_dirty_rows.csv'.format(table)  xf = open(dirty_data_file, 'w')  writer = csv.writer(xf, delimiter=reader.dialect.delimiter)  for line in reader:    # 记录并跳过脏数据: 键值数量不一致    if len(line) != len(cols):      writer.writerow(line)      continue    # 把 None 值替换为 'NULL'    clean_line = [None if x == 'NULL' else x for x in line]    # 往队列里写数据    queue.put(tuple(clean_line))    if reader.line_num % 500000 == 0:      logging.info('put %s tasks into queue.', reader.line_num)  xf.close()  # 给每个 worker 发送任务结束的信号  logging.info('send close signal to worker processes')  for i in range(w):    queue.put(None)  for p in workers:    p.join()def convert_file_to_utf8(f, rv_file=None):  if not rv_file:    name, ext = os.path.splitext(f)    if isinstance(name, unicode):      name = name.encode('utf8')    rv_file = '{}_utf8{}'.format(name, ext)  logging.info('start to process file %s', f)  with open(f) as infd:    with open(rv_file, 'w') as outfd:      lines = []      loop = 0      chunck = 200000      first_line = infd.readline().strip(codecs.BOM_UTF8).strip() + ''      lines.append(first_line)      for line in infd:        clean_line = line.decode('gb18030').encode('utf8')        clean_line = clean_line.rstrip() + ''        lines.append(clean_line)        if len(lines) == chunck:          outfd.writelines(lines)          lines = []          loop += 1          logging.info('processed %s lines.', loop * chunck)      outfd.writelines(lines)      logging.info('processed %s lines.', loop * chunck + len(lines))@click.group()def cli():  logging.basicConfig(level=logging.INFO,            format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')@cli.command('gbk_to_utf8')@click.argument('f')def convert_gbk_to_utf8(f):  convert_file_to_utf8(f)@cli.command('load')@click.option('-t', '--table', required=True, help='表名')@click.option('-i', '--filename', required=True, help='输入文件')@click.option('-w', '--workers', default=10, help='worker 数量，默认 10')def load_fac_day_pro_nos_sal_table(table, filename, workers):  with open(filename) as fd:    fd.readline()  # skip header    reader = csv.reader(fd)    insert_parallel(table, reader, w=workers)if name == 'main':  cli()