Psycopg2：将 CSV 数据复制到具有额外列值的表

我正在使用copy_expert中的方法psycopg2将数据从 CSV 文件复制到 postgresql 表。我有这样的 postgres 表：

create table my_table (
  cola          text,
  colb          text,
  colc          text,
  cold          text,
  cole          text,
  colf          text,
  colg          text
)

还有一个包含前五列数据的 CSV，如下所示：

cola,colb,colc,cold,cole
1,foo,a,10,vvv
2,bar,b,20,www
3,baz,c,30,xxx
4,boo,d,40,yyy
5,baa,e,50,zzz

我想复制前五列的 CSV 数据，同时还指定colf and colg（每行应该具有相同的值colf and colg).

我可以将前五列复制到我的表中，如下所示：

conn = psycopg2.connect('dbname=name user=username')
cur = conn.cursor()
copy_sql = """
  copy my_table (cola, colb, colc, cold, cole)
  from stdin with
    csv
    header
    delimiter as ','
"""
from_csv = '/path/to/data.csv'
with open(from_csv, 'r') as f:
  cur.copy_expert(sql=copy_sql, file=f)
  conn.commit()
  cur.close()

如何使用 python 指定最后两列的值？我知道我可以在表 DDL 中指定默认值，如下所示：

create table my_table (
  cola          text,
  colb          text,
  colc          text,
  cold          text,
  cole          text,
  colf          text default 'foo',
  colg          text default 'bar'
)

但我想使用 python 添加值，因为每个 CSV 上传都有自己的值colf and colg，这些值是由我的 python 代码中的逻辑确定的。

看起来有几种方法可以做到这一点，首先将我需要的列添加到数据中，然后上传更新的数据。

使用petl包裹：

import psycopg2
from petl import fromcsv, addfield, todb

csv_file = '/path/to/data.csv'
table = fromcsv(csv_file)
table = addfield(table, 'colf', 'Some value')
table = addfield(table, 'colg', 'Another value')

conn = psycopg2.connect('dbname=test user=user')
todb(table, conn, 'my_table')

这对于小数据来说效果很好，但是对于大数据来说却非常慢。这psycopg2 copy_from and copy_expert命令似乎运行得更快，因为它们使用了postgresql批量复制。我能够使用复制我的数据copy_from首先将我的 csv 文件转换为pandas dataframe:

import psycopg2
import pandas as pd
from io import StringIO

csv_file = '/path/to/file'
df = pd.read_csv(csv_file)
df['colf'] = 'My value'
df['colg'] = 'Foobar'

为了利用psycopg2 copy_命令，我需要转换dataframe到一个类似文件的对象中read() and readline()方法，我可以使用StringIO:

buf = StringIO()
df.to_csv(buf, header=False, index=False)
buf.pos = 0

请注意，您需要设置pos缓冲区为 0，因为pandas.to_csv似乎默认将 pos 设置为末尾。看这个答案 https://stackoverflow.com/a/38468155/1994092以获得解释。

然后我可以复制该缓冲区对象：

conn = psycopg2.connect('dbname=test user=user')
cur = conn.cursor()
cur.copy_from(buf, 'my_table', sep=',')
conn.commit()
cur.close()