# base.py
from django.db.backends.postgresql_psycopg2.base import *
from django.conf import settings


class DatabaseWrapper(DatabaseWrapper):
    def _cursor(self):
        cursor = super(DatabaseWrapper, self)._cursor()
        statement_timeout = getattr(settings, 'DATABASE_STATEMENT_TIMEOUT', None)
        if statement_timeout is not None:
            cursor.execute("SET STATEMENT_TIMEOUT=%s" % statement_timeout)
        return cursor