在Python中使用Postgres的COPY FROM文件查询而不写入临时文件

Question

我需要把一些数据从源头加载到Postgres数据库里。为此，我首先把数据写入一个临时的CSV文件，然后再通过COPY FROM查询把CSV文件里的数据加载到Postgres数据库中。我都是用Python来完成这个过程。

代码大概是这样的：

table_name = 'products'
temp_file = "'C:\\Users\\username\\tempfile.csv'"
db_conn = psycopg2.connect(host, port, user, password, database)
cursor = db_conn.cursor()
query = """COPY """ + table_name + """ FROM """ + temp_file + " WITH NULL AS ''; """
cursor.execute(query)

我想省去写中间文件这一步。相反，我想直接写入一个Python对象，然后用COPY FROM文件的方法把数据加载到Postgres数据库里。

我知道有一种方法是使用psycopg2的copy_from方法，它可以把数据从StringIO对象复制到Postgres数据库。不过，由于某些原因，我不能使用psycopg2，所以我希望我的COPY FROM操作不依赖于任何库。我想要的是一个Postgres查询，这样其他的Postgres驱动也能运行。

请给我一些建议，告诉我怎么做才能不写中间文件。

Answer 1

psycopg2 是一个可以直接支持 COPY 协议的工具，这样你就可以使用 COPY ... FROM STDIN 和 COPY ... TO STDOUT 了。

你可以查看 psycopg2 文档中的 COPY TO 和 COPY FROM 的使用方法。

不过你提到不能使用 psycopg2，那你就没办法了。驱动程序必须理解 COPY TO STDOUT 和 COPY FROM STDIN，才能使用这些功能，或者必须提供一种方法来直接写数据到网络连接上，这样你才能控制驱动程序的网络连接，自己实现 COPY 协议。要做到这一点，必须写特定于驱动程序的代码，不能仅仅依赖数据库API。

所以，虽然 khampson 的建议通常不是个好主意，但似乎是你唯一的选择。

（我发这个主要是为了确保那些没有限制使用 psycopg2 的人能做出明智的选择。）

如果你 必须 使用 psql，请注意：

• 使用 subprocess 模块 和 Popen 构造函数
• 给 psql 传递 -qAtX 和 -v ON_ERROR_STOP=1，这样可以确保批处理时的行为正常。
• 使用数组形式的命令，比如 ['psql', '-v', 'ON_ERROR_STOP=1', '-qAtX', '-c', '\copy mytable from stdin']，而不是使用 shell。
• 写入 psql 的标准输入，然后关闭它，等待 psql 完成。
• 记得处理命令失败时抛出的异常。让 subprocess 捕获错误输出，并将其包装在异常对象中。

这样做比旧式的 os.popen2 等方法更安全、更干净，也更容易出错。

Answer 2

你可以在你的脚本中调用 psql 这个命令行工具（也就是用 subprocess.call），然后利用它的 \copy 命令，把一个实例的输出直接传给另一个实例的输入，这样就不用临时文件了。比如：

psql -X -h from_host -U user -c "\copy from_table to stdout" | psql -X -h to_host -U user -c "\copy to_table from stdin"

这段话的意思是，假设目标数据库中已经有了你要的表。如果没有的话，你需要先用其他命令创建这个表。

另外，需要注意的是，这种方法有一个小问题，就是第一个 psql 调用中的错误可能会被传输过程给忽略掉。