estuary-cdk: add support for GoogleServiceAccount credentials

Google service account credentials will be needed for the upcoming
`source-google-play` connector. There are also other Google connectors
(like `source-google-analytics-data-api-native` and probably others)
that would could add support for service accounts.

I considered not relying on the `google-auth` package and manually
exchange the service account JSON for an access token. But after digging
into how the `google-auth` package handles that, it's a significant
amount of code that could be pretty fragile if Google decides to tweak
that process in any way.

Author: Alex-Bair

estuary-cdk/estuary_cdk/flow.py

@@ -303,3 +303,20 @@ class _OAuth2Credentials(RotatingOAuth2Credentials):
             def _you_must_build_oauth2_credentials_for_a_provider(self): ...
 
         return _OAuth2Credentials
+
+
+class GoogleServiceAccountSpec(BaseModel):
+    scopes: list[str]
+
+
+class GoogleServiceAccount(BaseModel):
+    credentials_title: Literal["Google Service Account"] = Field(
+        default="Google Service Account",
+        json_schema_extra={"type": "string", "order": 0}
+    )
+    service_account: str = Field(
+        title="Google Service Account",
+        description="Service account JSON key",
+        json_schema_extra={"secret": True, "multiline": True, "order": 1},
+    )
+

estuary-cdk/estuary_cdk/http.py

@@ -6,8 +6,14 @@
 import aiohttp
 import asyncio
 import base64
+import json
 import time
 
+
+from google.auth.credentials import TokenState as GoogleTokenState
+from google.auth.transport.requests import Request as GoogleAuthRequest
+from google.oauth2.service_account import Credentials as GoogleServiceAccountCredentials
+
 from . import Mixin
 from .flow import (
     AccessToken,
@@ -21,6 +27,8 @@
     RotatingOAuth2Credentials,
     OAuth2Spec,
     OAuth2RotatingTokenSpec,
+    GoogleServiceAccount,
+    GoogleServiceAccountSpec,
 )
 
 DEFAULT_AUTHORIZATION_HEADER = "Authorization"
@@ -205,9 +213,11 @@ class AccessTokenResponse(BaseModel):
         | LongLivedClientCredentialsOAuth2Credentials
         | AccessToken
         | BasicAuth
+        | GoogleServiceAccount
     )
     authorization_header: str = DEFAULT_AUTHORIZATION_HEADER
-    _access_token: AccessTokenResponse | None = None
+    google_spec: GoogleServiceAccountSpec | None = None
+    _access_token: AccessTokenResponse | GoogleServiceAccountCredentials | None = None
     _fetched_at: int = 0
 
     async def fetch_token(self, log: Logger, session: HTTPSession) -> tuple[str, str]:
@@ -230,6 +240,25 @@ async def fetch_token(self, log: Logger, session: HTTPSession) -> tuple[str, str
                     f"{self.credentials.username}:{self.credentials.password}".encode()
                 ).decode(),
             )
+        elif isinstance(self.credentials, GoogleServiceAccount):
+            assert isinstance(self.google_spec, GoogleServiceAccountSpec)
+            if self._access_token is None:
+                self._access_token = GoogleServiceAccountCredentials.from_service_account_info(
+                    json.loads(self.credentials.service_account),
+                    scopes=self.google_spec.scopes,
+                )
+
+            assert isinstance(self._access_token, GoogleServiceAccountCredentials)
+
+            match self._access_token.token_state:
+                case GoogleTokenState.FRESH:
+                    pass
+                case GoogleTokenState.STALE | GoogleTokenState.INVALID:
+                    self._access_token.refresh(GoogleAuthRequest())
+                case _:
+                    raise RuntimeError(f"Unknown GoogleTokenState: {self._access_token.token_state}")
+
+            return ("Bearer", self._access_token.token)
 
         assert (
             isinstance(self.credentials, BaseOAuth2Credentials)
@@ -240,6 +269,7 @@ async def fetch_token(self, log: Logger, session: HTTPSession) -> tuple[str, str
         current_time = time.time()
 
         if self._access_token is not None:
+            assert isinstance(self._access_token, self.AccessTokenResponse)
             horizon = self._fetched_at + self._access_token.expires_in * 0.75
 
             if current_time < horizon:

estuary-cdk/poetry.lock

@@ -19,6 +19,7 @@ pycron = "^3.1.2"
 stream-unzip = "0.0.99"
 
 airbyte-cdk = "<=6.56.0"
+google-auth = ">=2.40.3"
 pendulum = "<=2.0.0 | >=3.0.0"
 nltk = "^3.9.1"