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

1from datetime import datetime 2 ctx1ab

2from enum import StrEnum 2 ctx1ab

3from typing import TYPE_CHECKING, Any, Self, TypedDict 2 ctx1ab

4from uuid import UUID 2 ctx1ab

5 

6from sqlalchemy import ( 2 ctx1ab

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 2 ctx1ab

20from sqlalchemy.ext.hybrid import hybrid_property 2 ctx1ab

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

22 

23from polar.config import settings 2 ctx1ab

24from polar.email.sender import DEFAULT_REPLY_TO_EMAIL_ADDRESS, EmailFromReply 2 ctx1ab

25from polar.enums import InvoiceNumbering, SubscriptionProrationBehavior 2 ctx1ab

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

27from polar.kit.extensions.sqlalchemy import StringEnum 2 ctx1ab

28 

29from .account import Account 2 ctx1ab

30 

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

32 from .organization_review import OrganizationReview 

33 from .product import Product 

34 

35 

36class OrganizationSocials(TypedDict): 2 ctx1ab

37 platform: str 2 ctx1ab

38 url: str 2 ctx1ab

39 

40 

41class OrganizationDetails(TypedDict): 2 ctx1ab

42 about: str 2 ctx1ab

43 product_description: str 2 ctx1ab

44 intended_use: str 2 ctx1ab

45 customer_acquisition: list[str] 2 ctx1ab

46 future_annual_revenue: int 2 ctx1ab

47 switching: bool 2 ctx1ab

48 switching_from: str | None 2 ctx1ab

49 previous_annual_revenue: int 2 ctx1ab

50 

51 

52class OrganizationNotificationSettings(TypedDict): 2 ctx1ab

53 new_order: bool 2 ctx1ab

54 new_subscription: bool 2 ctx1ab

55 

56 

57_default_notification_settings: OrganizationNotificationSettings = { 2 ctx1ab

58 "new_order": True, 

59 "new_subscription": True, 

60} 

61 

62 

63class OrganizationSubscriptionSettings(TypedDict): 2 ctx1ab

64 allow_multiple_subscriptions: bool 2 ctx1ab

65 allow_customer_updates: bool 2 ctx1ab

66 proration_behavior: SubscriptionProrationBehavior 2 ctx1ab

67 benefit_revocation_grace_period: int 2 ctx1ab

68 prevent_trial_abuse: bool 2 ctx1ab

69 

70 

71_default_subscription_settings: OrganizationSubscriptionSettings = { 2 ctx1ab

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): 2 ctx1ab

81 invoice_numbering: InvoiceNumbering 2 ctx1ab

82 

83 

84_default_order_settings: OrganizationOrderSettings = { 2 ctx1ab

85 "invoice_numbering": InvoiceNumbering.customer, 

86} 

87 

88 

89class OrganizationCustomerEmailSettings(TypedDict): 2 ctx1ab

90 order_confirmation: bool 2 ctx1ab

91 subscription_cancellation: bool 2 ctx1ab

92 subscription_confirmation: bool 2 ctx1ab

93 subscription_cycled: bool 2 ctx1ab

94 subscription_past_due: bool 2 ctx1ab

95 subscription_revoked: bool 2 ctx1ab

96 subscription_uncanceled: bool 2 ctx1ab

97 subscription_updated: bool 2 ctx1ab

98 

99 

100_default_customer_email_settings: OrganizationCustomerEmailSettings = { 2 ctx1ab

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): 2 ctx1ab

113 CREATED = "created" 2 ctx1ab

114 ONBOARDING_STARTED = "onboarding_started" 2 ctx1ab

115 INITIAL_REVIEW = "initial_review" 2 ctx1ab

116 ONGOING_REVIEW = "ongoing_review" 2 ctx1ab

117 DENIED = "denied" 2 ctx1ab

118 ACTIVE = "active" 2 ctx1ab

119 

120 def get_display_name(self) -> str: 2 ctx1ab

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 2 ctx1ab

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

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

133 

134 @classmethod 2 ctx1ab

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

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

137 

138 

139class Organization(RateLimitGroupMixin, RecordModel): 2 ctx1ab

140 __tablename__ = "organizations" 2 ctx1ab

141 __table_args__ = ( 2 ctx1ab

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) 2 ctx1ab

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

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

151 

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

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

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

155 JSONB, nullable=False, default=list 

156 ) 

157 details: Mapped[OrganizationDetails] = mapped_column( 2 ctx1ab

158 JSONB, nullable=False, default=dict 

159 ) 

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

161 TIMESTAMP(timezone=True) 

162 ) 

163 

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

165 customer_invoice_next_number: Mapped[int] = mapped_column( 2 ctx1ab

166 Integer, nullable=False, default=1 

167 ) 

168 

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

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

171 ) 

