Coverage for polar/organization/tasks.py: 27%

1import uuid 

2 

3from sqlalchemy.orm import joinedload 

4 

5from polar.account.repository import AccountRepository 

6from polar.email.react import render_email_template 

7from polar.email.schemas import ( 

8 OrganizationReviewedEmail, 

9 OrganizationReviewedProps, 

10 OrganizationUnderReviewEmail, 

11 OrganizationUnderReviewProps, 

12) 

13from polar.email.sender import enqueue_email 

14from polar.exceptions import PolarTaskError 

15from polar.held_balance.service import held_balance as held_balance_service 

16from polar.integrations.plain.service import plain as plain_service 

17from polar.models import Organization 

18from polar.models.organization import OrganizationStatus 

19from polar.user.repository import UserRepository 

20from polar.worker import AsyncSessionMaker, TaskPriority, actor 

21 

22from .repository import OrganizationRepository 

23 

24 

25class OrganizationTaskError(PolarTaskError): ... 25 ↛ 28line 25 didn't jump to line 28 because

26 

27 

28class OrganizationDoesNotExist(OrganizationTaskError): 

29 def __init__(self, organization_id: uuid.UUID) -> None: 

30 self.organization_id = organization_id 

31 message = f"The organization with id {organization_id} does not exist." 

32 super().__init__(message) 

33 

34 

35class OrganizationAccountNotSet(OrganizationTaskError): 

36 def __init__(self, organization_id: uuid.UUID) -> None: 

37 self.organization_id = organization_id 

38 message = ( 

39 f"The organization with id {organization_id} does not have an account set." 

40 ) 

41 super().__init__(message) 

42 

43 

44class AccountDoesNotExist(OrganizationTaskError): 

45 def __init__(self, account_id: uuid.UUID) -> None: 

46 self.account_id = account_id 

47 message = f"The account with id {account_id} does not exist." 

48 super().__init__(message) 

49 

50 

51class UserDoesNotExist(OrganizationTaskError): 

52 def __init__(self, user_id: uuid.UUID) -> None: 

53 self.user_id = user_id 

54 message = f"The user with id {user_id} does not exist." 

55 super().__init__(message) 

56 

57 

58@actor(actor_name="organization.created", priority=TaskPriority.LOW) 

59async def organization_created(organization_id: uuid.UUID) -> None: 

60 async with AsyncSessionMaker() as session: 

61 repository = OrganizationRepository.from_session(session) 

62 organization = await repository.get_by_id(organization_id) 

63 if organization is None: 

64 raise OrganizationDoesNotExist(organization_id) 

65 

66 

67@actor(actor_name="organization.account_set", priority=TaskPriority.LOW) 

68async def organization_account_set(organization_id: uuid.UUID) -> None: 

69 async with AsyncSessionMaker() as session: 

70 repository = OrganizationRepository.from_session(session) 

71 organization = await repository.get_by_id(organization_id) 

72 if organization is None: 

73 raise OrganizationDoesNotExist(organization_id) 

74 

75 if organization.account_id is None: 

76 raise OrganizationAccountNotSet(organization_id) 

77 

78 account_repository = AccountRepository.from_session(session) 

79 account = await account_repository.get_by_id(organization.account_id) 

80 if account is None: 

81 raise AccountDoesNotExist(organization.account_id) 

82 

83 await held_balance_service.release_account(session, account) 

84 

85 

86@actor(actor_name="organization.under_review", priority=TaskPriority.LOW) 

87async def organization_under_review(organization_id: uuid.UUID) -> None: 

88 async with AsyncSessionMaker() as session: 

89 repository = OrganizationRepository.from_session(session) 

90 organization = await repository.get_by_id( 

91 organization_id, options=(joinedload(Organization.account),) 

92 ) 

93 if organization is None: 

94 raise OrganizationDoesNotExist(organization_id) 

95 

96 await plain_service.create_organization_review_thread(session, organization) 

97 

98 # Send an email for the initial review 

99 if organization.status == OrganizationStatus.INITIAL_REVIEW: 

100 admin_user = await repository.get_admin_user(session, organization) 

101 if admin_user: 

102 email = OrganizationUnderReviewEmail( 

103 props=OrganizationUnderReviewProps.model_validate( 

104 {"email": admin_user.email, "organization": organization} 

105 ) 

106 ) 

107 enqueue_email( 

108 to_email_addr=admin_user.email, 

109 subject="Your organization is under review", 

110 html_content=render_email_template(email), 

111 ) 

112 

113 

114@actor(actor_name="organization.reviewed", priority=TaskPriority.LOW) 

115async def organization_reviewed( 

116 organization_id: uuid.UUID, initial_review: bool = False 

117) -> None: 

118 async with AsyncSessionMaker() as session: 

119 repository = OrganizationRepository.from_session(session) 

120 organization = await repository.get_by_id(organization_id) 

121 if organization is None: 

122 raise OrganizationDoesNotExist(organization_id) 

123 

124 # Release held balance if account exists 

125 if organization.account_id: 

126 account_repository = AccountRepository.from_session(session) 

127 account = await account_repository.get_by_id(organization.account_id) 

128 if account: 

129 await held_balance_service.release_account(session, account) 

130 

131 # Send an email after the initial review 

132 if initial_review: 

133 admin_user = await repository.get_admin_user(session, organization) 

134 if admin_user: 

135 email = OrganizationReviewedEmail( 

136 props=OrganizationReviewedProps.model_validate( 

137 {"email": admin_user.email, "organization": organization} 

138 ) 

139 ) 

140 enqueue_email( 

141 to_email_addr=admin_user.email, 

142 subject="Your organization review is complete", 

143 html_content=render_email_template(email), 

144 ) 

145 

146 

147@actor(actor_name="organization.deletion_requested", priority=TaskPriority.HIGH) 

148async def organization_deletion_requested( 

149 organization_id: uuid.UUID, 

150 user_id: uuid.UUID, 

151 blocked_reasons: list[str], 

152) -> None: 

153 """Handle organization deletion request that requires support review.""" 

154 async with AsyncSessionMaker() as session: 

155 repository = OrganizationRepository.from_session(session) 

156 organization = await repository.get_by_id(organization_id) 

157 if organization is None: 

158 raise OrganizationDoesNotExist(organization_id) 

159 

160 user_repository = UserRepository.from_session(session) 

161 user = await user_repository.get_by_id(user_id) 

162 if user is None: 

163 raise UserDoesNotExist(user_id) 

164 

165 # Create Plain ticket for support handling 

166 await plain_service.create_organization_deletion_thread( 

167 session, organization, user, blocked_reasons 

168 )