python - 为什么当我将光标传递给 StreamingHttpResponse 时，我的光标在我的生成器函数中关闭了？

Question

我有一个可能很大的查询集，我不想加载到内存中。这是一条自定义 SQL 语句。

django.http.StreamingHttpResponse 采用 iterator(生成器)参数。为避免将所有内容加载到内存中，我使用了 Postgres 服务器端游标和 fetchmany(尽管我尚未验证这是否真的有效)。

这是我传递的生成器函数:

def queryset_generator(cursor, chunk_size=CHUNK_SIZE):
    while True:
        if cursor.closed:
            yield "cursor closed!"
            break
        rows = cursor.fetchmany(chunk_size)
        if not rows:
            break
        yield rows

我测试游标是否关闭，否则 psycopg2 会在之后的代码尝试访问关闭的游标时提示。

下面是我在 View 中传递它的方式(简化 SQL):

with connections['mydb'].cursor() as cursor:
    cursor.execute("SELECT * FROM foobar;")
    return StreamingHttpResponse(queryset_generator(cursor))

这始终给了我

cursor closed!

为什么游标在我的生成器函数中关闭？如果我这样做，我认为效果很好:

with connections['mydb'].cursor() as cursor:
    cursor.execute("SELECT * FROM foobar;")
    return StreamingHttpResponse(cursor.fetchall())

还可能值得注意的是，这在 shell 中运行良好:

cursor = connections['mydb'].cursor()
cursor.execute(...)
for x in StreamingHttpResponse(queryset_generator(cursor))._iterator:
    print(x)

Answer 1

Why is the cursor closed in my generator function?

因为您使用return 退出上下文管理器:

return StreamingHttpResponse(queryset_generator(cursor))

这会退出 with block ，触发上下文管理器上的 __exit__ 方法，然后此方法关闭游标。 with 语句或上下文管理器无法知道您刚刚将对 cursor 对象的引用传递给了其他仍然需要它保持打开状态的对象。 with 不关心引用，只关心语义 block 结束。

如果您需要在 StreamingHttpResponse() 实例完成数据流之前保持游标打开，则不能在 return 语句周围使用上下文管理器。

在上下文管理器中不使用游标，并让queryset_generator()函数负责使用with代替:

def queryset_generator(cursor, chunk_size=CHUNK_SIZE):
    with cursor:
        while True:
            if cursor.closed:
                yield "cursor closed!"
                break
            rows = cursor.fetchmany(chunk_size)
            if not rows:
                break
            yield rows

和

cursor = connections['mydb'].cursor()
cursor.execute("SELECT * FROM foobar;")
return StreamingHttpResponse(queryset_generator(cursor))

现在游标保持打开状态，直到 queryset_generator() 中的 while 循环完成。