Coverage for polar/models/customer_seat.py: 77%

1from datetime import datetime 

2from enum import StrEnum 

3from typing import TYPE_CHECKING, Any 

4from uuid import UUID 

5 

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

7from sqlalchemy.dialects.postgresql import JSONB 

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

9 

10from polar.kit.db.models import RecordModel 

11from polar.product.guard import is_metered_price 

12 

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

14 from .customer import Customer 

15 from .order import Order 

16 from .subscription import Subscription 

17 

18 

19class SeatStatus(StrEnum): 

20 pending = "pending" 

21 claimed = "claimed" 

22 revoked = "revoked" 

23 

24 

25class CustomerSeat(RecordModel): 

26 """ 

27 Represents a seat that can be assigned to a customer for access to benefits. It's 

28 mainly used for the seat-based billing. 

29 

30 Seats can be associated with either: 

31 - A recurring subscription (subscription_id) for ongoing seat-based billing 

32 - A one-time order (order_id) for perpetual seat-based purchases 

33 

34 Lifecycle: 

35 1. pending: Seat created but not yet assigned/claimed by an end customer 

36 2. claimed: Customer has accepted the invitation and benefits are granted 

37 3. revoked: Seat has been revoked, benefits removed, can be reassigned 

38 

39 Key constraints: 

40 - Exactly one of subscription_id OR order_id must be set (enforced by seat_source_check) 

41 - For subscriptions: seats are tied to billing cycle, revoked when subscription ends 

42 - For orders: seats are perpetual, never expire once claimed 

43 """ 

44 

45 __tablename__ = "customer_seats" 

46 __table_args__ = ( 

47 CheckConstraint( 

48 "(subscription_id IS NOT NULL AND order_id IS NULL) OR " 

49 "(subscription_id IS NULL AND order_id IS NOT NULL)", 

50 name="seat_source_check", 

51 ), 

52 ) 

53 

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

55 Uuid, 

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

57 nullable=True, 

58 index=True, 

59 ) 

60 

61 order_id: Mapped[UUID | None] = mapped_column( 

62 Uuid, 

63 ForeignKey("orders.id", ondelete="cascade"), 

64 nullable=True, 

65 index=True, 

66 ) 

67 

68 status: Mapped[SeatStatus] = mapped_column( 

69 String, nullable=False, default=SeatStatus.pending 

70 ) 

71 

72 customer_id: Mapped[UUID | None] = mapped_column( 

73 Uuid, 

74 ForeignKey("customers.id", ondelete="set null"), 

75 nullable=True, 

76 index=True, 

77 ) 

78 

79 invitation_token: Mapped[str | None] = mapped_column( 

80 String, nullable=True, default=None, index=True 

81 ) 

82 

83 invitation_token_expires_at: Mapped[datetime | None] = mapped_column( 

84 TIMESTAMP(timezone=True), nullable=True, default=None 

85 ) 

86 

87 claimed_at: Mapped[datetime | None] = mapped_column( 

88 TIMESTAMP(timezone=True), nullable=True, default=None 

89 ) 

90 

91 revoked_at: Mapped[datetime | None] = mapped_column( 

92 TIMESTAMP(timezone=True), nullable=True, default=None 

93 ) 

94 

95 seat_metadata: Mapped[dict[str, Any] | None] = mapped_column( 

96 "metadata", JSONB, nullable=True, default=None 

97 ) 

98 

99 @declared_attr 

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

101 return relationship("Subscription", lazy="raise") 

102 

103 @declared_attr 

104 def order(cls) -> Mapped["Order | None"]: 

105 return relationship("Order", lazy="raise") 

106 

107 @declared_attr 

108 def customer(cls) -> Mapped["Customer | None"]: 

109 return relationship("Customer", lazy="raise") 

110 

111 def is_pending(self) -> bool: 

112 return self.status == SeatStatus.pending 

113 

114 def is_claimed(self) -> bool: 

115 return self.status == SeatStatus.claimed 

116 

117 def is_revoked(self) -> bool: 

118 return self.status == SeatStatus.revoked 

119 

120 def has_metered_pricing(self) -> bool: 

121 if not self.subscription: 

122 return False 

123 

124 return any( 

125 is_metered_price(spp.product_price) 

126 for spp in self.subscription.subscription_product_prices 

127 )