Repository: ansible/awx
File: awx/main/utils/pglock.py (Python)

Code before:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.

from contextlib import contextmanager

from django_pglocks import advisory_lock as django_pglocks_advisory_lock
from django.db import connection


@contextmanager
def advisory_lock(*args, lock_session_timeout_milliseconds=0, **kwargs):
    if connection.vendor == 'postgresql':
        cur = None
        idle_in_transaction_session_timeout = None
        idle_session_timeout = None
        if lock_session_timeout_milliseconds > 0:
            cur = connection.cursor()
            idle_in_transaction_session_timeout = cur.execute('SHOW idle_in_transaction_session_timeout').fetchone()[0]
            idle_session_timeout = cur.execute('SHOW idle_session_timeout').fetchone()[0]
            cur.execute(f"SET idle_in_transaction_session_timeout = {lock_session_timeout_milliseconds}")
            cur.execute(f"SET idle_session_timeout = {lock_session_timeout_milliseconds}")
        with django_pglocks_advisory_lock(*args, **kwargs) as internal_lock:
            yield internal_lock
            if lock_session_timeout_milliseconds > 0:
                cur.execute(f"SET idle_in_transaction_session_timeout = {idle_in_transaction_session_timeout}")
                cur.execute(f"SET idle_session_timeout = {idle_session_timeout}")
                cur.close()
    else:
        yield True


Review comment (suggestion):
in the case that we are closed, this will hit a `InterfaceError: the cursor is closed` error

in the case of task managers, we probably don't care because we are willing for the entire pid to go away, but if someone later on uses `lock_session_timeout_milliseconds` in another case there might be adverse affects

Code after:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.

from contextlib import contextmanager

from django_pglocks import advisory_lock as django_pglocks_advisory_lock
from django.db import connection


@contextmanager
def advisory_lock(*args, lock_session_timeout_milliseconds=0, **kwargs):
    if connection.vendor == 'postgresql':
        cur = None
        idle_in_transaction_session_timeout = None
        idle_session_timeout = None
        if lock_session_timeout_milliseconds > 0:
            with connection.cursor() as cur:
                idle_in_transaction_session_timeout = cur.execute('SHOW idle_in_transaction_session_timeout').fetchone()[0]
                idle_session_timeout = cur.execute('SHOW idle_session_timeout').fetchone()[0]
                cur.execute(f"SET idle_in_transaction_session_timeout = {lock_session_timeout_milliseconds}")
                cur.execute(f"SET idle_session_timeout = {lock_session_timeout_milliseconds}")
        with django_pglocks_advisory_lock(*args, **kwargs) as internal_lock:
            yield internal_lock
            if lock_session_timeout_milliseconds > 0:
                with connection.cursor() as cur:
                    cur.execute(f"SET idle_in_transaction_session_timeout = {idle_in_transaction_session_timeout}")
                    cur.execute(f"SET idle_session_timeout = {idle_session_timeout}")
    else:
        yield True