Coverage for polar/customer_seat/schemas.py: 53%

79 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-12-05 16:17 +0000

1from datetime import datetime 1a

2from typing import Any 1a

3from uuid import UUID 1a

4 

5from pydantic import EmailStr, Field, field_validator, model_validator 1a

6from sqlalchemy import inspect 1a

7 

8from polar.kit.schemas import Schema, TimestampedSchema 1a

9from polar.models.customer_seat import SeatStatus 1a

10 

11 

12class SeatAssign(Schema): 1a

13 subscription_id: UUID | None = Field( 1a

14 None, 

15 description="Subscription ID. Required if checkout_id and order_id are not provided.", 

16 ) 

17 checkout_id: UUID | None = Field( 1a

18 None, 

19 description="Checkout ID. Used to look up subscription or order from the checkout page.", 

20 ) 

21 order_id: UUID | None = Field( 1a

22 None, 

23 description="Order ID for one-time purchases. Required if subscription_id and checkout_id are not provided.", 

24 ) 

25 email: EmailStr | None = Field( 1a

26 None, description="Email of the customer to assign the seat to" 

27 ) 

28 external_customer_id: str | None = Field( 1a

29 None, description="External customer ID for the seat assignment" 

30 ) 

31 customer_id: UUID | None = Field( 1a

32 None, description="Customer ID for the seat assignment" 

33 ) 

34 metadata: dict[str, Any] | None = Field( 1a

35 None, description="Additional metadata for the seat (max 10 keys, 1KB total)" 

36 ) 

37 immediate_claim: bool = Field( 1a

38 default=False, 

39 description="If true, the seat will be immediately claimed without sending an invitation email. API-only feature.", 

40 ) 

41 

42 @field_validator("metadata") 1a

43 @classmethod 1a

44 def validate_metadata(cls, v: dict[str, Any] | None) -> dict[str, Any] | None: 1a

45 if v is None: 

46 return v 

47 if len(v) > 10: 

48 raise ValueError("Metadata cannot have more than 10 keys") 

49 import json 

50 

51 if len(json.dumps(v)) > 1024: # 1KB limit 

52 raise ValueError("Metadata size cannot exceed 1KB") 

53 return v 

54 

55 @model_validator(mode="after") 1a

56 def validate_identifiers(self) -> "SeatAssign": 1a

57 seat_source_identifiers = [ 

58 self.subscription_id, 

59 self.checkout_id, 

60 self.order_id, 

61 ] 

62 seat_source_count = sum( 

63 1 for identifier in seat_source_identifiers if identifier is not None 

64 ) 

65 

66 if seat_source_count != 1: 

67 raise ValueError( 

68 "Exactly one of subscription_id, checkout_id, or order_id must be provided" 

69 ) 

70 

71 customer_identifiers = [self.email, self.external_customer_id, self.customer_id] 

72 customer_count = sum( 

73 1 for identifier in customer_identifiers if identifier is not None 

74 ) 

75 

76 if customer_count != 1: 

77 raise ValueError( 

78 "Exactly one of email, external_customer_id, or customer_id must be provided" 

79 ) 

80 

81 return self 

82 

83 

84class SeatClaim(Schema): 1a

85 invitation_token: str = Field(..., description="Invitation token to claim the seat") 1a

86 

87 

88class CustomerSeat(TimestampedSchema): 1a

89 id: UUID = Field(..., description="The seat ID") 1a

90 subscription_id: UUID | None = Field( 1a

91 None, description="The subscription ID (for recurring seats)" 

92 ) 

93 order_id: UUID | None = Field( 1a

94 None, description="The order ID (for one-time purchase seats)" 

95 ) 

96 status: SeatStatus = Field(..., description="Status of the seat") 1a

97 customer_id: UUID | None = Field(None, description="The assigned customer ID") 1a

98 customer_email: str | None = Field(None, description="The assigned customer email") 1a

99 invitation_token_expires_at: datetime | None = Field( 1a

100 None, description="When the invitation token expires" 

101 ) 

102 claimed_at: datetime | None = Field(None, description="When the seat was claimed") 1a

103 revoked_at: datetime | None = Field(None, description="When the seat was revoked") 1a

104 seat_metadata: dict[str, Any] | None = Field( 1a

105 None, description="Additional metadata for the seat" 

106 ) 

107 

108 @model_validator(mode="before") 1a

109 @classmethod 1a

110 def extract_customer_email(cls, data: Any) -> Any: 1a

111 if isinstance(data, dict): 

112 # For dict data 

113 if "customer" in data and data["customer"]: 

114 data["customer_email"] = data.get("customer", {}).get("email") 

115 return data 

116 elif hasattr(data, "__dict__"): 

117 # For SQLAlchemy models - check if customer is loaded 

118 state = inspect(data) 

119 if "customer" not in state.unloaded: 

120 # Customer is loaded, we can extract the email 

121 # But we need to let Pydantic handle the model conversion 

122 # We'll just add the customer_email field if customer is available 

123 if hasattr(data, "customer") and data.customer: 

124 # Add customer_email as a temporary attribute 

125 object.__setattr__(data, "customer_email", data.customer.email) 

126 return data 

127 

128 

129class SeatsList(Schema): 1a

130 seats: list[CustomerSeat] = Field(..., description="List of seats") 1a

131 available_seats: int = Field(..., description="Number of available seats") 1a

132 total_seats: int = Field( 1a

133 ..., description="Total number of seats for the subscription" 

134 ) 

135 

136 

137class SeatClaimInfo(Schema): 1a

138 """ 

139 Read-only information about a seat claim invitation. 

140 Safe for email scanners - no side effects when fetched. 

141 """ 

142 

143 product_name: str = Field(..., description="Name of the product") 1a

144 product_id: UUID = Field(..., description="ID of the product") 1a

145 organization_name: str = Field(..., description="Name of the organization") 1a

146 organization_slug: str = Field(..., description="Slug of the organization") 1a

147 customer_email: str = Field( 1a

148 ..., description="Email of the customer assigned to this seat" 

149 ) 

150 can_claim: bool = Field(..., description="Whether the seat can be claimed") 1a

151 

152 

153class CustomerSeatClaimResponse(Schema): 1a

154 """Response after successfully claiming a seat.""" 

155 

156 seat: CustomerSeat = Field(..., description="The claimed seat") 1a

157 customer_session_token: str = Field( 1a

158 ..., description="Session token for immediate customer portal access" 

159 ) 

160 

161 

162CustomerSeatID = UUID 1a