Coverage for polar/organization/resolver.py: 81%

1from typing import Protocol 

2 

3from pydantic import UUID4 

4from sqlalchemy import select 

5 

6from polar.auth.models import AuthSubject, is_organization 

7from polar.exceptions import PolarRequestValidationError 

8from polar.models import Organization, User, UserOrganization 

9from polar.postgres import AsyncSession 

10 

11 

12class _OrganizationIDModelNone(Protocol): 

13 organization_id: UUID4 | None 

14 

15 

16class _OrganizationIDModel(Protocol): 

17 organization_id: UUID4 

18 

19 

20OrganizationIDModel = _OrganizationIDModelNone | _OrganizationIDModel 

21 

22 

23async def get_payload_organization( 

24 session: AsyncSession, 

25 auth_subject: AuthSubject[User | Organization], 

26 model: OrganizationIDModel, 

27) -> Organization: 

28 # Avoids a circular import :( 

29 

30 if is_organization(auth_subject): 30 ↛ 31line 30 didn't jump to line 31 because the condition on line 30 was never true

31 if model.organization_id is not None: 

32 raise PolarRequestValidationError( 

33 [ 

34 { 

35 "type": "organization_token", 

36 "msg": ( 

37 "Setting organization_id is disallowed " 

38 "when using an organization token." 

39 ), 

40 "loc": ( 

41 "body", 

42 "organization_id", 

43 ), 

44 "input": model.organization_id, 

45 } 

46 ] 

47 ) 

48 return auth_subject.subject 

49 

50 if model.organization_id is None: 

51 raise PolarRequestValidationError( 

52 [ 

53 { 

54 "type": "missing", 

55 "msg": "organization_id is required.", 

56 "loc": ( 

57 "body", 

58 "organization_id", 

59 ), 

60 "input": None, 

61 } 

62 ] 

63 ) 

64 

65 statement = select(Organization).where( 

66 Organization.id == model.organization_id, 

67 Organization.id.in_( 

68 select(UserOrganization.organization_id).where( 

69 UserOrganization.user_id == auth_subject.subject.id, 

70 UserOrganization.deleted_at.is_(None), 

71 ) 

72 ), 

73 ) 

74 result = await session.execute(statement) 

75 organization = result.scalar_one_or_none() 

76 

77 if organization is None: 

78 raise PolarRequestValidationError( 

79 [ 

80 { 

81 "loc": ( 

82 "body", 

83 "organization_id", 

84 ), 

85 "msg": "Organization not found.", 

86 "type": "value_error", 

87 "input": model.organization_id, 

88 } 

89 ] 

90 ) 

91 

92 return organization