Почему мой курсор закрывается в моей функции генератора, когда я передаю его StreamingHttpResponse?

У меня есть потенциально большой набор запросов, который я не хочу загружать в память. Это пользовательская инструкция 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 жалуется, когда код впоследствии пытается получить доступ закрытый курсор.

И вот как я передаю это в моем представлении (упрощая SQL): 

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

Что последовательно дает мне

Курсор закрыт!

Почему курсор закрыт в моей функции генератора? Если я делаю это, на мой взгляд, это работает нормально: 

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

Также потенциально примечательно, что это прекрасно работает в оболочке: 

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

1 ответ:

Почему курсор закрыт в моей функции генератора?

Потому что вы выходите из контекстного менеджера с помощью return: 

return StreamingHttpResponse(queryset_generator(cursor))

Это завершает работу блока with, запускает метод __exit__ в контекстном менеджере, и этот метод закрывает курсор. Оператор with или менеджер контекста не могут знать, что вы только что передали ссылку на Объект cursor чему-то другому, что все еще нуждается в том, чтобы он оставался открытым. with не заботится о ссылках, только о семантическом блоке окончание.

Если вам нужно держать курсор открытым до тех пор, пока экземпляр 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))
Теперь курсор остается открытым до тех пор, пока не будет выполнен цикл while в queryset_generator().
3