使用python读写postgres脚本

Question

我有 postgres 表，我想使用 python 在这些表上运行 PostgreSQL 脚本文件，然后将查询结果写入 csv 文件。脚本文件有多个查询，用分号; 分隔。示例脚本如下所示

脚本文件：

--Duplication Check
select p.*, c.name
from scale_polygons_v3 c inner join cartographic_v3 p
on (metaphone(c.name_displ, 20) LIKE metaphone(p.name, 20)) AND c.kind NOT IN (9,10)
where ST_Contains(c.geom, p.geom);

--Area Check
select sp.areaid,sp.name_displ,p.road_id,p.name
from scale_polygons_v3 sp, pak_roads_20162207 p
where st_contains(sp.geom,p.geom) and sp.kind = 1
and p.areaid != sp.areaid;

当我运行 python 代码时，它成功执行而没有任何错误，但我面临的问题是，在将查询结果写入 csv 文件期间。只有最后执行的查询的结果被写入 csv 文件。这意味着第一个查询结果被第二个查询覆盖，第二个被第三个覆盖，依此类推，直到最后一个查询。

这是我的 python 代码：

import psycopg2
import sys
import csv
import datetime, time

def run_sql_file(filename, connection):
'''
    The function takes a filename and a connection as input
    and will run the SQL query on the given connection  
'''
    start = time.time()

    file = open(filename, 'r')
    sql = s = " ".join(file.readlines())
   #sql = sql1[3:]
    print "Start executing: " + " at " + str(datetime.datetime.now().strftime("%Y-%m-%d %H:%M")) + "\n"
    print "Query:\n", sql + "\n"
    cursor = connection.cursor()
    cursor.execute(sql)
    records = cursor.fetchall()
    with open('Report.csv', 'a') as f:
        writer = csv.writer(f, delimiter=',')
        for row in records:
            writer.writerow(row)
    connection.commit()
    end = time.time()
    row_count = sum(1 for row in records)
    print "Done Executing:", filename
    print "Number of rows returned:", row_count
    print "Time elapsed to run the query:",str((end - start)*1000) + ' ms'
    print "\t ==============================="

def main():    
    connection = psycopg2.connect("host='localhost' dbname='central' user='postgres' password='tpltrakker'")
    run_sql_file("script.sql", connection)
    connection.close()

if __name__ == "__main__":
    main()

我的代码有什么问题？

Answer 1

如果您能够稍微更改 SQL 脚本，那么这里有一个解决方法：

#!/usr/bin/env python

import psycopg2

script = '''
    declare cur1 cursor for
        select * from (values(1,2),(3,4)) as t(x,y);

    declare cur2 cursor for
        select 'a','b','c';
    '''
print script

conn = psycopg2.connect('');

# Cursors exists and available only inside the transaction
conn.autocommit = False;

# Create cursors from script
conn.cursor().execute(script);

# Read names of cursors
cursors = conn.cursor();
cursors.execute('select name from pg_cursors;')
cur_names = cursors.fetchall()

# Read data from each available cursor
for cname in cur_names:
    print cname[0]
    cur = conn.cursor()
    cur.execute('fetch all from ' + cname[0])
    rows = cur.fetchall()
    # Here you can save the data to the file
    print rows


conn.rollback()

print 'done'

免责声明：我完全是 Python 新手。

Answer 2

这是将每个查询输出为不同文件的最简单方法。 copy_expert

query = '''
    select p.*, c.name
    from
        scale_polygons_v3 c
        inner join
        cartographic_v3 p on metaphone(c.name_displ, 20) LIKE metaphone(p.name, 20) and c.kind not in (9,10)
    where ST_Contains(c.geom, p.geom)
'''
copy = "copy ({}) to stdout (format csv)".format(query)
f = open('Report.csv', 'wb')
cursor.copy_expert(copy, f, size=8192)
f.close()

query = '''
    select sp.areaid,sp.name_displ,p.road_id,p.name
    from scale_polygons_v3 sp, pak_roads_20162207 p
    where st_contains(sp.geom,p.geom) and sp.kind = 1 and p.areaid != sp.areaid;
'''
copy = "copy ({}) to stdout (format csv)".format(query)
f = open('Report2.csv', 'wb')
cursor.copy_expert(copy, f, size=8192)
f.close()

如果您想将第二个输出附加到同一个文件，那么只需保持第一个文件对象处于打开状态。

请注意，copy 必须输出到 stdout 以使其对 copy_expert 可用