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

时间:2016-10-15 21:32:00

标签: python postgresql multiprocessing psycopg2

我有一个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)

感谢您提供的任何帮助。

1 个答案:

答案 0 :(得分:1)

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

查看文档:{​​{3}}