pg_lock.py

"""
The module implements PostgreSQL Advisory Locks.
https://www.postgresql.org/docs/10/static/functions-admin.html#FUNCTIONS-ADVISORY-LOCKS
"""

from zlib import crc32

from django.db import connection


def _encode_lock_id(lock_id):
    pos = crc32(lock_id.encode('utf-8'))
    lock_id = (2 ** 31 - 1) & pos
    if pos & 2 ** 31:
        lock_id -= 2 ** 31
    return lock_id


def _pg_lock(func_name, lock_id, fetch=False):
    lock_id = _encode_lock_id(lock_id)
    query = "SELECT {}({});".format(func_name, lock_id)
    cursor = connection.cursor()
    cursor.execute(query)
    if fetch:
        return cursor.fetchone()[0]


def pg_advisory_xact_lock(lock_id):
    func_name = 'pg_advisory_xact_lock'
    return _pg_lock(func_name, lock_id)


def pg_try_advisory_xact_lock(lock_id):
    func_name = 'pg_try_advisory_xact_lock'
    return _pg_lock(func_name, lock_id, fetch=True)