Coverage for polar/license_key/repository.py: 42%

1from uuid import UUID 1 ctx1a

2 

3from sqlalchemy import Select, select 1 ctx1a

4from sqlalchemy.orm import joinedload 1 ctx1a

5 

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

7from polar.kit.repository import ( 1 ctx1a

8 Options, 

9 RepositoryBase, 

10 RepositorySoftDeletionIDMixin, 

11 RepositorySoftDeletionMixin, 

12) 

13from polar.models import LicenseKey, Organization, UserOrganization 1 ctx1a

14 

15 

16class LicenseKeyRepository( 1 ctx1a

17 RepositorySoftDeletionIDMixin[LicenseKey, UUID], 

18 RepositorySoftDeletionMixin[LicenseKey], 

19 RepositoryBase[LicenseKey], 

20): 

21 model = LicenseKey 1 ctx1a

22 

23 async def get_by_organization_and_key( 1 ctx1a

24 self, 

25 organization: UUID, 

26 key: str, 

27 *, 

28 include_deleted: bool = False, 

29 options: Options = (), 

30 ) -> LicenseKey | None: 

31 statement = ( 

32 self.get_base_statement(include_deleted=include_deleted) 

33 .where(LicenseKey.organization_id == organization, LicenseKey.key == key) 

34 .options(*options) 

35 ) 

36 return await self.get_one_or_none(statement) 

37 

38 async def get_by_id_organization_customer_and_benefit( 1 ctx1a

39 self, 

40 id: UUID, 

41 organization: UUID, 

42 customer: UUID, 

43 benefit: UUID, 

44 *, 

45 include_deleted: bool = False, 

46 options: Options = (), 

47 ) -> LicenseKey | None: 

48 statement = ( 

49 self.get_base_statement(include_deleted=include_deleted) 

50 .where( 

51 LicenseKey.id == id, 

52 LicenseKey.organization_id == organization, 

53 LicenseKey.customer_id == customer, 

54 LicenseKey.benefit_id == benefit, 

55 ) 

56 .options(*options) 

57 ) 

58 return await self.get_one_or_none(statement) 

59 

60 async def get_readable_by_key( 1 ctx1a

61 self, 

62 key: str, 

63 organization_id: UUID, 

64 auth_subject: AuthSubject[User | Organization], 

65 *, 

66 options: Options = (), 

67 ) -> LicenseKey | None: 

68 statement = ( 

69 self.get_readable_statement(auth_subject) 

70 .where( 

71 LicenseKey.key == key, 

72 LicenseKey.organization_id == organization_id, 

73 ) 

74 .options(*options) 

75 ) 

76 return await self.get_one_or_none(statement) 

77 

78 def get_eager_options(self) -> Options: 1 ctx1a

79 return ( 

80 joinedload(LicenseKey.customer), 

81 joinedload(LicenseKey.activations), 

82 joinedload(LicenseKey.benefit), 

83 ) 

84 

85 def get_readable_statement( 1 ctx1a

86 self, auth_subject: AuthSubject[User | Organization] 

87 ) -> Select[tuple[LicenseKey]]: 

88 statement = self.get_base_statement() 

89 

90 if is_user(auth_subject): 

91 user = auth_subject.subject 

92 statement = statement.where( 

93 LicenseKey.organization_id.in_( 

94 select(UserOrganization.organization_id).where( 

95 UserOrganization.user_id == user.id, 

96 UserOrganization.deleted_at.is_(None), 

97 ) 

98 ) 

99 ) 

100 elif is_organization(auth_subject): 

101 statement = statement.where( 

102 LicenseKey.organization_id == auth_subject.subject.id, 

103 ) 

104 

105 return statement