如何使用 SQLAlchemy 将数据直接从一个 postgresql 数据库传输到另一个？

Question

在 ETL 过程中，我想定期查询数据库“A”（例如，所有时间戳大于程序上次运行的行）并将该数据移动到数据库“B”中以进行进一步处理。两者都是 PostgreSQL 数据库。我想在 Python 脚本中执行此数据传输，使用 SQLAlchemy 连接到两个数据库。什么是最不混乱、最不脆弱的方法？

我知道 Postgres 的 COPY TO 和 COPY FROM 命令允许通过中间文件 (see here) 将表行和查询结果从一个数据库服务器传输到另一个数据库服务器。从 Unix 命令行，您甚至可以通过管道将数据库 A 的输出作为输入到数据库 B，而无需潜在的大中间文件 (see excellent instructions here)。我想知道的是如何使用两个 SQLAlchemy 连接在 Python 脚本中完成最后一个技巧，而不是使用 subprocess 运行 shell 命令。

import sqlalchemy
dbA = sqlalchemy.create_engine(connection_string_A)
dbB = sqlalchemy.create_engine(connection_string_B)

# how do I do this part?
dbA.execute('SELECT (column) FROM widgets...') # somehow pipe output into...
dbB.execute('INSERT INTO widgets (column) ...') # without holding lots of data in memory or on disk

为了记录，我现在没有使用 SQLAlchemy 的任何 ORM 功能，只是裸 SQL 查询。

Answer 1

您在问题中询问了两个不同的问题。一个是如何将 CSV 从COPY FROM 传送到COPY TO；另一个是如何将行从SELECT 查询传递到INSERT。

将SELECT 查询中的行通过管道传输到INSERT 是一种谎言，因为虽然您可以从SELECT 查询中流式传输行，但您不能将行流式传输到INSERT，所以您'必须批量执行多个INSERTs。由于INSERTs，这种方法具有很高的开销，但由于往返 CSV 导致的数据丢失问题较少。我将重点介绍为什么将 CSV 从 COPY FROM 管道传输到 COPY TO 很棘手，以及如何完成它。

psycopg2 允许您通过（同步）copy_expert 函数执行COPY 命令。它要求您为COPY FROM 传入一个可读文件对象，为COPY TO 传入一个可写文件对象。要完成您所描述的，您需要两个单独的线程来运行这两个命令中的每一个，一个带有write() 方法的文件对象，如果COPY FROM 命令无法跟上，则该方法会阻塞，以及一个带有@987654341 的文件对象@ 方法在 COPY TO 命令跟不上时阻塞。这是一个典型的生产者-消费者问题，很难解决。

这是我快速编写的一个（Python 3）。它可能充满了错误。如果您发现死锁，请告诉我（欢迎编辑）。

from threading import Lock, Condition, Thread


class Output(object):
    def __init__(self, pipe):
        self.pipe = pipe

    def read(self, count):
        with self.pipe.lock:
            # wait until pipe is still closed or buffer is not empty
            while not self.pipe.closed and len(self.pipe.buffer) == 0:
                self.pipe.empty_cond.wait()

            if len(self.pipe.buffer) == 0:
                return ""

            count = max(count, len(self.pipe.buffer))
            res, self.pipe.buffer = \
                self.pipe.buffer[:count], self.pipe.buffer[count:]
            self.pipe.full_cond.notify()
        return res

    def close(self):
        with self.pipe.lock:
            self.pipe.closed = True
            self.pipe.full_cond.notify()


class Input(object):
    def __init__(self, pipe):
        self.pipe = pipe

    def write(self, s):
        with self.pipe.lock:
            # wait until pipe is closed or buffer is not full
            while not self.pipe.closed \
                    and len(self.pipe.buffer) > self.pipe.bufsize:
                self.pipe.full_cond.wait()

            if self.pipe.closed:
                raise Exception("pipe closed")

            self.pipe.buffer += s
            self.pipe.empty_cond.notify()

    def close(self):
        with self.pipe.lock:
            self.pipe.closed = True
            self.pipe.empty_cond.notify()


class FilePipe(object):
    def __init__(self, bufsize=4096):
        self.buffer = b""
        self.bufsize = 4096
        self.input = Input(self)
        self.output = Output(self)
        self.lock = Lock()
        self.full_cond = Condition(self.lock)
        self.empty_cond = Condition(self.lock)
        self.closed = False

使用示例：

def read_thread(conn, f):
    conn.cursor().copy_expert("COPY foo TO STDIN;", f)
    f.close()
    conn.close()

engine.execute(
    "CREATE TABLE foo(id int);"
    "CREATE TABLE bar(id int);"
    "INSERT INTO foo (SELECT generate_series(1, 100000) AS id);"
    "COMMIT;")
input_conn = engine.raw_connection()
output_conn = engine.raw_connection()
pipe = FilePipe()

t = Thread(target=read_thread, args=(input_conn, pipe.input))
t.start()
output_cur = output_conn.cursor()
output_cur.copy_expert("COPY bar FROM STDIN;", pipe.output)
output_conn.commit()
output_conn.close()
t.join()

print(list(engine.execute("SELECT count(*) FROM bar;")))  # 100000

Answer 2

如果数据不是很大（可以保存在单个主机的主内存中），你可以试试我的基于pandas/python3/sqlalchemy的开源ETL工具包，bailaohe/parade，我提供了一个tutorial .您可以使用 pandas 对数据进行转换并直接返回结果数据框。只需稍加配置，即可将 pandas 数据帧转储到不同的目标连接。

对于你的问题，你可以使用 parade 生成一个简单的 sql 类型的任务，如下：

# -*- coding:utf-8 -*-
from parade.core.task import SqlETLTask
from parade.type import stdtypes


class CopyPostgres(SqlETLTask):

    @property
    def target_conn(self):
        """
        the target connection to write the result
        :return:
        """
        return 'target_postgres'

    @property
    def source_conn(self):
        """
        the source connection to write the result
        :return:
        """
        return 'source_postgres'

    @property
    def etl_sql(self):
        """
        the single sql statement to process etl
        :return:
        """
        return """SELECT (column) FROM widgets"""

您甚至可以组成一个包含多个任务的 DAG 工作流，并直接使用 Parade 安排工作流。希望这会有所帮助。