psycopg2 export DB to csv with column names-CodePudding

I'm using psycopg2 to connect to postgre DB, and to export the data into CSV file.

This is how I made the export DB to csv:

def export_table_to_csv(self, table, csv_path):
    sql = "COPY (SELECT * FROM %s) TO STDOUT WITH CSV DELIMITER ','" % table
    self.cur.execute(sql)

    with open(csv_path, "w") as file:
        self.cur.copy_expert(sql, file)

But the data is just the rows - without the column names.

How can I export the data with the column names?

P.S. I am able to print the column names:

sql = '''SELECT * FROM test'''
self.cur.execute(sql)
column_names = [desc[0] for desc in self.cur.description]
for i in column_names:
    print(i)

I want the cleanest way to do export the DB with columns name (i.e. I prefer to do this in one method, and not rename columns In retrospect).

CodePudding user response：

As I said in my comment, you can add HEADER to the WITH clause of your SQL:

sql = "COPY (SELECT * FROM export_test) TO STDOUT WITH CSV HEADER"

By default, comma delimiters are used with CSV option so you don't need to specify.

For future Questions, you should submit a minimal reproducible example. That is, code we can directly copy and paste and run. I was curious if this would work so I made one and tried it:

import psycopg2

conn = psycopg2.connect('host=<host> dbname=<dbname> user=<user>')
cur = conn.cursor()

# create test table
cur.execute('DROP TABLE IF EXISTS export_test')
sql = '''CREATE TABLE export_test
(
  id integer,
  uname text,
  fruit1 text,
  fruit2 text,
  fruit3 text
)'''
cur.execute(sql)

# insert data into table
sql = '''BEGIN;
insert into export_test
(id, uname, fruit1, fruit2, fruit3)
values(1, 'tom jones', 'apple', 'banana', 'pear');
insert into export_test
(id, uname, fruit1, fruit2, fruit3)
values(2, 'billy idol', 'orange', 'cherry', 'strawberry');
COMMIT;'''
cur.execute(sql)

# export to csv
fid = open('export_test.csv', 'w')
sql = "COPY (SELECT * FROM export_test) TO STDOUT WITH CSV HEADER"
cur.copy_expert(sql, fid)
fid.close()

And the resultant file is:

id,uname,fruit1,fruit2,fruit3
1,tom jones,apple,banana,pear
2,billy idol,orange,cherry,strawberry

CodePudding user response：

The easiest way for you is to use pandas. There are nice instruments for reading from database and writing to csv.

import pandas as pd

df = pd.read_sql("SELECT ...", db_connection)
df.to_csv('path/to/file.csv')