Coverage for polar/models/external_event.py: 91%

1from datetime import datetime 

2from enum import StrEnum 

3from typing import Any, Literal 

4 

5import stripe as stripe_lib 

6from sqlalchemy import ( 

7 TIMESTAMP, 

8 Boolean, 

9 ColumnElement, 

10 String, 

11 UniqueConstraint, 

12 type_coerce, 

13) 

14from sqlalchemy.dialects.postgresql import JSONB 

15from sqlalchemy.ext.hybrid import hybrid_property 

16from sqlalchemy.orm import Mapped, mapped_column 

17 

18from polar.kit.db.models.base import RecordModel 

19from polar.kit.extensions.sqlalchemy.types import StrEnumType 

20 

21 

22class ExternalEventSource(StrEnum): 

23 stripe = "stripe" 

24 

25 

26class ExternalEvent(RecordModel): 

27 __tablename__ = "external_events" 

28 __table_args__ = (UniqueConstraint("source", "external_id"),) 

29 

30 source: Mapped[ExternalEventSource] = mapped_column( 

31 StrEnumType(ExternalEventSource), nullable=False, index=True 

32 ) 

33 handled_at: Mapped[datetime | None] = mapped_column( 

34 TIMESTAMP(timezone=True), nullable=True, default=None, index=True 

35 ) 

36 task_name: Mapped[str] = mapped_column(String, nullable=False, index=True) 

37 external_id: Mapped[str] = mapped_column(String, nullable=False, index=True) 

38 data: Mapped[dict[str, Any]] = mapped_column("data", JSONB, nullable=False) 

39 

40 @hybrid_property 

41 def is_handled(self) -> bool: 

42 return self.handled_at is not None 

43 

44 @is_handled.inplace.expression 

45 @classmethod 

46 def _is_handled_expression(cls) -> ColumnElement[bool]: 

47 return type_coerce(cls.handled_at.is_not(None), Boolean) 

48 

49 __mapper_args__ = { 

50 "polymorphic_on": "source", 

51 } 

52 

53 

54class StripeEvent(ExternalEvent): 

55 source: Mapped[Literal[ExternalEventSource.stripe]] = mapped_column( # pyright: ignore 

56 use_existing_column=True, default=ExternalEventSource.stripe 

57 ) 

58 

59 @property 

60 def stripe_data(self) -> stripe_lib.Event: 

61 return stripe_lib.Event.construct_from(self.data, key=None) 

62 

63 __mapper_args__ = { 

64 "polymorphic_identity": ExternalEventSource.stripe, 

65 "polymorphic_load": "inline", 

66 }