Coverage for object_streams/postgres.py: 91%

1"""PostgreSQL wakeup helpers for object stream outbox rows.""" 

2 

3from __future__ import annotations 

4 

5import re 

6from collections.abc import Callable 

7from collections.abc import Iterator 

8 

9from django.conf import settings 

10from django.db import DEFAULT_DB_ALIAS 

11from django.db import connections 

12 

13 

14__all__ = ( 

15 "DEFAULT_NOTIFY_CHANNEL", 

16 "get_notify_channel", 

17 "listen_outbox_event_ids", 

18 "notify_outbox_event", 

19 "validate_notify_channel", 

20) 

21 

22 

23DEFAULT_NOTIFY_CHANNEL = "object_streams_events" 

24MAX_NOTIFY_CHANNEL_BYTES = 63 

25_CHANNEL_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") 

26 

27 

28def get_notify_channel() -> str: 

29 """Return the PostgreSQL NOTIFY channel used for outbox wakeups.""" 

30 

31 return str(getattr(settings, "OBJECT_STREAMS_NOTIFY_CHANNEL", DEFAULT_NOTIFY_CHANNEL)) 

32 

33 

34def validate_notify_channel(channel: str) -> str: 

35 """Validate a PostgreSQL notification channel name.""" 

36 

37 if not _CHANNEL_RE.fullmatch(channel): 

38 msg = "PostgreSQL notification channels must be unquoted identifier names." 

39 raise ValueError(msg) 

40 if len(channel.encode("utf-8")) > MAX_NOTIFY_CHANNEL_BYTES: 40 ↛ 41line 40 didn't jump to line 41 because the condition on line 40 was never true

41 msg = "PostgreSQL notification channels must be 63 bytes or shorter." 

42 raise ValueError(msg) 

43 return channel 

44 

45 

46def notify_outbox_event(event_id: int, *, using: str | None = None, channel: str | None = None) -> None: 

47 """Send a PostgreSQL wakeup notification for an outbox event id.""" 

48 

49 channel_name = validate_notify_channel(channel or get_notify_channel()) 

50 connection = connections[using or DEFAULT_DB_ALIAS] 

51 with connection.cursor() as cursor: 

52 cursor.execute("SELECT pg_notify(%s, %s)", [channel_name, str(event_id)]) 

53 

54 

55def listen_outbox_event_ids( 

56 *, 

57 using: str = DEFAULT_DB_ALIAS, 

58 channel: str | None = None, 

59 timeout: float | None = None, 

60 stop_after: int | None = None, 

61 on_listening: Callable[[], bool] | None = None, 

62) -> Iterator[int]: 

63 """Yield outbox event ids from PostgreSQL notifications.""" 

64 

65 channel_name = validate_notify_channel(channel or get_notify_channel()) 

66 connection = connections[using] 

67 connection.ensure_connection() 

68 raw_connection = connection.connection 

69 notifies = getattr(raw_connection, "notifies", None) 

70 if notifies is None: 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true

71 msg = "Object stream listening requires a psycopg 3 PostgreSQL connection." 

72 raise RuntimeError(msg) 

73 

74 previous_autocommit = connection.get_autocommit() 

75 connection.set_autocommit(True) 

76 try: 

77 with connection.cursor() as cursor: 

78 cursor.execute(f"LISTEN {channel_name}") 

79 

80 if on_listening is not None and on_listening(): 

81 return 

82 

83 for notification in notifies(timeout=timeout, stop_after=stop_after): 

84 if notification.channel != channel_name: 

85 continue 

86 try: 

87 yield int(notification.payload) 

88 except ValueError: 

89 continue 

90 finally: 

91 with connection.cursor() as cursor: 

92 cursor.execute(f"UNLISTEN {channel_name}") 

93 connection.set_autocommit(previous_autocommit)