psycopg2 使用多个线程或进程处理游标结果答案

【问题标题】：psycopg2 process cursor results with muliple threads or processespsycopg2 使用多个线程或进程处理游标结果
【发布时间】：2020-02-10 15:18:20
【问题描述】：

我有一个函数可以查询一个大表以对其进行索引...它创建一个名为“all_accounts”的服务器端游标。

def get_all_accounts(self):
    cursor = self.get_cursor('all_accounts')
    cursor.execute("SELECT * FROM account_summary LIMIT 20000;")

然后我一次处理这 2,000 个左右，以插入 NoSQL 解决方案：

def index_docs(self, cursor):
  while True:
    # consume result over a series of iterations
    # with each iteration fetching 2000 records
    record_count = cursor.rowcount
    records = cursor.fetchmany(size=2000)

    if not records:
        break

    for r in records:
        # do stuff

我希望 index_docs 函数能够并行 x10 使用游标 fetchmany() 调用，因为我的瓶颈不是由目标系统引起的，而是由我的脚本的单线程性质引起的。我过去做过一些异步/工作人员的事情，但是 psycopg2 光标似乎可能是个问题。想法？

【问题讨论】：

标签： python postgresql

【解决方案1】：

如果单个进程/线程访问游标并将工作分配给推送到另一个数据库的多个工作进程，我认为您将是安全的。（快速浏览一下，服务器端游标不能在连接之间共享，但我可能错了。）

就是这样。通常你会使用imap_unordered 来迭代单个项目的集合（并使用比默认值 1 更高的 chunksize），但我认为我们也可以在这里使用批次......

import multiprocessing

def get_batches(conn):
    cursor = conn.get_cursor('all_accounts')
    cursor.execute("SELECT * FROM account_summary LIMIT 20000;")
    while True:
        records = cursor.fetchmany(size=500)
        if not records:
            break
        yield list(records)


def process_batch(batch):
    # (this function is run in child processes)
    for r in batch:
        # ...
    return "some arbitrary result"


def main():
    conn = connect...()
    with multiprocessing.Pool() as p:
        batch_generator = get_batches(conn)
        for result in p.imap_unordered(process_batch, get_batches):
            print(result)  # doesn't really matter

【讨论】：