Coverage for polar/integrations/google/endpoints.py: 37%

66 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-12-05 16:17 +0000

1import uuid 1a

2from typing import Any 1a

3 

4from fastapi import Depends, Request 1a

5from fastapi.responses import RedirectResponse 1a

6from httpx_oauth.integrations.fastapi import OAuth2AuthorizeCallback 1a

7from httpx_oauth.oauth2 import OAuth2Token 1a

8 

9from polar.auth.dependencies import WebUserOrAnonymous 1a

10from polar.auth.models import is_user 1a

11from polar.auth.service import auth as auth_service 1a

12from polar.config import settings 1a

13from polar.exceptions import PolarRedirectionError 1a

14from polar.integrations.loops.service import loops as loops_service 1a

15from polar.kit import jwt 1a

16from polar.kit.http import ReturnTo 1a

17from polar.openapi import APITag 1a

18from polar.postgres import AsyncSession, get_db_session 1a

19from polar.posthog import posthog 1a

20from polar.routing import APIRouter 1a

21from polar.user.schemas import UserSignupAttribution, UserSignupAttributionQuery 1a

22 

23from .service import GoogleServiceError, google_oauth_client 1a

24from .service import google as google_service 1a

25 

26oauth2_authorize_callback = OAuth2AuthorizeCallback( 1a

27 google_oauth_client, route_name="integrations.google.callback" 

28) 

29 

30 

31class OAuthCallbackError(PolarRedirectionError): ... 1a

32 

33 

34router = APIRouter( 1a

35 prefix="/integrations/google", 

36 tags=["integrations_google", APITag.private], 

37) 

38 

39 

40@router.get("/authorize", name="integrations.google.authorize") 1a

41async def google_authorize( 1a

42 request: Request, 

43 auth_subject: WebUserOrAnonymous, 

44 return_to: ReturnTo, 

45 signup_attribution: UserSignupAttributionQuery, 

46) -> RedirectResponse: 

47 state: dict[str, Any] = {} 

48 

49 state["return_to"] = return_to 

50 

51 if signup_attribution: 

52 state["signup_attribution"] = signup_attribution.model_dump(exclude_unset=True) 

53 

54 if is_user(auth_subject): 

55 state["user_id"] = str(auth_subject.subject.id) 

56 

57 encoded_state = jwt.encode(data=state, secret=settings.SECRET, type="google_oauth") 

58 redirect_uri = str(request.url_for("integrations.google.callback")) 

59 authorization_url = await google_oauth_client.get_authorization_url( 

60 redirect_uri=redirect_uri, 

61 state=encoded_state, 

62 scope=[ 

63 "https://www.googleapis.com/auth/userinfo.profile", 

64 "https://www.googleapis.com/auth/userinfo.email", 

65 ], 

66 ) 

67 return RedirectResponse(authorization_url, 303) 

68 

69 

70@router.get("/callback", name="integrations.google.callback") 1a

71async def google_callback( 1a

72 request: Request, 

73 auth_subject: WebUserOrAnonymous, 

74 session: AsyncSession = Depends(get_db_session), 

75 access_token_state: tuple[OAuth2Token, str | None] = Depends( 

76 oauth2_authorize_callback 

77 ), 

78) -> RedirectResponse: 

79 token_data, state = access_token_state 

80 error_description = token_data.get("error_description") 

81 if error_description: 

82 raise OAuthCallbackError(error_description) 

83 if not state: 

84 raise OAuthCallbackError("No state") 

85 

86 try: 

87 state_data = jwt.decode( 

88 token=state, secret=settings.SECRET, type="google_oauth" 

89 ) 

90 except jwt.DecodeError as e: 

91 raise OAuthCallbackError("Invalid state") from e 

92 

93 return_to = state_data.get("return_to", None) 

94 state_user_id = state_data.get("user_id") 

95 

96 state_signup_attribution = state_data.get("signup_attribution") 

97 if state_signup_attribution: 

98 state_signup_attribution = UserSignupAttribution.model_validate( 

99 state_signup_attribution 

100 ) 

101 

102 try: 

103 if ( 

104 is_user(auth_subject) 

105 and state_user_id is not None 

106 and auth_subject.subject.id == uuid.UUID(state_user_id) 

107 ): 

108 is_signup = False 

109 user = await google_service.link_user( 

110 session, user=auth_subject.subject, token=token_data 

111 ) 

112 else: 

113 user, is_signup = await google_service.get_updated_or_create( 

114 session, 

115 token=token_data, 

116 signup_attribution=state_signup_attribution, 

117 ) 

118 except GoogleServiceError as e: 

119 raise OAuthCallbackError(e.message, e.status_code, return_to=return_to) from e 

120 

121 # Event tracking last to ensure business critical data is stored first 

122 if is_signup: 

123 posthog.user_signup(user, "google") 

124 await loops_service.user_signup(user, googleLogin=True) 

125 else: 

126 posthog.user_login(user, "google") 

127 await loops_service.user_update(session, user, googleLogin=True) 

128 

129 return await auth_service.get_login_response( 

130 session, request, user, return_to=return_to 

131 )