Coverage for polar/models/billing_entry.py: 95%

1from datetime import datetime 

2from enum import StrEnum 

3from typing import TYPE_CHECKING, Self 

4from uuid import UUID 

5 

6from sqlalchemy import TIMESTAMP, ForeignKey, Index, String, Uuid 

7from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship 

8from sqlalchemy.types import Integer 

9 

10from polar.kit.db.models import RecordModel 

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

12 

13if TYPE_CHECKING: 13 ↛ 14line 13 didn't jump to line 14 because the condition on line 13 was never true

14 from polar.models import ( 

15 Customer, 

16 Discount, 

17 Event, 

18 OrderItem, 

19 ProductPrice, 

20 Subscription, 

21 SubscriptionProductPrice, 

22 ) 

23 

24 

25class BillingEntryDirection(StrEnum): 

26 debit = "debit" 

27 credit = "credit" 

28 

29 

30class BillingEntryType(StrEnum): 

31 cycle = "cycle" 

32 proration = "proration" 

33 metered = "metered" 

34 subscription_seats_increase = "subscription_seats_increase" 

35 subscription_seats_decrease = "subscription_seats_decrease" 

36 

37 

38class BillingEntry(RecordModel): 

39 __tablename__ = "billing_entry" 

40 __table_args__ = ( 

41 Index( 

42 "ix_billing_entries_s_oi_pp", 

43 "subscription_id", 

44 "order_item_id", 

45 "product_price_id", 

46 ), 

47 Index( 

48 "ix_billing_entries_s_d_oi_pp", 

49 "subscription_id", 

50 "deleted_at", 

51 "order_item_id", 

52 "product_price_id", 

53 ), 

54 ) 

55 

56 start_timestamp: Mapped[datetime] = mapped_column( 

57 TIMESTAMP(timezone=True), nullable=False, index=True 

58 ) 

59 end_timestamp: Mapped[datetime] = mapped_column( 

60 TIMESTAMP(timezone=True), nullable=False, index=True 

61 ) 

62 type: Mapped[BillingEntryType] = mapped_column( 

63 StrEnumType(BillingEntryType), nullable=False, index=True 

64 ) 

65 direction: Mapped[BillingEntryDirection] = mapped_column( 

66 StrEnumType(BillingEntryDirection), nullable=False 

67 ) 

68 amount: Mapped[int | None] = mapped_column(Integer, nullable=True, default=None) 

69 discount_amount: Mapped[int | None] = mapped_column( 

70 Integer, nullable=True, default=None 

71 ) 

72 currency: Mapped[str | None] = mapped_column(String(3), nullable=True, default=None) 

73 customer_id: Mapped[UUID] = mapped_column( 

74 Uuid, ForeignKey("customers.id", ondelete="cascade"), nullable=False, index=True 

75 ) 

76 product_price_id: Mapped[UUID] = mapped_column( 

77 Uuid, 

78 ForeignKey("product_prices.id", ondelete="restrict"), 

79 nullable=False, 

80 index=True, 

81 ) 

82 subscription_id: Mapped[UUID | None] = mapped_column( 

83 Uuid, 

84 ForeignKey("subscriptions.id", ondelete="cascade"), 

85 nullable=True, 

86 index=True, 

87 ) 

88 discount_id: Mapped[UUID | None] = mapped_column( 

89 Uuid, ForeignKey("discounts.id", ondelete="restrict"), nullable=True 

90 ) 

91 event_id: Mapped[UUID] = mapped_column( 

92 Uuid, ForeignKey("events.id", ondelete="cascade"), nullable=False 

93 ) 

94 order_item_id: Mapped[UUID | None] = mapped_column( 

95 Uuid, 

96 ForeignKey("order_items.id", ondelete="cascade"), 

97 nullable=True, 

98 index=True, 

99 ) 

100 

101 @declared_attr 

102 def customer(cls) -> Mapped["Customer"]: 

103 return relationship("Customer", lazy="raise_on_sql") 

104 

105 @declared_attr 

106 def product_price(cls) -> Mapped["ProductPrice"]: 

107 return relationship("ProductPrice", lazy="raise_on_sql") 

108 

109 @declared_attr 

110 def subscription(cls) -> Mapped["Subscription | None"]: 

111 return relationship("Subscription", lazy="raise_on_sql") 

112 

113 @declared_attr 

114 def discount(cls) -> Mapped["Discount | None"]: 

115 return relationship("Discount", lazy="raise_on_sql") 

116 

117 @declared_attr 

118 def event(cls) -> Mapped["Event"]: 

119 return relationship("Event", lazy="raise_on_sql") 

120 

121 @declared_attr 

122 def order_item(cls) -> Mapped["OrderItem | None"]: 

123 return relationship("OrderItem", lazy="raise_on_sql") 

124 

125 @classmethod 

126 def from_metered_event( 

127 cls, 

128 customer: "Customer", 

129 subscription_product_price: "SubscriptionProductPrice", 

130 event: "Event", 

131 ) -> Self: 

132 return cls( 

133 start_timestamp=event.timestamp, 

134 end_timestamp=event.timestamp, 

135 type=BillingEntryType.metered, 

136 direction=BillingEntryDirection.debit, 

137 customer=customer, 

138 product_price=subscription_product_price.product_price, 

139 subscription=subscription_product_price.subscription, 

140 event=event, 

141 )