在进程池中共享数据库连接

Question

我有一个 Python 3 程序，它根据行的 id（在 Postgres 9.5 数据库中的表中）更新一个大的行列表。

我使用多处理来加速这个过程。由于 Psycopg 的连接不能跨进程共享，我为每一行创建一个连接，然后关闭它。

总体而言，多处理比单处理快（8 个 CPU 时快 5 倍）。但是，创建连接很慢：我只想创建几个连接，并在需要时保持打开状态。

由于 .map() 将 ids_list 切割成多个块并提交给进程池，是否可以为同一块/进程中的所有 id 共享数据库连接？

示例代码：

from multiprocessing import Pool
import psycopg2


def create_db_connection():
    conn = psycopg2.connect(database=database,
                            user=user,
                            password=password,
                            host=host)
    return conn


def my_function(item_id):
    conn = create_db_connection()

    # Other CPU-intensive operations are done here

    cur = conn.cursor()
    cur.execute("""
        UPDATE table
        SET
        my_column = 1
        WHERE id = %s;
        """,
        (item_id, ))
    cur.close()
    conn.commit()


if __name__ == '__main__':
    ids_list = []  # Long list of ids

    pool = Pool()  # os.cpu_count() processes
    pool.map(my_function, ids_list)

感谢您的任何帮助，您可以提供。

Answer 1

您可以使用Pool 构造函数的初始值设定项参数。在初始化函数中设置 DB 连接。也许将连接凭据作为参数传递。

查看文档：https://docs.python.org/3/library/multiprocessing.html#module-multiprocessing.pool