[Python] 调取MYSQL数据并插入到CSV文件

showerlee 2015-04-17 14:35 Programming, PYTHON, 其他阅读 (9,403) 抢沙发

如何利用python脚本将远程数据库查询值,并将该值按照csv中"column A"对应关系整体插入到"column B",最近handle了一个case,联想到python天生对数据处理较shell有明显的优势,最后尝试用脚本搞定这个看起来逻辑很简单,但又不简单的data processing.

Target:

column A 是DB内存储的若干Project name,需要通过查询数据库,获取这些project name在DB对应的Project key的具体值,并插入column A后保存.

#!/usr/bin/env python
# encoding: utf-8

import MySQLdb
import csv
import sys

# Define csv list
Csv_content = []
Csv_content_edited = []

# Define project list
Project_names = []

# Define db list
Db_content = []

# Define the file that needs to be handled.
try:
	file_name = sys.argv[1]
	new_file_name = file_name.split('.')[0] + '_new.' + file_name.split('.')[1]
except IndexError:
	pass

def selectDB():
	# Open db connection
	db = MySQLdb.connect("test.com", "testuser", "testuser", "testdb")

	# Use cursor() fuction to get current db cursor
	cursor = db.cursor()

	Csv_content_edited = readContent(file_name)

	for c in range(1,len(Csv_content_edited)):
		Project_names.append(Csv_content_edited[c][0])

	Project_name_string = ",".join(['"' + p + '"' for p in Project_names])

	# SQL "SELECT" statement
	sql = 'select pname,pkey from project where pname in (%s)' %Project_name_string
	
	try:
		# Execute SQL
		cursor.execute(sql)
		# Obtain all the record list
		results = cursor.fetchall()

		for row in results:
			#lower_user_name = row[3]
			Db_content.append(row)

		return Db_content
	except:
		print "Error: unable to fecth data"

	# Close connection
	db.close()


def readContent(file_name):
	# Read the csv file,then put it into list.
	with open(file_name, 'r') as csvfile:
		csv_reader = csv.reader(csvfile, delimiter=',')
		for row in csv_reader:			 
			if row[0]:
				Csv_content.append(row)
	return Csv_content


def insert_col():
	Csv_content_edited = readContent(file_name)
	# Insert null value to each components of "Csv_content_edited" afterward.
	for i in range(0,len(Csv_content_edited)):
		Csv_content_edited[i].insert(1,'')

	# Define the second inserted column title.
	Csv_content_edited[0][1] = "Pkey"
	# Grab the users data from db.
	Db_content = selectDB()
	# print Db_content
	
	for d in range(0,len(Db_content)):
		Pkey = Db_content[d][1]
		Pname = Db_content[d][0]
		for c in range(0,len(Csv_content_edited)):  
			if Csv_content_edited[c][0] == Pname:
				Csv_content_edited[c][1] = Pkey
	# print Csv_content_edited

	pname_list = []
	Csv_content_edited_new = []
	
	for c in range(0,len(Csv_content_edited)):
		if not Csv_content_edited[c][0] in pname_list:
			pname_list.append(Csv_content_edited[c][0])
			Csv_content_edited_new.append(Csv_content_edited[c])
	# print Csv_content_edited_new
	
	return Csv_content_edited_new



# Write the csv file. 
def writeContent():
	with open(new_file_name,'wb') as csvfile:
		csv_writer = csv.writer(csvfile)
		csv_writer.writerows(insert_col())


# Execute the finnal function.	
if __name__ == '__main__':
	try:
		writeContent()
	except (IOError,NameError,IndexError):
		print "Please type the correct file name. e.g: '" + sys.argv[0] + " testfile.csv'"
	else:
		print 'The result file is: %s' %new_file_name

Result:

Before the change: