如何在 sqlalchemy 中使用 psycopg2.extras？答案

【问题标题】：How can I use psycopg2.extras in sqlalchemy?如何在 sqlalchemy 中使用 psycopg2.extras？
【发布时间】：2017-09-05 03:43:45
【问题描述】：

我想将大量条目（~600k）上传到 PostgreSQL 数据库中的一个简单表中，每个条目有一个外键、一个时间戳和 3 个浮点数。但是，每个条目执行here 描述的核心批量插入需要 60 毫秒，因此整个执行需要 10 小时。我发现这是executemany() 方法的性能问题，但是已经用psycopg2 2.7 中的execute_values() 方法解决了。

我运行的代码如下：

#build a huge list of dicts, one dict for each entry
engine.execute(SimpleTable.__table__.insert(),
               values) # around 600k dicts in a list

我发现这是一个常见问题，但是我还没有设法在 sqlalchemy 本身中找到解决方案。有没有办法告诉 sqlalchemy 在某些情况下调用execute_values()？有没有其他方法可以在不自己构建SQL语句的情况下实现大插入？

感谢您的帮助！

【问题讨论】：

我正要建议使用SimpleTable.__table__.insert().values(values)，它会编译成带有多个VALUES 元组的单个INSERT 语句，但事实证明它在我的机器上实际上更慢。编译本身与使用依赖于executemany() 的方法一样慢。

标签： python python-3.x sqlalchemy psycopg2

【解决方案1】：

同时，通过 create_engine() 函数上的 use_batch_mode 标志（从 SqlAlchemy 1.2.0 开始）成为可能。请参阅docs。它使用来自psycopg.extras 的execute_batch() 函数。

【讨论】：

他们进一步扩展了对快速执行助手的支持：“在 1.3.7 版中已更改：use_batch_mode 标志已被新参数 executemany_mode 取代，该参数同时提供支持用于 psycopg2 的 execute_batch 助手以及 execute_values 助手。” docs.sqlalchemy.org/en/13/dialects/…
从 1.4 开始，values_only 是默认值。 psycopg2 execute_values() 扩展用于限定 INSERT 语句，它重写 INSERT 以包含多个 VALUES 子句，以便可以使用一个语句插入多个参数集。

【解决方案2】：

这不是您正在寻找的答案，因为这不解决尝试指示 SQLAlchemy 使用 psycopg 附加功能，并且需要 - 有点 - 手动 SQL，但是：您可以从引擎访问底层 psycopg 连接raw_connection()，允许使用COPY FROM：

import io
import csv
from psycopg2 import sql

def bulk_copy(engine, table, values):
    csv_file = io.StringIO()
    headers = list(values[0].keys())
    writer = csv.DictWriter(csv_file, headers)
    writer.writerows(values)

    csv_file.seek(0)

    # NOTE: `format()` here is *not* `str.format()`, but
    # `SQL.format()`. Never use plain string formatting.
    copy_stmt = sql.SQL("COPY {} (" +
                        ",".join(["{}"] * len(headers)) +
                        ") FROM STDIN CSV").\
        format(sql.Identifier(str(table.name)),
               *(sql.Identifier(col) for col in headers))

    # Fetch a raw psycopg connection from the SQLAlchemy engine
    conn = engine.raw_connection()
    try:
        with conn.cursor() as cur:
            cur.copy_expert(copy_stmt, csv_file)

        conn.commit()

    except:
        conn.rollback()
        raise

    finally:
        conn.close()

然后

bulk_copy(engine, SimpleTable.__table__, values)

与执行 INSERT 语句相比，这应该快得多。在这台机器上移动 600,000 条记录大约需要 8 秒，约 13µs/条记录。您还可以将原始连接和光标与 extras 包一起使用。

【讨论】：