Server Side Cursors for Django's psycopg2 Backend

Author:: ryanbutterfield
Posted:: March 2, 2011
Language:: Python
Version:: 1.2
Score:: 2 (after 2 ratings)

Download
Raw

Based on discussion on http://thebuild.com/blog/2010/12/14/using-server-side-postgresql-cursors-in-django/ and http://thebuild.com/blog/2010/12/13/very-large-result-sets-in-django-using-postgresql/ but instead implements them via extending the psycopg2 backend which allows you to use all of django's machinery without having to resort to using raw cursors.

Usage:

qs = Model.objects.all()
with server_side_cursors(qs, itersize=100):
    for item in qs.iterator():
        item.value
        if random_reason_to_break:
            break

Setup:

In your own project create the package hierarchy myproject.db.backends.postgresql_psycopg2 and place the code in base.py.

In your settings.py set the database ENGINE to be 'myproject.db.backends.postgresql_psycopg2'.

If you using south you'll have to let it know its a postgresql_psycopg2 backend by adding to SOUTH_DATABASE_ADAPTERS (see south documentation).

Note:

Make sure your using psycopg >= 2.4 for efficient named (server side) cursor iteration.

import uuid

from django.db.backends.postgresql_psycopg2.base import *
from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper as BaseDatabaseWrapper

class server_side_cursors(object):
    """
    With block helper that enables and disables server side cursors.
    """
    def __init__(self, qs_or_using_or_connection, itersize=None):
        from django.db import connections
        from django.db.models.query import QuerySet
        
        self.itersize = itersize
        if isinstance(qs_or_using_or_connection, QuerySet):
            self.connection = connections[qs_or_using_or_connection.db]
        elif isinstance(qs_or_using_or_connection, basestring):
            self.connection = connections[qs_or_using_or_connection]
        else:
            self.connection = qs_or_using_or_connection
    
    def __enter__(self):
        self.connection.server_side_cursors = True
        self.connection.server_side_cursor_itersize = self.itersize
    
    def __exit__(self, type, value, traceback):
        self.connection.server_side_cursors = False
        self.connection.server_side_cursor_itersize = None

class DatabaseWrapper(BaseDatabaseWrapper):
    """
    Psycopg2 database backend that allows the use of server side cursors.
    
    Usage:
    
    qs = Model.objects.all()
    with server_side_cursors(qs, itersize=x):
        for item in qs.iterator():
            item.value
    
    """
    def __init__(self, *args, **kwargs):
        self.server_side_cursors = False
        self.server_side_cursor_itersize = None
        super(DatabaseWrapper, self).__init__(*args, **kwargs)
    
    def _cursor(self):
        """
        Returns a unique server side cursor if they are enabled, 
        otherwise falls through to the default client side cursors.
        """
        if self.server_side_cursors:
            # intialise the connection if we haven't already
            # this will waste a client side cursor, but only on the first call
            if self.connection is None:
                super(DatabaseWrapper, self)._cursor()
            
            # give the cursor a unique name which will invoke server side cursors
            cursor = self.connection.cursor(name='cur%s' % str(uuid.uuid4()).replace('-', ''))
            cursor.tzinfo_factory = None
            
            if self.server_side_cursor_itersize is not None:
                cursor.itersize = self.server_side_cursor_itersize
            
            return CursorWrapper(cursor)
        
        return super(DatabaseWrapper, self)._cursor()

Comments

cwisecarver (on May 26, 2016):

This doesn't work, at least, on 1.9. I modified it a little bit and got it working.

Here's a gist, the markdown formatting on this site was too much for me to bear.

Please login first before commenting.

Server Side Cursors for Django's psycopg2 Backend

More like this

Comments