Coverage for polar/models/organization.py: 85%

1from datetime import datetime 

2from enum import StrEnum 

3from typing import TYPE_CHECKING, Any, Self, TypedDict 

4from uuid import UUID 

5 

6from sqlalchemy import ( 

7 TIMESTAMP, 

8 Boolean, 

9 CheckConstraint, 

10 ColumnElement, 

11 ForeignKey, 

12 Integer, 

13 String, 

14 Text, 

15 UniqueConstraint, 

16 Uuid, 

17 and_, 

18) 

19from sqlalchemy.dialects.postgresql import CITEXT, JSONB 

20from sqlalchemy.ext.hybrid import hybrid_property 

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

22 

23from polar.config import settings 

24from polar.email.sender import DEFAULT_REPLY_TO_EMAIL_ADDRESS, EmailFromReply 

25from polar.enums import InvoiceNumbering, SubscriptionProrationBehavior 

26from polar.kit.db.models import RateLimitGroupMixin, RecordModel 

27from polar.kit.extensions.sqlalchemy import StringEnum 

28 

29from .account import Account 

30 

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

32 from .organization_review import OrganizationReview 

33 from .product import Product 

34 

35 

36class OrganizationSocials(TypedDict): 

37 platform: str 

38 url: str 

39 

40 

41class OrganizationDetails(TypedDict): 

42 about: str 

43 product_description: str 

44 intended_use: str 

45 customer_acquisition: list[str] 

46 future_annual_revenue: int 

47 switching: bool 

48 switching_from: str | None 

49 previous_annual_revenue: int 

50 

51 

52class OrganizationNotificationSettings(TypedDict): 

53 new_order: bool 

54 new_subscription: bool 

55 

56 

57_default_notification_settings: OrganizationNotificationSettings = { 

58 "new_order": True, 

59 "new_subscription": True, 

60} 

61 

62 

63class OrganizationSubscriptionSettings(TypedDict): 

64 allow_multiple_subscriptions: bool 

65 allow_customer_updates: bool 

66 proration_behavior: SubscriptionProrationBehavior 

67 benefit_revocation_grace_period: int 

68 prevent_trial_abuse: bool 

69 

70 

71_default_subscription_settings: OrganizationSubscriptionSettings = { 

72 "allow_multiple_subscriptions": False, 

73 "allow_customer_updates": True, 

74 "proration_behavior": SubscriptionProrationBehavior.prorate, 

75 "benefit_revocation_grace_period": 0, 

76 "prevent_trial_abuse": False, 

77} 

78 

79 

80class OrganizationOrderSettings(TypedDict): 

81 invoice_numbering: InvoiceNumbering 

82 

83 

84_default_order_settings: OrganizationOrderSettings = { 

85 "invoice_numbering": InvoiceNumbering.customer, 

86} 

87 

88 

89class OrganizationCustomerEmailSettings(TypedDict): 

90 order_confirmation: bool 

91 subscription_cancellation: bool 

92 subscription_confirmation: bool 

93 subscription_cycled: bool 

94 subscription_past_due: bool 

95 subscription_revoked: bool 

96 subscription_uncanceled: bool 

97 subscription_updated: bool 

98 

99 

100_default_customer_email_settings: OrganizationCustomerEmailSettings = { 

101 "order_confirmation": True, 

102 "subscription_cancellation": True, 

103 "subscription_confirmation": True, 

104 "subscription_cycled": True, 

105 "subscription_past_due": True, 

106 "subscription_revoked": True, 

107 "subscription_uncanceled": True, 

108 "subscription_updated": True, 

109} 

110 

111 

112class OrganizationStatus(StrEnum): 

113 CREATED = "created" 

114 ONBOARDING_STARTED = "onboarding_started" 

115 INITIAL_REVIEW = "initial_review" 

116 ONGOING_REVIEW = "ongoing_review" 

117 DENIED = "denied" 

118 ACTIVE = "active" 

119 

120 def get_display_name(self) -> str: 

121 return { 

122 OrganizationStatus.CREATED: "Created", 

123 OrganizationStatus.ONBOARDING_STARTED: "Onboarding Started", 

124 OrganizationStatus.INITIAL_REVIEW: "Initial Review", 

125 OrganizationStatus.ONGOING_REVIEW: "Ongoing Review", 

126 OrganizationStatus.DENIED: "Denied", 

127 OrganizationStatus.ACTIVE: "Active", 

128 }[self] 

129 

130 @classmethod 

131 def review_statuses(cls) -> set[Self]: 

132 return {cls.INITIAL_REVIEW, cls.ONGOING_REVIEW} # type: ignore 

133 

134 @classmethod 

135 def payment_ready_statuses(cls) -> set[Self]: 

