1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

#!/usr/bin/env python

# -*- coding: utf-8 -*-

import codecs

import csv

import logging

import multiprocessing

import os

import warnings

import click

import MySQLdb

import sqlalchemy

warnings.filterwarnings('ignore', category=MySQLdb.Warning)

# 批量插入的记录数量

BATCH = 5000

DB_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() + '\n'

lines.append(first_line)

for line in infd:

clean_line = line.decode('gb18030').encode('utf8')

clean_line = clean_line.rstrip() + '\n'

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()