Coverage for polar/eventstream/service.py: 71%

1from typing import Any 

2from uuid import UUID 

3 

4import structlog 

5from pydantic import BaseModel 

6 

7from polar.kit.utils import generate_uuid 

8from polar.logging import Logger 

9from polar.postgres import AsyncSession 

10from polar.redis import Redis 

11from polar.user_organization.service import ( 

12 user_organization as user_organization_service, 

13) 

14from polar.worker import enqueue_job 

15 

16log: Logger = structlog.get_logger() 

17 

18 

19class Receivers(BaseModel): 

20 user_id: UUID | None = None 

21 organization_id: UUID | None = None 

22 checkout_client_secret: str | None = None 

23 customer_id: UUID | None = None 

24 

25 def generate_channel_name(self, scope: str, resource_id: UUID | str) -> str: 

26 return f"{scope}:{resource_id}" 

27 

28 def get_channels(self) -> list[str]: 

29 channels = [] 

30 if self.user_id: 

31 channels.append(self.generate_channel_name("user", self.user_id)) 

32 

33 if self.organization_id: 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 channels.append(self.generate_channel_name("org", self.organization_id)) 

35 

36 if self.checkout_client_secret: 

37 channels.append( 

38 self.generate_channel_name("checkout", self.checkout_client_secret) 

39 ) 

40 

41 if self.customer_id: 41 ↛ 42line 41 didn't jump to line 42 because the condition on line 41 was never true

42 channels.append(self.generate_channel_name("customer", self.customer_id)) 

43 

44 return channels 

45 

46 

47class Event(BaseModel): 

48 id: UUID 

49 key: str 

50 payload: dict[str, Any] 

51 

52 

53async def send_event(redis: Redis, event_json: str, channels: list[str]) -> None: 

54 for channel in channels: 

55 await redis.publish(channel, event_json) 

56 log.debug( 

57 "Published event to eventstream", event_json=event_json, channels=channels 

58 ) 

59 

60 

61async def publish( 

62 key: str, 

63 payload: dict[str, Any], 

64 user_id: UUID | None = None, 

65 organization_id: UUID | None = None, 

66 checkout_client_secret: str | None = None, 

67 customer_id: UUID | None = None, 

68) -> None: 

69 receivers = Receivers( 

70 user_id=user_id, 

71 organization_id=organization_id, 

72 checkout_client_secret=checkout_client_secret, 

73 customer_id=customer_id, 

74 ) 

75 channels = receivers.get_channels() 

76 event = Event( 

77 id=generate_uuid(), 

78 key=key, 

79 payload=payload, 

80 ).model_dump_json() 

81 

82 enqueue_job("eventstream.publish", event, channels) 

83 

84 

85async def publish_members( 

86 session: AsyncSession, key: str, payload: dict[str, Any], organization_id: UUID 

87) -> None: 

88 members = await user_organization_service.list_by_org( 

89 session, org_id=organization_id 

90 ) 

91 

92 for m in members: 

93 receivers = Receivers(user_id=m.user_id) 

94 channels = receivers.get_channels() 

95 event = Event( 

96 id=generate_uuid(), 

97 key=key, 

98 payload=payload, 

99 ).model_dump_json() 

100 

101 enqueue_job("eventstream.publish", event, channels)