Coverage for polar/benefit/grant/repository.py: 70%

1from collections.abc import Sequence 

2from typing import Unpack 

3from uuid import UUID 

4 

5from sqlalchemy import select 

6 

7from polar.kit.repository import ( 

8 Options, 

9 RepositoryBase, 

10 RepositorySoftDeletionIDMixin, 

11 RepositorySoftDeletionMixin, 

12 RepositorySortingMixin, 

13 SortingClause, 

14) 

15from polar.models import Benefit, BenefitGrant, Customer, Product, ProductBenefit 

16from polar.models.benefit import BenefitType 

17from polar.models.benefit_grant import BenefitGrantScope 

18 

19from .sorting import BenefitGrantSortProperty 

20 

21 

22class BenefitGrantRepository( 

23 RepositorySortingMixin[BenefitGrant, BenefitGrantSortProperty], 

24 RepositorySoftDeletionIDMixin[BenefitGrant, UUID], 

25 RepositorySoftDeletionMixin[BenefitGrant], 

26 RepositoryBase[BenefitGrant], 

27): 

28 model = BenefitGrant 

29 

30 async def get_by_benefit_and_scope( 

31 self, customer: Customer, benefit: Benefit, **scope: Unpack[BenefitGrantScope] 

32 ) -> BenefitGrant | None: 

33 statement = self.get_base_statement().where( 

34 BenefitGrant.customer_id == customer.id, 

35 BenefitGrant.benefit_id == benefit.id, 

36 BenefitGrant.deleted_at.is_(None), 

37 BenefitGrant.scope == scope, 

38 ) 

39 return await self.get_one_or_none(statement) 

40 

41 async def list_granted_by_scope( 

42 self, **scope: Unpack[BenefitGrantScope] 

43 ) -> Sequence[BenefitGrant]: 

44 statement = self.get_base_statement().where( 

45 BenefitGrant.scope == scope, 

46 BenefitGrant.is_granted.is_(True), 

47 BenefitGrant.deleted_at.is_(None), 

48 ) 

49 return await self.get_all(statement) 

50 

51 async def list_granted_by_benefit( 

52 self, 

53 benefit: Benefit, 

54 *, 

55 options: Options = (), 

56 ) -> Sequence[BenefitGrant]: 

57 statement = ( 

58 self.get_base_statement() 

59 .where( 

60 BenefitGrant.benefit_id == benefit.id, 

61 BenefitGrant.is_granted.is_(True), 

62 BenefitGrant.deleted_at.is_(None), 

63 ) 

64 .options(*options) 

65 ) 

66 return await self.get_all(statement) 

67 

68 async def list_granted_by_customer( 

69 self, 

70 customer_id: UUID, 

71 *, 

72 options: Options = (), 

73 ) -> Sequence[BenefitGrant]: 

74 statement = ( 

75 self.get_base_statement() 

76 .where( 

77 BenefitGrant.customer_id == customer_id, 

78 BenefitGrant.is_granted.is_(True), 

79 BenefitGrant.deleted_at.is_(None), 

80 ) 

81 .options(*options) 

82 ) 

83 return await self.get_all(statement) 

84 

85 async def list_granted_by_benefit_and_customer( 

86 self, 

87 benefit: Benefit, 

88 customer: Customer, 

89 *, 

90 options: Options = (), 

91 ) -> Sequence[BenefitGrant]: 

92 statement = ( 

93 self.get_base_statement() 

94 .where( 

95 BenefitGrant.benefit_id == benefit.id, 

96 BenefitGrant.customer_id == customer.id, 

97 BenefitGrant.is_granted.is_(True), 

98 BenefitGrant.deleted_at.is_(None), 

99 ) 

100 .options(*options) 

101 ) 

102 return await self.get_all(statement) 

103 

104 async def list_by_customer_and_benefit_type( 

105 self, 

106 customer: Customer, 

107 benefit_type: BenefitType, 

108 *, 

109 options: Options = (), 

110 ) -> Sequence[BenefitGrant]: 

111 statement = ( 

112 self.get_base_statement() 

113 .join(Benefit) 

114 .where( 

115 BenefitGrant.customer_id == customer.id, 

116 Benefit.type == benefit_type, 

117 ) 

118 ).options(*options) 

119 return await self.get_all(statement) 

120 

121 async def list_outdated_grants( 

122 self, product: Product, **scope: Unpack[BenefitGrantScope] 

123 ) -> Sequence[BenefitGrant]: 

124 product_benefits_statement = ( 

125 select(Benefit.id) 

126 .join(ProductBenefit) 

127 .where(ProductBenefit.product_id == product.id) 

128 ) 

129 statement = self.get_base_statement().where( 

130 BenefitGrant.scope == scope, 

131 BenefitGrant.benefit_id.not_in(product_benefits_statement), 

132 BenefitGrant.is_granted.is_(True), 

133 BenefitGrant.deleted_at.is_(None), 

134 ) 

135 return await self.get_all(statement) 

136 

137 def get_sorting_clause(self, property: BenefitGrantSortProperty) -> SortingClause: 

138 match property: 

139 case BenefitGrantSortProperty.created_at: 

140 return BenefitGrant.created_at 

141 case BenefitGrantSortProperty.granted_at: 

142 return BenefitGrant.granted_at 

143 case BenefitGrantSortProperty.revoked_at: 143 ↛ exitline 143 didn't return from function 'get_sorting_clause' because the pattern on line 143 always matched

144 return BenefitGrant.revoked_at