Skip to content

Commit d32e032

Browse files
committed
Method to list active pg locks
1 parent 2486153 commit d32e032

File tree

2 files changed

+210
-9
lines changed

2 files changed

+210
-9
lines changed

ansible_base/lib/utils/db.py

Lines changed: 95 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,10 @@
66

77
import psycopg
88
from django.conf import settings
9-
from django.db import DEFAULT_DB_ALIAS, OperationalError, connection, connections, transaction
9+
from django.db import DEFAULT_DB_ALIAS, OperationalError, connection, transaction
1010
from django.db.backends.postgresql.base import DatabaseWrapper as PsycopgDatabaseWrapper
1111
from django.db.migrations.executor import MigrationExecutor
12+
from django.db.transaction import get_connection
1213

1314
logger = logging.getLogger(__name__)
1415

@@ -40,6 +41,97 @@ def migrations_are_complete() -> bool:
4041
# that was licensed under the MIT license
4142

4243

44+
def string_to_advisory_lock_id(lock_string: str) -> int:
45+
"""Convert a string to a PostgreSQL advisory lock ID integer.
46+
47+
Generates an id within postgres integer range (-2^31 to 2^31 - 1).
48+
crc32 generates an unsigned integer in Py3, we convert it into
49+
a signed integer using 2's complement (this is a noop in Py2).
50+
51+
Args:
52+
lock_string: The string to convert to a lock ID
53+
54+
Returns:
55+
Integer lock ID suitable for PostgreSQL advisory locks
56+
"""
57+
pos = crc32(lock_string.encode("utf-8"))
58+
lock_id = (2**31 - 1) & pos
59+
if pos & 2**31:
60+
lock_id -= 2**31
61+
return lock_id
62+
63+
64+
def advisory_lock_id_to_debug_info(lock_id: int) -> dict:
65+
"""Convert an advisory lock ID back to debug information.
66+
67+
Note: This cannot reverse the string due to CRC32 hash collisions,
68+
but provides debugging information about the lock ID.
69+
70+
Args:
71+
lock_id: The integer lock ID
72+
73+
Returns:
74+
Dictionary with debug information about the lock ID
75+
"""
76+
# Convert back to unsigned for analysis
77+
if lock_id < 0:
78+
unsigned_value = lock_id + 2**31
79+
had_high_bit = True
80+
else:
81+
unsigned_value = lock_id
82+
had_high_bit = False
83+
84+
return {
85+
'lock_id': lock_id,
86+
'unsigned_crc32': unsigned_value,
87+
'had_high_bit_set': had_high_bit,
88+
'hex_representation': hex(lock_id),
89+
}
90+
91+
92+
def get_active_advisory_locks(using=None) -> list:
93+
"""Get a list of all currently held advisory locks.
94+
95+
Args:
96+
using: Database alias to use (defaults to DEFAULT_DB_ALIAS)
97+
98+
Returns:
99+
List of dictionaries containing lock information
100+
"""
101+
if using is None:
102+
using = DEFAULT_DB_ALIAS
103+
104+
conn = get_connection(using)
105+
if conn.vendor != "postgresql":
106+
return []
107+
108+
with conn.cursor() as cursor:
109+
cursor.execute(
110+
"""
111+
SELECT locktype, classid, objid, objsubid, pid, mode, granted
112+
FROM pg_locks
113+
WHERE locktype = 'advisory'
114+
"""
115+
)
116+
117+
locks = []
118+
for row in cursor.fetchall():
119+
locktype, classid, objid, objsubid, pid, mode, granted = row
120+
locks.append(
121+
{
122+
'locktype': locktype,
123+
'classid': classid,
124+
'objid': objid,
125+
'objsubid': objsubid,
126+
'pid': pid,
127+
'mode': mode,
128+
'granted': granted,
129+
}
130+
)
131+
132+
return locks
133+
134+
43135
@contextmanager
44136
def django_pglocks_advisory_lock(lock_id, shared=False, wait=True, using=None):
45137

@@ -81,13 +173,7 @@ def django_pglocks_advisory_lock(lock_id, shared=False, wait=True, using=None):
81173

82174
tuple_format = True
83175
elif isinstance(lock_id, str):
84-
# Generates an id within postgres integer range (-2^31 to 2^31 - 1).
85-
# crc32 generates an unsigned integer in Py3, we convert it into
86-
# a signed integer using 2's complement (this is a noop in Py2)
87-
pos = crc32(lock_id.encode("utf-8"))
88-
lock_id = (2**31 - 1) & pos
89-
if pos & 2**31:
90-
lock_id -= 2**31
176+
lock_id = string_to_advisory_lock_id(lock_id)
91177
elif not isinstance(lock_id, int):
92178
raise ValueError("Cannot use %s as a lock id" % lock_id)
93179

@@ -104,7 +190,7 @@ def django_pglocks_advisory_lock(lock_id, shared=False, wait=True, using=None):
104190
acquire_params = (function_name,) + params
105191

106192
command = base % acquire_params
107-
cursor = connections[using].cursor()
193+
cursor = get_connection(using).cursor()
108194

109195
cursor.execute(command)
110196

test_app/tests/lib/utils/test_db.py

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,14 @@
1010

1111
from ansible_base.lib.utils.db import (
1212
advisory_lock,
13+
advisory_lock_id_to_debug_info,
14+
get_active_advisory_locks,
1315
get_pg_notify_params,
1416
migrations_are_complete,
1517
psycopg_conn_string_from_settings_dict,
1618
psycopg_connection_from_django,
1719
psycopg_kwargs_from_settings_dict,
20+
string_to_advisory_lock_id,
1821
)
1922

