Coverage for polar/license_key/schemas.py: 58%

1from datetime import datetime 

2from typing import Annotated, Literal, Self 

3 

4from dateutil.relativedelta import relativedelta 

5from pydantic import UUID4, AliasChoices, AliasPath, Field 

6from pydantic.json_schema import SkipJsonSchema 

7 

8from polar.benefit.schemas import BenefitID 

9from polar.benefit.strategies.license_keys.properties import ( 

10 BenefitLicenseKeyActivationProperties, 

11 BenefitLicenseKeyExpirationProperties, 

12) 

13from polar.customer.schemas.customer import CustomerBase 

14from polar.exceptions import NotPermitted, ResourceNotFound, Unauthorized 

15from polar.kit.metadata import ( 

16 MAXIMUM_KEYS, 

17 METADATA_DESCRIPTION, 

18 MetadataKey, 

19 MetadataValue, 

20) 

21from polar.kit.schemas import IDSchema, Schema, TimestampedSchema 

22from polar.kit.utils import generate_uuid, utc_now 

23from polar.models.license_key import LicenseKeyStatus 

24 

25############################################################################### 

26# RESPONSES 

27############################################################################### 

28 

29NotFoundResponse = { 

30 "description": "License key not found.", 

31 "model": ResourceNotFound.schema(), 

32} 

33 

34UnauthorizedResponse = { 

35 "description": "Not authorized to manage license key.", 

36 "model": Unauthorized.schema(), 

37} 

38 

39 

40ActivationNotPermitted = { 

41 "description": "License key activation not supported or limit reached. Use /validate endpoint for licenses without activations.", 

42 "model": NotPermitted.schema(), 

43} 

44 

45 

46############################################################################### 

47# RESPONSES 

48############################################################################### 

49 

50ConditionsField = Annotated[ 

51 dict[MetadataKey, MetadataValue], 

52 Field( 

53 default_factory=dict, 

54 max_length=MAXIMUM_KEYS, 

55 description=METADATA_DESCRIPTION.format( 

56 heading=( 

57 "Key-value object allowing you to set conditions " 

58 "that must match when validating the license key." 

59 ) 

60 ), 

61 ), 

62] 

63 

64MetaField = Annotated[ 

65 dict[MetadataKey, MetadataValue], 

66 Field( 

67 default_factory=dict, 

68 max_length=MAXIMUM_KEYS, 

69 description=METADATA_DESCRIPTION.format( 

70 heading=( 

71 "Key-value object allowing you to store additional information " 

72 "about the activation" 

73 ) 

74 ), 

75 ), 

76] 

77 

78 

79class LicenseKeyValidate(Schema): 

80 key: str 

81 organization_id: UUID4 

82 activation_id: UUID4 | None = None 

83 benefit_id: BenefitID | None = None 

84 customer_id: UUID4 | None = None 

85 increment_usage: int | None = None 

86 conditions: ConditionsField 

87 

88 

89class LicenseKeyActivate(Schema): 

90 key: str 

91 organization_id: UUID4 

92 label: str 

93 conditions: ConditionsField 

94 meta: MetaField 

95 

96 

97class LicenseKeyDeactivate(Schema): 

98 key: str 

99 organization_id: UUID4 

100 activation_id: UUID4 

101 

102 

103class LicenseKeyCustomer(CustomerBase): ... 103 ↛ 106line 103 didn't jump to line 106 because

104 

105 

106class LicenseKeyUser(Schema): 

107 id: UUID4 = Field( 

108 validation_alias=AliasChoices( 

109 # Validate from ORM model 

110 "legacy_user_id", 

111 # Validate from stored webhook payload 

112 "id", 

113 ) 

114 ) 

115 email: str 

116 public_name: str = Field( 

117 validation_alias=AliasChoices( 

118 # Validate from ORM model 

119 "legacy_user_public_name", 

120 # Validate from stored webhook payload 

121 "public_name", 

122 ) 

123 ) 

124 avatar_url: str | None = Field(None) 

125 

126 

127class LicenseKeyRead(TimestampedSchema, IDSchema): 

128 organization_id: UUID4 