172 status: Mapped[OrganizationStatus] = mapped_column( 2 ctx1ab

173 StringEnum(OrganizationStatus), 

174 nullable=False, 

175 default=OrganizationStatus.CREATED, 

176 ) 

177 next_review_threshold: Mapped[int] = mapped_column( 2 ctx1ab

178 Integer, nullable=False, default=0 

179 ) 

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

181 TIMESTAMP(timezone=True), nullable=True 

182 ) 

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

184 TIMESTAMP(timezone=True), nullable=True 

185 ) 

186 

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

188 

189 @declared_attr 2 ctx1ab

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

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

192 

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

194 

195 # Time of blocking traffic/activity to given organization 

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

197 TIMESTAMP(timezone=True), 

198 nullable=True, 

199 default=None, 

200 ) 

201 

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

203 JSONB, nullable=False, default=dict 

204 ) 

205 

206 subscription_settings: Mapped[OrganizationSubscriptionSettings] = mapped_column( 2 ctx1ab

207 JSONB, nullable=False, default=_default_subscription_settings 

208 ) 

209 

210 order_settings: Mapped[OrganizationOrderSettings] = mapped_column( 2 ctx1ab

211 JSONB, nullable=False, default=_default_order_settings 

212 ) 

213 

214 notification_settings: Mapped[OrganizationNotificationSettings] = mapped_column( 2 ctx1ab

215 JSONB, nullable=False, default=_default_notification_settings 

216 ) 

217 

218 customer_email_settings: Mapped[OrganizationCustomerEmailSettings] = mapped_column( 2 ctx1ab

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( 2 ctx1ab

227 JSONB, nullable=False, default=dict 

228 ) 

229 subscriptions_billing_engine: Mapped[bool] = mapped_column( 2 ctx1ab

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) 2 ctx1ab

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

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

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

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

243 String, nullable=True, default=None 

244 ) 

245 

246 # 

247 # End: Fields synced from GitHub 

248 # 

249 

250 @hybrid_property 2 ctx1ab

251 def can_authenticate(self) -> bool: 2 ctx1ab

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

253 

254 @can_authenticate.inplace.expression 2 ctx1ab

255 @classmethod 2 ctx1ab

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

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

258 

259 @hybrid_property 2 ctx1ab

260 def storefront_enabled(self) -> bool: 2 ctx1ab

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

262 

263 @storefront_enabled.inplace.expression 2 ctx1ab

264 @classmethod 2 ctx1ab

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

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

267 

268 @hybrid_property 2 ctx1ab

269 def is_under_review(self) -> bool: 2 ctx1ab

270 return self.status in OrganizationStatus.review_statuses() 

271 

272 @is_under_review.inplace.expression 2 ctx1ab

273 @classmethod 2 ctx1ab

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

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

276 

277 @property 2 ctx1ab

278 def polar_site_url(self) -> str: 2 ctx1ab

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

280 

281 @property 2 ctx1ab

282 def account_url(self) -> str: 2 ctx1ab

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

284 

285 @property 2 ctx1ab

286 def allow_multiple_subscriptions(self) -> bool: 2 ctx1ab

287 return self.subscription_settings["allow_multiple_subscriptions"] 

288 

289 @property 2 ctx1ab

290 def allow_customer_updates(self) -> bool: 2 ctx1ab

291 return self.subscription_settings["allow_customer_updates"] 

292 

293 @property 2 ctx1ab

294 def proration_behavior(self) -> SubscriptionProrationBehavior: 2 ctx1ab

295 return SubscriptionProrationBehavior( 

296 self.subscription_settings["proration_behavior"] 

297 ) 

298 

299 @property 2 ctx1ab

300 def benefit_revocation_grace_period(self) -> int: 2 ctx1ab

301 return self.subscription_settings["benefit_revocation_grace_period"] 

302 

303 @property 2 ctx1ab

304 def prevent_trial_abuse(self) -> bool: 2 ctx1ab

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

306 

307 @property 2 ctx1ab

308 def invoice_numbering(self) -> InvoiceNumbering: 2 ctx1ab

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

310 

311 @declared_attr 2 ctx1ab

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

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

314 

315 @declared_attr 2 ctx1ab

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

317 return relationship( 2 ctx1ab

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 2 ctx1ab

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

331 return relationship( 2 ctx1ab

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: 2 ctx1ab

340 if self.blocked_at is not None: 

341 return True 

342 return False 

343 

344 def is_active(self) -> bool: 2 ctx1ab

345 return self.status == OrganizationStatus.ACTIVE 

346 

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

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 2 ctx1ab

355 def statement_descriptor_prefixed(self) -> str: 2 ctx1ab

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

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

358 

359 @property 2 ctx1ab

360 def email_from_reply(self) -> EmailFromReply: 2 ctx1ab

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 }