psycopgのCOPYサポート

Using COPY TO and COPY FROMを読むとCOPYが使えるとあるが、write_row()とwrite()は（当たり前と言えば当たり前だが）挙動が異なる。

records = [(10, 20, "hello"), (40, None, "world")]

with cursor.copy("COPY sample (col1, col2, col3) FROM STDIN") as copy:
    for record in records:
        copy.write_row(record)

with open("data", "r") as f:
    with cursor.copy("COPY data FROM STDIN") as copy:
        while data := f.read(BLOCK_SIZE):
            copy.write(data)

write_row()はリストを行として扱うが、write()はファイルからread()したものをそのまま流し込むので、ファイルの行の終端をexpectしている。なので、write()でPythonのリストやstrを書き込む時は、最後に\nが必要。

with cursor.copy("COPY data FROM STDIN") as copy:
    for i in range(10)
        copy.write(i + '\n')

追記（2024.11.11）：テストしていて気付いたが、psycopgはCOPYコマンドのFREEZEオプションをサポートしてないっぽい。んー…

追記（2024.11.12）：FREEZEサポートしてた。つまり、トランザクション内でテーブルを作成するかTRUNCATEしないとダメ。psycopgはコンテキストマネージャに対応しているので、withブロック内でそれらを実行してからCOPYすれば良かった。

with conn.cur() as cur:
    cur.execute("TRUNCATE data")
    with cursor.copy("COPY data FROM STDIN FREEZE") as copy:
        for i in range(10)
            copy.write(i + '\n')

関連