Coverage for polar/models/benefit_grant.py: 64%

1from datetime import UTC, datetime 

2from typing import TYPE_CHECKING, Any, TypedDict 

3from uuid import UUID 

4 

5from sqlalchemy import ( 

6 TIMESTAMP, 

7 Boolean, 

8 ColumnElement, 

9 ForeignKey, 

10 UniqueConstraint, 

11 Uuid, 

12 and_, 

13 type_coerce, 

14) 

15from sqlalchemy.dialects.postgresql import JSONB 

16from sqlalchemy.ext.hybrid import hybrid_property 

17from sqlalchemy.orm import ( 

18 CompositeProperty, 

19 Mapped, 

20 composite, 

21 declared_attr, 

22 mapped_column, 

23 relationship, 

24) 

25 

26from polar.kit.db.models import RecordModel 

27 

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

29 from polar.benefit.strategies import BenefitGrantProperties 

30 from polar.models import Benefit, Customer, Member, Order, Subscription 

31 

32 

33class BenefitGrantError(TypedDict): 

34 message: str 

35 type: str 

36 timestamp: str 

37 

38 

39class BenefitGrantScope(TypedDict, total=False): 

40 subscription: "Subscription" 

41 order: "Order" 

42 

43 

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

45 

46 class BenefitGrantScopeArgs(TypedDict, total=False): 

47 subscription_id: UUID 

48 order_id: UUID 

49 

50else: 

51 

52 class BenefitGrantScopeArgs(dict): 

53 def __init__( 

54 self, subscription_id: UUID | None = None, order_id: UUID | None = None 

55 ) -> None: 

56 d = {} 

57 if subscription_id is not None: 

58 d["subscription_id"] = subscription_id 

59 if order_id is not None: 

60 d["order_id"] = order_id 

61 super().__init__(d) 

62 

63 def __composite_values__(self) -> tuple[UUID | None, UUID | None]: 

64 return self.get("subscription_id"), self.get("order_id") 

65 

66 

67class BenefitGrantScopeComparator(CompositeProperty.Comparator[BenefitGrantScopeArgs]): 

68 def __eq__(self, other: Any) -> ColumnElement[bool]: # type: ignore[override] 

69 if not isinstance(other, dict) or other == {}: 

70 raise ValueError("A non-empty dictionary scope must be provided.") 

71 clauses = [] 

72 composite_columns = self.__clause_element__().columns 

73 for key, value in other.items(): 

74 clauses.append(composite_columns[f"{key}_id"] == value.id) 

75 return and_(*clauses) 

76 

77 

78class BenefitGrant(RecordModel): 

79 """ 

80 Represents a benefit granted to a customer or member. 

81 

82 Unique constraints: 

83 - benefit_grants_sbc_key: Ensures one grant per (subscription, customer, benefit) 

84 - benefit_grants_smb_key: Ensures one grant per (subscription, member, benefit) 

85 

86 These constraints allow both customer-level and member-level benefit grants. 

87 """ 

88 

89 __tablename__ = "benefit_grants" 

90 __table_args__ = ( 

91 UniqueConstraint( 

92 "subscription_id", 

93 "customer_id", 

94 "benefit_id", 

95 name="benefit_grants_sbc_key", 

96 ), 

97 UniqueConstraint( 

98 "subscription_id", 

99 "member_id", 

100 "benefit_id", 

101 name="benefit_grants_smb_key", 

102 ), 

103 ) 

104 

105 granted_at: Mapped[datetime | None] = mapped_column( 

106 TIMESTAMP(timezone=True), nullable=True 

107 ) 

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

109 TIMESTAMP(timezone=True), nullable=True 

110 ) 

111 

112 customer_id: Mapped[UUID] = mapped_column( 

113 Uuid, ForeignKey("customers.id", ondelete="cascade"), nullable=False, index=True 

114 ) 

115 

116 @declared_attr 

117 def customer(cls) -> Mapped["Customer"]: 

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

119 

120 member_id: Mapped[UUID | None] = mapped_column( 

121 Uuid, 

122 ForeignKey("members.id", ondelete="cascade"), 

123 nullable=True, 

124 index=True, 

125 ) 

126 

127 @declared_attr 

128 def member(cls) -> Mapped["Member | None"]: 

129 return relationship("Member", lazy="raise") 

130 

131 benefit_id: Mapped[UUID] = mapped_column( 

132 Uuid, 

133 ForeignKey("benefits.id", ondelete="cascade"), 

134 nullable=False, 

135 index=True, 

136 ) 

137 

138 @declared_attr 

139 def benefit(cls) -> Mapped["Benefit"]: 

140 return relationship("Benefit", lazy="raise", back_populates="grants") 

141 

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

143 Uuid, 

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

145 nullable=True, 

146 # Don't create an index for subscription_id 

147 # as it's covered by the unique constraint, being the leading column of it 

148 index=False, 

149 ) 

150 

151 @declared_attr 

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

153 return relationship("Subscription", lazy="raise", back_populates="grants") 

154 

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

156 Uuid, 

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

158 nullable=True, 

159 index=True, 

160 ) 

161 

162 @declared_attr 

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

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

165 

166 scope: Mapped[BenefitGrantScopeArgs] = composite( 

167 "subscription_id", "order_id", comparator_factory=BenefitGrantScopeComparator 

168 ) 

169 

170 @declared_attr 

171 def properties(cls) -> Mapped["BenefitGrantProperties"]: 

172 return mapped_column("properties", JSONB, nullable=False, default=dict) 

173 

174 @declared_attr 

175 def error(cls) -> Mapped[BenefitGrantError | None]: 

176 return mapped_column( 

177 "error", JSONB(none_as_null=True), nullable=True, default=None 

178 ) 

179 

180 @hybrid_property 

181 def is_granted(self) -> bool: 

182 return self.granted_at is not None 

183 

184 @is_granted.inplace.expression 

185 @classmethod 

186 def _is_granted_expression(cls) -> ColumnElement[bool]: 

187 return type_coerce(cls.granted_at.is_not(None), Boolean) 

188 

189 @hybrid_property 

190 def is_revoked(self) -> bool: 

191 return self.revoked_at is not None 

192 

193 @is_revoked.inplace.expression 

194 @classmethod 

195 def _is_revoked_expression(cls) -> ColumnElement[bool]: 

196 return type_coerce(cls.revoked_at.is_not(None), Boolean) 

197 

198 def set_granted(self) -> None: 

199 self.granted_at = datetime.now(UTC) 

200 self.revoked_at = None 

201 self.error = None 

202 

203 def set_revoked(self) -> None: 

204 self.granted_at = None 

205 self.revoked_at = datetime.now(UTC) 

206 

207 def set_grant_failed(self, error: Exception) -> None: 

208 self.granted_at = None 

209 self.revoked_at = None 

210 self.error = BenefitGrantError( 

211 message=str(error), 

212 type=error.__class__.__name__, 

213 timestamp=datetime.now(UTC).isoformat(), 

214 ) 

215 

216 @property 

217 def previous_properties(self) -> "BenefitGrantProperties | None": 

218 return getattr(self, "_previous_properties", None) 

219 

220 @previous_properties.setter 

221 def previous_properties(self, value: "BenefitGrantProperties") -> None: 

222 self._previous_properties = value