psycopg2 connection retry monkey patch

import psycopg2

original_connect = psycopg2.connect


def connect_with_retry(*args, **kwargs):
    retries = 3

    while True:
        retries -= 1
        try:
            return original_connect(*args, **kwargs)
        except psycopg2.OperationalError as exc:
            # Unfortunately exc.pgcode is None for connect_timeout.
            if exc.args != (u'timeout expired\n',):
                raise

            if retries <= 0:
                raise


def patch_psycopg2():
    psycopg2.connect = connect_with_retry


patch_psycopg2()