Coverage for polar/customer_portal/repository/customer_meter.py: 98%

1from uuid import UUID 

2 

3from sqlalchemy import Select, or_, select 

4from sqlalchemy.dialects.postgresql import UUID as PGUUID 

5 

6from polar.auth.models import AuthSubject, Customer 

7from polar.kit.repository import ( 

8 RepositoryBase, 

9 RepositorySoftDeletionIDMixin, 

10 RepositorySoftDeletionMixin, 

11) 

12from polar.kit.repository.base import RepositorySortingMixin, SortingClause 

13from polar.models import ( 

14 Benefit, 

15 BenefitGrant, 

16 CustomerMeter, 

17 Meter, 

18 Subscription, 

19 SubscriptionMeter, 

20) 

21from polar.models.benefit import BenefitType 

22from polar.models.subscription import SubscriptionStatus 

23 

24from ..sorting.customer_meter import CustomerCustomerMeterSortProperty 

25 

26 

27class CustomerMeterRepository( 

28 RepositorySortingMixin[CustomerMeter, CustomerCustomerMeterSortProperty], 

29 RepositorySoftDeletionIDMixin[CustomerMeter, UUID], 

30 RepositorySoftDeletionMixin[CustomerMeter], 

31 RepositoryBase[CustomerMeter], 

32): 

33 model = CustomerMeter 

34 

35 def get_readable_statement( 

36 self, auth_subject: AuthSubject[Customer] 

37 ) -> Select[tuple[CustomerMeter]]: 

38 customer_id = auth_subject.subject.id 

39 

40 # Subquery for meters from active subscriptions 

41 subscription_meters = ( 

42 select(SubscriptionMeter.meter_id) 

43 .select_from(SubscriptionMeter) 

44 .join(Subscription, Subscription.id == SubscriptionMeter.subscription_id) 

45 .where( 

46 Subscription.customer_id == customer_id, 

47 Subscription.status.in_(SubscriptionStatus.active_statuses()), 

48 ) 

49 ) 

50 

51 # Subquery for meters from benefit grants (one-time purchases) 

52 benefit_meters = ( 

53 select(Benefit.properties["meter_id"].astext.cast(PGUUID)) 

54 .select_from(BenefitGrant) 

55 .join(Benefit, Benefit.id == BenefitGrant.benefit_id) 

56 .where( 

57 BenefitGrant.customer_id == customer_id, 

58 BenefitGrant.granted_at.is_not(None), 

59 BenefitGrant.revoked_at.is_(None), 

60 Benefit.type == BenefitType.meter_credit.value, 

61 ) 

62 ) 

63 

64 return ( 

65 self.get_base_statement() 

66 .join(CustomerMeter.meter) 

67 .where( 

68 CustomerMeter.customer_id == customer_id, 

69 or_( 

70 CustomerMeter.meter_id.in_(subscription_meters), 

71 CustomerMeter.meter_id.in_(benefit_meters), 

72 ), 

73 Meter.archived_at.is_(None), 

74 ) 

75 ) 

76 

77 def get_sorting_clause( 

78 self, property: CustomerCustomerMeterSortProperty 

79 ) -> SortingClause: 

80 match property: 

81 case CustomerCustomerMeterSortProperty.created_at: 

82 return self.model.created_at 

83 case CustomerCustomerMeterSortProperty.modified_at: 

84 return self.model.modified_at 

85 case CustomerCustomerMeterSortProperty.meter_id: 

86 return self.model.meter_id 

87 case CustomerCustomerMeterSortProperty.meter_name: 

88 return Meter.name 

89 case CustomerCustomerMeterSortProperty.consumed_units: 

90 return self.model.consumed_units 

91 case CustomerCustomerMeterSortProperty.credited_units: 

92 return self.model.credited_units 

93 case CustomerCustomerMeterSortProperty.balance: 93 ↛ exitline 93 didn't return from function 'get_sorting_clause' because the pattern on line 93 always matched

94 return self.model.balance