129 user_id: SkipJsonSchema[UUID4] = Field( 

130 validation_alias=AliasChoices( 

131 # Validate from stored webhook payload 

132 "user_id", 

133 # Validate from ORM model 

134 AliasPath("customer", "legacy_user_id"), 

135 ), 

136 deprecated="Use `customer_id`.", 

137 ) 

138 customer_id: UUID4 

139 user: SkipJsonSchema[LicenseKeyUser] = Field( 

140 validation_alias=AliasChoices( 

141 # Validate from stored webhook payload 

142 "user", 

143 # Validate from ORM model 

144 "customer", 

145 ), 

146 deprecated="Use `customer`.", 

147 ) 

148 customer: LicenseKeyCustomer 

149 benefit_id: BenefitID 

150 key: str 

151 display_key: str 

152 status: LicenseKeyStatus 

153 limit_activations: int | None 

154 usage: int 

155 limit_usage: int | None 

156 validations: int 

157 last_validated_at: datetime | None 

158 expires_at: datetime | None 

159 

160 

161class LicenseKeyActivationBase(Schema): 

162 id: UUID4 

163 license_key_id: UUID4 

164 label: str 

165 meta: dict[str, str | int | float | bool] 

166 created_at: datetime 

167 modified_at: datetime | None 

168 

169 

170class LicenseKeyWithActivations(LicenseKeyRead): 

171 activations: list[LicenseKeyActivationBase] 

172 

173 

174class ValidatedLicenseKey(LicenseKeyRead): 

175 activation: LicenseKeyActivationBase | None = None 

176 

177 

178class LicenseKeyActivationRead(LicenseKeyActivationBase): 

179 license_key: LicenseKeyRead 

180 

181 

182class LicenseKeyUpdate(Schema): 

183 status: LicenseKeyStatus | None = None 

184 usage: int = 0 

185 limit_activations: int | None = Field(gt=0, le=1000, default=None) 

186 limit_usage: int | None = Field(gt=0, default=None) 

187 expires_at: datetime | None = None 

188 

189 

190class LicenseKeyCreate(LicenseKeyUpdate): 

191 organization_id: UUID4 

192 customer_id: UUID4 

193 benefit_id: BenefitID 

194 key: str 

195 

196 @classmethod 

197 def generate_key(cls, prefix: str | None = None) -> str: 

198 key = str(generate_uuid()).upper() 

199 if prefix is None: 

200 return key 

201 

202 prefix = prefix.strip().upper() 

203 return f"{prefix}-{key}" 

204 

205 @classmethod 

206 def generate_expiration_dt( 

207 cls, ttl: int, timeframe: Literal["year", "month", "day"] 

208 ) -> datetime: 

209 now = utc_now() 

210 match timeframe: 

211 case "year": 

212 return now + relativedelta(years=ttl) 

213 case "month": 

214 return now + relativedelta(months=ttl) 

215 case _: 

216 return now + relativedelta(days=ttl) 

217 

218 @classmethod 

219 def build( 

220 cls, 

221 organization_id: UUID4, 

222 customer_id: UUID4, 

223 benefit_id: UUID4, 

224 prefix: str | None = None, 

225 status: LicenseKeyStatus = LicenseKeyStatus.granted, 

226 limit_usage: int | None = None, 

227 activations: BenefitLicenseKeyActivationProperties | None = None, 

228 expires: BenefitLicenseKeyExpirationProperties | None = None, 

229 ) -> Self: 

230 expires_at = None 

231 if expires: 

232 ttl = expires.get("ttl", None) 

233 timeframe = expires.get("timeframe", None) 

234 if ttl and timeframe: 

235 expires_at = cls.generate_expiration_dt(ttl, timeframe) 

236 

237 limit_activations = None 

238 if activations: 

239 limit_activations = activations.get("limit", None) 

240 

241 key = cls.generate_key(prefix=prefix) 

242 return cls( 

243 organization_id=organization_id, 

244 customer_id=customer_id, 

245 benefit_id=benefit_id, 

246 key=key, 

247 status=status, 

248 limit_activations=limit_activations, 

249 limit_usage=limit_usage, 

250 expires_at=expires_at, 

251 )