136 return {cls.ACTIVE, *cls.review_statuses()} # type: ignore 

137 

138 

139class Organization(RateLimitGroupMixin, RecordModel): 

140 __tablename__ = "organizations" 

141 __table_args__ = ( 

142 UniqueConstraint("slug"), 

143 CheckConstraint( 

144 "next_review_threshold >= 0", name="next_review_threshold_positive" 

145 ), 

146 ) 

147 

148 name: Mapped[str] = mapped_column(String, nullable=False, index=True) 

149 slug: Mapped[str] = mapped_column(CITEXT, nullable=False, unique=True) 

150 avatar_url: Mapped[str | None] = mapped_column(String, nullable=True) 

151 

152 email: Mapped[str | None] = mapped_column(String, nullable=True, default=None) 

153 website: Mapped[str | None] = mapped_column(String, nullable=True, default=None) 

154 socials: Mapped[list[OrganizationSocials]] = mapped_column( 

155 JSONB, nullable=False, default=list 

156 ) 

157 details: Mapped[OrganizationDetails] = mapped_column( 

158 JSONB, nullable=False, default=dict 

159 ) 

160 details_submitted_at: Mapped[datetime | None] = mapped_column( 

161 TIMESTAMP(timezone=True) 

162 ) 

163 

164 customer_invoice_prefix: Mapped[str] = mapped_column(String, nullable=False) 

165 customer_invoice_next_number: Mapped[int] = mapped_column( 

166 Integer, nullable=False, default=1 

167 ) 

168 

169 account_id: Mapped[UUID | None] = mapped_column( 

170 Uuid, ForeignKey("accounts.id", ondelete="set null"), nullable=True 

171 ) 

172 status: Mapped[OrganizationStatus] = mapped_column( 

173 StringEnum(OrganizationStatus), 

174 nullable=False, 

175 default=OrganizationStatus.CREATED, 

176 ) 

177 next_review_threshold: Mapped[int] = mapped_column( 

178 Integer, nullable=False, default=0 

179 ) 

180 status_updated_at: Mapped[datetime | None] = mapped_column( 

181 TIMESTAMP(timezone=True), nullable=True 

182 ) 

183 initially_reviewed_at: Mapped[datetime | None] = mapped_column( 

184 TIMESTAMP(timezone=True), nullable=True 

185 ) 

186 

187 internal_notes: Mapped[str | None] = mapped_column(Text, nullable=True) 

188 

189 @declared_attr 

190 def account(cls) -> Mapped[Account | None]: 

191 return relationship(Account, lazy="raise", back_populates="organizations") 

192 

193 onboarded_at: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True)) 

194 

195 # Time of blocking traffic/activity to given organization 

196 blocked_at: Mapped[datetime | None] = mapped_column( 

197 TIMESTAMP(timezone=True), 

198 nullable=True, 

199 default=None, 

200 ) 

201 

202 profile_settings: Mapped[dict[str, Any]] = mapped_column( 

203 JSONB, nullable=False, default=dict 

204 ) 

205 

206 subscription_settings: Mapped[OrganizationSubscriptionSettings] = mapped_column( 

207 JSONB, nullable=False, default=_default_subscription_settings 

208 ) 

209 

210 order_settings: Mapped[OrganizationOrderSettings] = mapped_column( 

211 JSONB, nullable=False, default=_default_order_settings 

212 ) 

213 

214 notification_settings: Mapped[OrganizationNotificationSettings] = mapped_column( 

215 JSONB, nullable=False, default=_default_notification_settings 

216 ) 

217 

218 customer_email_settings: Mapped[OrganizationCustomerEmailSettings] = mapped_column( 

219 JSONB, nullable=False, default=_default_customer_email_settings 

220 ) 

221 

222 # 

223 # Feature Flags 

224 # 

225 

226 feature_settings: Mapped[dict[str, Any]] = mapped_column( 

227 JSONB, nullable=False, default=dict 

228 ) 

229 subscriptions_billing_engine: Mapped[bool] = mapped_column( 

230 Boolean, nullable=False, default=settings.ORGANIZATIONS_BILLING_ENGINE_DEFAULT 

231 ) 

232 

233 # 

234 # Fields synced from GitHub 

235 # 

236 

237 # Org description or user bio 

238 bio: Mapped[str | None] = mapped_column(String, nullable=True, default=None) 

239 company: Mapped[str | None] = mapped_column(String, nullable=True, default=None) 

240 blog: Mapped[str | None] = mapped_column(String, nullable=True, default=None) 

241 location: Mapped[str | None] = mapped_column(String, nullable=True, default=None) 

242 twitter_username: Mapped[str | None] = mapped_column( 

243 String, nullable=True, default=None 

244 ) 

245 

246 # 

