Coverage for object_streams/subscriptions.py: 86%

1"""Subscription request parsing and normalization.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Mapping 

6from collections.abc import Sequence 

7from dataclasses import dataclass 

8from dataclasses import field 

9from enum import StrEnum 

10from typing import Any 

11 

12 

13__all__ = ("ResyncRequired", "SubscriptionKind", "SubscriptionRequest") 

14 

15 

16class SubscriptionKind(StrEnum): 

17 OBJECT = "object" 

18 FILTER = "filter" 

19 MODEL = "model" 

20 

21 

22@dataclass(frozen=True, slots=True) 

23class ResyncRequired: 

24 """Subscription-level message for cases where precise replay is unavailable.""" 

25 

26 subscription_id: str 

27 cursor: int 

28 reason: str = "cursor_replay_unavailable" 

29 

30 def as_dict(self) -> dict[str, Any]: 

31 return { 

32 "type": "resync_required", 

33 "subscription_id": self.subscription_id, 

34 "cursor": self.cursor, 

35 "reason": self.reason, 

36 } 

37 

38 

39@dataclass(frozen=True, slots=True) 

40class SubscriptionRequest: 

41 """Normalized client subscription request.""" 

42 

43 kind: SubscriptionKind | str 

44 model: str 

45 pk: str | None = None 

46 filters: Mapping[str, Any] = field(default_factory=dict) 

47 search: str | None = None 

48 ordering: Sequence[str] = field(default_factory=tuple) 

49 shape: Mapping[str, Any] = field(default_factory=dict) 

50 cursor: int | None = None 

51 subscription_id: str | None = None 

52 

53 def __post_init__(self): 

54 kind = SubscriptionKind(str(self.kind)) 

55 object.__setattr__(self, "kind", kind) 

56 object.__setattr__(self, "filters", dict(self.filters)) 

57 object.__setattr__(self, "ordering", tuple(self.ordering)) 

58 object.__setattr__(self, "shape", dict(self.shape)) 

59 if self.pk is not None: 

60 object.__setattr__(self, "pk", str(self.pk)) 

61 if self.cursor is not None: 

62 cursor = int(self.cursor) 

63 if cursor < 0: 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true

64 msg = "Subscription cursors must be non-negative." 

65 raise ValueError(msg) 

66 object.__setattr__(self, "cursor", cursor) 

67 if kind == SubscriptionKind.OBJECT and self.pk is None: 

68 msg = "Object subscriptions require a primary key." 

69 raise ValueError(msg) 

70 

71 @classmethod 

72 def from_message(cls, message: Mapping[str, Any]) -> SubscriptionRequest: 

73 if message.get("op") != "subscribe": 73 ↛ 74line 73 didn't jump to line 74 because the condition on line 73 was never true

74 msg = "SubscriptionRequest only accepts subscribe messages." 

75 raise ValueError(msg) 

76 if not message.get("model"): 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true

77 msg = "Subscribe messages require a model label." 

78 raise ValueError(msg) 

79 

80 return cls( 

81 kind=message.get("kind", SubscriptionKind.OBJECT), 

82 model=str(message["model"]), 

83 pk=message.get("pk"), 

84 filters=message.get("filter") or message.get("filters") or {}, 

85 search=message.get("search"), 

86 ordering=message.get("ordering") or (), 

87 shape=message.get("shape") or {}, 

88 cursor=message.get("cursor"), 

89 subscription_id=message.get("subscription_id"), 

90 ) 

91 

92 def as_dict(self) -> dict[str, Any]: 

93 value: dict[str, Any] = { 

94 "op": "subscribe", 

95 "kind": str(self.kind), 

96 "model": self.model, 

97 "filter": dict(self.filters), 

98 } 

99 if self.pk is not None: 

100 value["pk"] = self.pk 

101 if self.search is not None: 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true

102 value["search"] = self.search 

103 if self.ordering: 

104 value["ordering"] = list(self.ordering) 

105 if self.shape: 105 ↛ 106line 105 didn't jump to line 106 because the condition on line 105 was never true

106 value["shape"] = dict(self.shape) 

107 if self.cursor is not None: 

108 value["cursor"] = self.cursor 

109 if self.subscription_id is not None: 

110 value["subscription_id"] = self.subscription_id 

111 return value