2020-08-21

Postgres advisory locks for python developers.

Postgres docs for locking in general, advisory locks in specific.

Everyone should use postgres' advisory locks, they're great!

Why? What are they? Who cares?

With an advisory lock, you can get a lock on an arbitrary integer for the duration of a transaction. This is very useful as you can:

Let's implemement advisory locks in python with SQLAlchemy.

def acquire_advisory_lock(session: Session, *ids: Any) -> None:
    # make an integer key
    key_str = "-".join([str(id_) for id_ in ids])
    key_bytes: bytes = key_str.encode("utf-8")
    m = hashlib.sha256()
    m.update(key_bytes)
    # pg_try_advisory_xact_lock is limited to an 8-byte signed integer
    key = int.from_bytes(m.digest()[:8], byteorder="big", signed=True)

    # get a lock on the db with the key
    rows = session.execute(sql.select([func.pg_try_advisory_xact_lock(key)]))
    locked = not next(rows)[0]

    # if it is already locked by another transaction, raise an error
    if locked:
        raise AdvisoryLockNotAvailable(key)

    logger.debug("Acquired lock on key: %s", key)

For good measure, here's an example of the kind of test one would write against said function:

@pytest.fixture
def two_sessions(engine):
    Base.metadata.create_all(bind=engine)
    session_1 = sessionmaker(bind=engine)()
    session_2 = sessionmaker(bind=engine)()
    yield session_1, session_2
    # make sure we're cleaned up
    session_1.close()
    session_2.close()


def test_acquire_advisory_lock(two_sessions):
    session_1, session_2 = two_sessions

    # lock with session_1
    acquire_advisory_lock(session_1, 123)
    # show that we can lock it again in the same session and it doesn't blow up
    acquire_advisory_lock(session_1, 123)

    # try and fail to lock with session_2
    with pytest.raises(AdvisoryLockNotAvailable):
        acquire_advisory_lock(session_2, 123)

    # commit session_1
    session_1.commit()

    # now we should be able to acquire the lock with session_2
    acquire_advisory_lock(session_2, 123)

Easy!