DB Timeouts from Django App

db_timeouts.py

from django.conf import settings

# NOTE
# The timeout values here only restrict the roles from application code. The actual timeout set in the DB could
# be a different value.
# Timeout set in DB on 27/03/2018 is 50s.

DEFAULT_DB_TIMEOUT_IN_MS = 50000

default_conn = settings.DJANGO_DEFAULT_DB_CONNECTION_NAME
explorer_conn = settings.EXPLORER_CONNECTION_NAME
slave_conn = settings.SLAVE_NAME


DB_IDENTIFIER_AND_CONNECTION_TO_TIMEOUT_MAP = {
    settings.PROD_APP_DB_TIMEOUT_IDENTIFIER: {
        default_conn: 20000,
        explorer_conn: 50000,
        slave_conn: 20000
    },
    settings.PROD_CELERY_DB_TIMEOUT_IDENTIFIER: {
        default_conn: 50000,
        explorer_conn: 50000,
        slave_conn: 50000
    },
    settings.DEFAULT_DB_TIMEOUT_IDENTIFIER: {
        default_conn: 50000,
        explorer_conn: 50000,
        slave_conn: 50000
    },
}

models.py

from django.db.backends.signals import connection_created

# ...

def set_timeout_on_new_conn(sender, connection, **kwargs):
    """
        Rig django to set statement timeout for each new connection based on the config
    """
    try:
        timeout_to_set = DB_IDENTIFIER_AND_CONNECTION_TO_TIMEOUT_MAP[settings.DB_TIMEOUT_IDENTIFIER][connection.alias]
    except KeyError as e:
        logger.error("KeyError while setting DB Timeout: {0}".format(e))
        timeout_to_set = DEFAULT_DB_TIMEOUT_IN_MS

    with connection.cursor() as cursor:
        cursor.execute("set statement_timeout={0}".format(timeout_to_set))


connection_created.connect(set_timeout_on_new_conn)

settings.py

# ...

PROD_APP_DB_TIMEOUT_IDENTIFIER = 'db_timeout:production:app'
PROD_CELERY_DB_TIMEOUT_IDENTIFIER = 'db_timeout:production:celery'
PROD_CRON_DB_TIMEOUT_IDENTIFIER = 'db_timeout:production:cron'
PROD_BIZ_API_DB_TIMEOUT_IDENTIFIER = 'db_timeout:production:biz'
PROD_ADMIN_DB_TIMEOUT_IDENTIFIER = 'db_timeout:production:admin'
PROD_EXPLORER_DB_TIMEOUT_IDENTIFIER = 'db_timeout:production:explorer'
STAGING_DB_TIMEOUT_IDENTIFIER = 'db_timeout:staging:*'
DEFAULT_DB_TIMEOUT_IDENTIFIER = 'db_timeout:*:*'

# ...

DB_TIMEOUT_IDENTIFIER = PROD_CELERY_DB_TIMEOUT_IDENTFIER

# ...