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

1from datetime import UTC, datetime 2 ctx1ab

2from typing import TYPE_CHECKING, Any, TypedDict 2 ctx1ab

3from uuid import UUID 2 ctx1ab

4 

5from sqlalchemy import ( 2 ctx1ab

6 TIMESTAMP, 

7 Boolean, 

8 ColumnElement, 

9 ForeignKey, 

10 UniqueConstraint, 

11 Uuid, 

12 and_, 

13 type_coerce, 

14) 

15from sqlalchemy.dialects.postgresql import JSONB 2 ctx1ab

16from sqlalchemy.ext.hybrid import hybrid_property 2 ctx1ab

17from sqlalchemy.orm import ( 2 ctx1ab

18 CompositeProperty, 

19 Mapped, 

20 composite, 

21 declared_attr, 

22 mapped_column, 

23 relationship, 

24) 

25 

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

27 

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

29 from polar.benefit.strategies import BenefitGrantProperties 

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

31 

32 

33class BenefitGrantError(TypedDict): 2 ctx1ab

34 message: str 2 ctx1ab

35 type: str 2 ctx1ab

36 timestamp: str 2 ctx1ab

37 

38 

39class BenefitGrantScope(TypedDict, total=False): 2 ctx1ab

40 subscription: "Subscription" 2 ctx1ab

41 order: "Order" 2 ctx1ab

42 

43 

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

45 

46 class BenefitGrantScopeArgs(TypedDict, total=False): 

47 subscription_id: UUID 

48 order_id: UUID 

49 

50else: 

51 

52 class BenefitGrantScopeArgs(dict): 2 ctx1ab

53 def __init__( 2 ctx1ab

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

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

65 

66 

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

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

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

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

90 __table_args__ = ( 2 ctx1ab

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

106 TIMESTAMP(timezone=True), nullable=True 

107 ) 

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

109 TIMESTAMP(timezone=True), nullable=True 

110 ) 

111 

112 customer_id: Mapped[UUID] = mapped_column( 2 ctx1ab

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

114 ) 

115 

116 @declared_attr 2 ctx1ab

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

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

119 

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

121 Uuid, 

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

123 nullable=True, 

124 index=True, 

125 ) 

126 

127 @declared_attr 2 ctx1ab

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

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

130 

131 benefit_id: Mapped[UUID] = mapped_column( 2 ctx1ab

132 Uuid, 

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

134 nullable=False, 

135 index=True, 

136 ) 

137 

138 @declared_attr 2 ctx1ab

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

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

141 

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

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

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

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

154 

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

156 Uuid, 

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

158 nullable=True, 

159 index=True, 

160 ) 

161 

162 @declared_attr 2 ctx1ab

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

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

165 

166 scope: Mapped[BenefitGrantScopeArgs] = composite( 2 ctx1ab

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

168 ) 

169 

170 @declared_attr 2 ctx1ab

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

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

173 

174 @declared_attr 2 ctx1ab

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

176 return mapped_column( 2 ctx1ab

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

178 ) 

179 

180 @hybrid_property 2 ctx1ab

181 def is_granted(self) -> bool: 2 ctx1ab

182 return self.granted_at is not None 

183 

184 @is_granted.inplace.expression 2 ctx1ab

185 @classmethod 2 ctx1ab

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

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

188 

189 @hybrid_property 2 ctx1ab

190 def is_revoked(self) -> bool: 2 ctx1ab

191 return self.revoked_at is not None 

192 

193 @is_revoked.inplace.expression 2 ctx1ab

194 @classmethod 2 ctx1ab

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

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

197 

198 def set_granted(self) -> None: 2 ctx1ab

199 self.granted_at = datetime.now(UTC) 

200 self.revoked_at = None 

201 self.error = None 

202 

203 def set_revoked(self) -> None: 2 ctx1ab

204 self.granted_at = None 

205 self.revoked_at = datetime.now(UTC) 

206 

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

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

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

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

219 

220 @previous_properties.setter 2 ctx1ab

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

222 self._previous_properties = value