Coverage for polar/customer_meter/repository.py: 22%

1from collections.abc import Sequence 1 ctx1a

2from uuid import UUID 1 ctx1a

3 

4from sqlalchemy import Select, select 1 ctx1a

5from sqlalchemy.orm import contains_eager 1 ctx1a

6 

7from polar.auth.models import AuthSubject, Organization, User, is_organization, is_user 1 ctx1a

8from polar.kit.repository import ( 1 ctx1a

9 RepositoryBase, 

10 RepositorySoftDeletionIDMixin, 

11 RepositorySoftDeletionMixin, 

12 RepositorySortingMixin, 

13 SortingClause, 

14) 

15from polar.kit.repository.base import Options 1 ctx1a

16from polar.models import Customer, CustomerMeter, Meter, UserOrganization 1 ctx1a

17 

18from .sorting import CustomerMeterSortProperty 1 ctx1a

19 

20 

21class CustomerMeterRepository( 1 ctx1a

22 RepositorySortingMixin[CustomerMeter, CustomerMeterSortProperty], 

23 RepositorySoftDeletionIDMixin[CustomerMeter, UUID], 

24 RepositorySoftDeletionMixin[CustomerMeter], 

25 RepositoryBase[CustomerMeter], 

26): 

27 model = CustomerMeter 1 ctx1a

28 

29 async def get_all_by_customer( 1 ctx1a

30 self, customer_id: UUID, *, options: Options = () 

31 ) -> Sequence[CustomerMeter]: 

32 statement = ( 

33 self.get_base_statement() 

34 .where(CustomerMeter.customer_id == customer_id) 

35 .options(*options) 

36 .order_by( 

37 CustomerMeter.created_at.asc(), 

38 ) 

39 ) 

40 return await self.get_all(statement) 

41 

42 async def get_by_customer_and_meter( 1 ctx1a

43 self, 

44 customer_id: UUID, 

45 meter_id: UUID, 

46 *, 

47 options: Options = (), 

48 ) -> CustomerMeter | None: 

49 statement = ( 

50 self.get_base_statement() 

51 .where( 

52 CustomerMeter.customer_id == customer_id, 

53 CustomerMeter.meter_id == meter_id, 

54 ) 

55 .options(*options) 

56 ) 

57 return await self.get_one_or_none(statement) 

58 

59 def get_readable_statement( 1 ctx1a

60 self, auth_subject: AuthSubject[User | Organization] 

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

62 statement = ( 

63 self.get_base_statement() 

64 .join(CustomerMeter.customer) 

65 .options( 

66 contains_eager(CustomerMeter.customer), 

67 ) 

68 ) 

69 

70 if is_user(auth_subject): 

71 user = auth_subject.subject 

72 statement = statement.where( 

73 Customer.organization_id.in_( 

74 select(UserOrganization.organization_id).where( 

75 UserOrganization.user_id == user.id, 

76 UserOrganization.deleted_at.is_(None), 

77 ) 

78 ) 

79 ) 

80 elif is_organization(auth_subject): 

81 statement = statement.where( 

82 Customer.organization_id == auth_subject.subject.id, 

83 ) 

84 

85 return statement 

86 

87 def get_sorting_clause(self, property: CustomerMeterSortProperty) -> SortingClause: 1 ctx1a

88 match property: 

89 case CustomerMeterSortProperty.created_at: 

90 return self.model.created_at 

91 case CustomerMeterSortProperty.modified_at: 

92 return self.model.modified_at 

93 case CustomerMeterSortProperty.customer_id: 

94 return self.model.customer_id 

95 case CustomerMeterSortProperty.customer_name: 

96 return Customer.name 

97 case CustomerMeterSortProperty.meter_id: 

98 return self.model.meter_id 

99 case CustomerMeterSortProperty.meter_name: 

100 return Meter.name 

101 case CustomerMeterSortProperty.consumed_units: 

102 return self.model.consumed_units 

103 case CustomerMeterSortProperty.credited_units: 

104 return self.model.credited_units 

105 case CustomerMeterSortProperty.balance: 

106 return self.model.balance