postgres_paginated_cursor.py

from psycopg2.extras import DictCursor
from django.db import connections

def get_cursor(alias='default', cursor_factory=None):
    """map from django's ORM layer to the raw DB cursor."""

    wrapped_conn = connections[alias]

    # hack to ensure connection is immediately opened:
    if wrapped_conn.connection is None:
        cursor = wrapped_conn.cursor()

    raw_conn = wrapped_conn.connection

    return raw_conn.cursor(cursor_factory=cursor_factory)

def exec_to_dict(sql, params=(), alias='default'):
    """
    Variable substitution under psycopg db driver are %s and positional (only).

    exec_to_dict('select 1 as x from a where a.b = %s', (3,))
      is equivalent to "select 1 from a where a.b = 3"
      and would return [{'x': 1}]
    """

    cursor = get_cursor(alias, cursor_factory=DictCursor)
    cursor.execute(sql, params)

    # yield the results
    return iter_cursor_results(cursor)

def iter_cursor_results(cursor, fetch_size=1000):
    """An iterator that uses fetchmany to keep memory usage down"""
    while True:
        results = cursor.fetchmany(fetch_size)
        if not results:
            break
        for result in results:
            yield result