247 # End: Fields synced from GitHub 

248 # 

249 

250 @hybrid_property 

251 def can_authenticate(self) -> bool: 

252 return self.deleted_at is None and self.blocked_at is None 

253 

254 @can_authenticate.inplace.expression 

255 @classmethod 

256 def _can_authenticate_expression(cls) -> ColumnElement[bool]: 

257 return and_(cls.deleted_at.is_(None), cls.blocked_at.is_(None)) 

258 

259 @hybrid_property 

260 def storefront_enabled(self) -> bool: 

261 return self.profile_settings.get("enabled", False) 

262 

263 @storefront_enabled.inplace.expression 

264 @classmethod 

265 def _storefront_enabled_expression(cls) -> ColumnElement[bool]: 

266 return Organization.profile_settings["enabled"].as_boolean() 

267 

268 @hybrid_property 

269 def is_under_review(self) -> bool: 

270 return self.status in OrganizationStatus.review_statuses() 

271 

272 @is_under_review.inplace.expression 

273 @classmethod 

274 def _is_under_review_expression(cls) -> ColumnElement[bool]: 

275 return cls.status.in_(OrganizationStatus.review_statuses()) 

276 

277 @property 

278 def polar_site_url(self) -> str: 

279 return f"{settings.FRONTEND_BASE_URL}/{self.slug}" 

280 

281 @property 

282 def account_url(self) -> str: 

283 return f"{settings.FRONTEND_BASE_URL}/dashboard/{self.slug}/finance/account" 

284 

285 @property 

286 def allow_multiple_subscriptions(self) -> bool: 

287 return self.subscription_settings["allow_multiple_subscriptions"] 

288 

289 @property 

290 def allow_customer_updates(self) -> bool: 

291 return self.subscription_settings["allow_customer_updates"] 

292 

293 @property 

294 def proration_behavior(self) -> SubscriptionProrationBehavior: 

295 return SubscriptionProrationBehavior( 

296 self.subscription_settings["proration_behavior"] 

297 ) 

298 

299 @property 

300 def benefit_revocation_grace_period(self) -> int: 

301 return self.subscription_settings["benefit_revocation_grace_period"] 

302 

303 @property 

304 def prevent_trial_abuse(self) -> bool: 

305 return self.subscription_settings.get("prevent_trial_abuse", False) 

306 

307 @property 

308 def invoice_numbering(self) -> InvoiceNumbering: 

309 return InvoiceNumbering(self.order_settings["invoice_numbering"]) 

310 

311 @declared_attr 

312 def all_products(cls) -> Mapped[list["Product"]]: 

313 return relationship("Product", lazy="raise", back_populates="organization") 

314 

315 @declared_attr 

316 def products(cls) -> Mapped[list["Product"]]: 

317 return relationship( 

318 "Product", 

319 lazy="raise", 

320 primaryjoin=( 

321 "and_(" 

322 "Product.organization_id == Organization.id, " 

323 "Product.is_archived.is_(False)" 

324 ")" 

325 ), 

326 viewonly=True, 

327 ) 

328 

329 @declared_attr 

330 def review(cls) -> Mapped["OrganizationReview | None"]: 

331 return relationship( 

332 "OrganizationReview", 

333 lazy="raise", 

334 back_populates="organization", 

335 cascade="delete, delete-orphan", 

336 uselist=False, # This makes it a one-to-one relationship 

337 ) 

338 

339 def is_blocked(self) -> bool: 

340 if self.blocked_at is not None: 340 ↛ 341line 340 didn't jump to line 341 because the condition on line 340 was never true

341 return True 

342 return False 

343 

344 def is_active(self) -> bool: 

345 return self.status == OrganizationStatus.ACTIVE 

346 

347 def statement_descriptor(self, suffix: str = "") -> str: 

348 max_length = settings.stripe_descriptor_suffix_max_length 

349 if suffix: 

350 space_for_slug = max_length - len(suffix) 

351 return self.slug[:space_for_slug] + suffix 

352 return self.slug[:max_length] 

353 

354 @property 

355 def statement_descriptor_prefixed(self) -> str: 

356 # Cannot use *. Setting separator to # instead. 

357 return f"{settings.STRIPE_STATEMENT_DESCRIPTOR}# {self.statement_descriptor()}" 

358 

359 @property 

360 def email_from_reply(self) -> EmailFromReply: 

361 return { 

362 "from_name": f"{self.name} (via {settings.EMAIL_FROM_NAME})", 

363 "from_email_addr": f"{self.slug}@{settings.EMAIL_FROM_DOMAIN}", 

364 "reply_to_name": self.name, 

365 "reply_to_email_addr": self.email or DEFAULT_REPLY_TO_EMAIL_ADDRESS, 

366 }