import csv
import logging
import itertools
from random import choice
from string import ascii_lowercase
from collections import namedtuple
from tempfile import NamedTemporaryFile
DELIMITER = '\t'
ENCODING = 'utf-8'
def get_random_cursor_name(length=7):
return ''.join(choice(ascii_lowercase) for i in range(length))
def extract_table(conn, query, itersize=2000, use_named_cursor=True):
"""
:type conn: psycopg2.connection
:param conn: connection object returned by psycopg2.connect.
:type query: string
:param query: SQL query string.
:type itersize: int
:param itersize: If itersize == -1, then fetchall, else fetchmany(itersize)
:type use_named_cursor: boolean
:param use_named_cursor: If true, then use server side cursor, else client side cursor.
"""
if use_named_cursor:
cursor_name = get_random_cursor_name()
with conn.cursor(cursor_name) as cursor:
cursor.itersize = 2000 if itersize == -1 else itersize
cursor.execute(query)
row = cursor.fetchone()
header = [desc[0] for desc in cursor.description]
Row = namedtuple('Row', header)
yield Row(*row)
for row in cursor:
yield Row(*row)
else:
with conn.cursor() as cursor:
cursor.execute(query)
header = [desc[0] for desc in cursor.description]
Row = namedtuple('Row', header)
if itersize == -1:
rows = cursor.fetchall()
else:
cursor.arraysize = itersize
rows = itertools.chain.from_iterable(iter(cursor.fetchmany, []))
for row in rows:
yield Row(*row)
def copy_table(src_conn, dst_conn, src_query, dst_table_name, use_named_cursor=False, truncate_table=False):
"""
Run query and import result to another database.
We use null="" in copy_from, because csv library dumps None as "" (empty string).
:type src_conn: psycopg2.connection
:param src_conn: Source database connection.
:type dst_conn: psycopg2.connection
:param dst_conn: Destination database connection.
:type src_query: string
:param src_query: Query to be executed on source database.
:type dst_table_name: string
:param dst_table_name: Table name in destination database.
"""
rows = extract_table(src_conn, src_query, use_named_cursor=use_named_cursor)
with NamedTemporaryFile('w+t', encoding=ENCODING, newline='') as fp, dst_conn.cursor() as cursor:
logging.info("Writing result to a temporary file.")
csv.writer(fp, delimiter=DELIMITER).writerows(rows)
fp.flush()
fp.seek(0)
if truncate_table:
logging.info("Truncating {}.".format(dst_table_name))
cursor.execute('TRUNCATE {};'.format(dst_table_name))
logging.info("Copying result to {}".format(dst_table_name))
cursor.copy_from(file=fp, table=dst_table_name, sep=DELIMITER, null="")
dst_conn.commit()
This code, in particular copy_table
, is used to run a query, write the result to a temporary CSV file and copy the result to a table.
An example of the usage is as follows:
Let's say you want to copy a subset of the table public.my_table
in DB1 to DB2.
import psycopg2
import credentials
src_conn = psycopg2.connect(
database=credentials.DB1_DB,
user=credentials.DB1_USER,
password=credentials.DB1_PASS,
host=credentials.DB1_HOST,
port=credentials.DB1_PORT
)
dst_conn = psycopg2.connect(
database=credentials.DB2_DB,
user=credentials.DB2_USER,
password=credentials.DB2_PASS,
host=credentials.DB2_HOST,
port=credentials.DB2_PORT
)
query = 'SELECT * FROM public.my_table where product_count > 200'
copy_table(src_conn, dst_conn, src_query=query, dst_table_name='public.my_table_subset')
How could I write this code to be more Pythonic?
Note that I'm using Python 3.5.