2023

@@ -148,6 +151,118 @@ def test_psycopg_connection_from_django_new_conn(self):
148151
assert isinstance(psycopg_connection_from_django(), psycopg.Connection)
149152

150153

154+
class TestStringToAdvisoryLockId:
155+
"""Test the string to advisory lock ID conversion function.
156+
157+
Generated by Claude Code (Sonnet 4)
158+
"""
159+
160+
def test_string_to_advisory_lock_id_basic(self):
161+
"""Test basic string to lock ID conversion."""
162+
lock_id = string_to_advisory_lock_id("test_string")
163+
assert isinstance(lock_id, int)
164+
assert -(2**31) <= lock_id <= 2**31 - 1
165+
166+
def test_string_to_advisory_lock_id_consistency(self):
167+
"""Test that the same string always produces the same lock ID."""
168+
test_string = "consistent_test"
169+
lock_id1 = string_to_advisory_lock_id(test_string)
170+
lock_id2 = string_to_advisory_lock_id(test_string)
171+
assert lock_id1 == lock_id2
172+
173+
def test_string_to_advisory_lock_id_different_strings(self):
174+
"""Test that different strings produce different lock IDs."""
175+
lock_id1 = string_to_advisory_lock_id("string1")
176+
lock_id2 = string_to_advisory_lock_id("string2")
177+
assert lock_id1 != lock_id2
178+
179+
def test_string_to_advisory_lock_id_unicode(self):
180+
"""Test string to lock ID conversion with unicode characters."""
181+
lock_id = string_to_advisory_lock_id("test_🔒_unicode")
182+
assert isinstance(lock_id, int)
183+
assert -(2**31) <= lock_id <= 2**31 - 1
184+
185+
186+
class TestAdvisoryLockIdToDebugInfo:
187+
"""Test the advisory lock ID to debug info function.
188+
189+
Generated by Claude Code (Sonnet 4)
190+
"""
191+
192+
def test_advisory_lock_id_to_debug_info_positive(self):
193+
"""Test debug info for positive lock ID."""
194+
lock_id = 12345
195+
debug_info = advisory_lock_id_to_debug_info(lock_id)
196+
197+
assert debug_info['lock_id'] == lock_id
198+
assert debug_info['unsigned_crc32'] == lock_id
199+
assert debug_info['had_high_bit_set'] is False
200+
assert debug_info['hex_representation'] == hex(lock_id)
201+
202+
def test_advisory_lock_id_to_debug_info_negative(self):
203+
"""Test debug info for negative lock ID."""
204+
lock_id = -12345
205+
debug_info = advisory_lock_id_to_debug_info(lock_id)
206+
207+
assert debug_info['lock_id'] == lock_id
208+
assert debug_info['unsigned_crc32'] == lock_id + 2**31
209+
assert debug_info['had_high_bit_set'] is True
210+
assert debug_info['hex_representation'] == hex(lock_id)
211+
212+
def test_advisory_lock_id_to_debug_info_roundtrip(self):
213+
"""Test debug info for a string-generated lock ID."""
214+
test_string = "test_debug_roundtrip"
215+
lock_id = string_to_advisory_lock_id(test_string)
216+
debug_info = advisory_lock_id_to_debug_info(lock_id)
217+
218+
assert debug_info['lock_id'] == lock_id
219+
assert isinstance(debug_info['unsigned_crc32'], int)
220+
assert isinstance(debug_info['had_high_bit_set'], bool)
221+
assert debug_info['hex_representation'] == hex(lock_id)
222+
223+
224+
class TestGetActiveAdvisoryLocks(SkipIfSqlite):
225+
"""Test the get active advisory locks function.
226+
227+
Generated by Claude Code (Sonnet 4)
228+
"""
229+
230+
@pytest.mark.django_db
231+
def test_get_active_advisory_locks_empty(self):
232+
"""Test getting active locks when none are held."""
233+
locks = get_active_advisory_locks()
234+
assert isinstance(locks, list)
235+
# We can't guarantee no locks since other tests might be running
236+
237+
@pytest.mark.django_db
238+
def test_get_active_advisory_locks_with_lock(self):
239+
"""Test getting active locks when we hold one."""
240+
test_lock_name = "test_get_active_locks"
241+
242+
with advisory_lock(test_lock_name):
243+
locks = get_active_advisory_locks()
244+
assert isinstance(locks, list)
245+
# Should have at least our lock
246+
assert len(locks) >= 1
247+
248+
# Check that lock entries have expected structure
249+
for lock in locks:
250+
assert 'locktype' in lock
251+
assert 'classid' in lock
252+
assert 'objid' in lock
253+
assert 'objsubid' in lock
254+
assert 'pid' in lock
255+
assert 'mode' in lock
256+
assert 'granted' in lock
257+
assert lock['locktype'] == 'advisory'
258+
259+
@pytest.mark.django_db
260+
def test_get_active_advisory_locks_sqlite_returns_empty(self):
261+
"""Test that SQLite returns empty list."""
262+
# This test will be skipped by SkipIfSqlite for completeness
263+
pass
264+
265+
151266
class TestAdvisoryLock(SkipIfSqlite):
152267
THREAD_WAIT_TIME = 0.1
153268

0 commit comments

Comments
 (0)