如何根据日期使用Python将数据导入多个CSV文件

Question

I have created the below code to import data in CSV file from PostgreSQL DB. However, I want to create multiple files based on date.

import psycopg2
import csv

conn_string = "host='' port='5432' user='' password='' dbname=''"

conn = psycopg2.connect(conn_string)

cur=conn.cursor()

query="select * from sample where date between '' and ''"

cur.execute(query)

title=[i[0] for i in cur.description]

result=cur.fetchall()

csvfile=open('filename.csv','w')

if result:
    c = csv.writer(csvfile)
    c.writerow(title)
    c.writerows(result)


cur.close()
conn.close()

The files should be split similar to the below format:

01jan.csv
02jan.csv 
etc.

Answer 1

You can loop over the query results and open a new file whenever the row date changes. The results must be ordered by date, otherwise you can lose some data.

import psycopg2
import psycopg2.extras
import csv
import datetime

# conn_string = ...
conn = psycopg2.connect(conn_string)

# we need results in dict
cur = conn.cursor(cursor_factory = psycopg2.extras.DictCursor)

# order by date - important!
query = "select * from sample where date between '2018-01-01' and '2018-01-10' order by date"
cur.execute(query)
title = [i[0] for i in cur.description]

date = None
writer = None
csvfile = None

for row in cur:
    if date != row['date']:
        # when date changes we should close current file (if opened)
        # and open a new one with name based on date
        if csvfile:
            csvfile.close()
        date = row['date']
        filename = date.strftime("%d%b")+ '.csv'
        csvfile = open(filename, 'w', newline='')
        writer = csv.writer(csvfile)
        writer.writerow(title)
    writer.writerow(row)

cur.close()
conn.close()

The above solution is acceptable for rather small datasets. If the amount of data for one day is large, you should rather use copy_expert()

cur = conn.cursor()

# example loop for ten days of Jan 2018
for day in range(1, 10):
    date = datetime.date(2018, 1, day)
    filename = date.strftime("%d%b")+ '.csv'
    command = 'copy (select * from sample where date = %s) to stdout with csv header'
    sql = cur.mogrify(command, [date])
    with open(filename, 'w') as file:
        cur.copy_expert(sql, file)