From d2504ac5ed0e2e56aaabb0a1a0d1a2cca5cc326c Mon Sep 17 00:00:00 2001 From: Chris Bartholomew Date: Sat, 7 Mar 2026 02:59:51 -0500 Subject: [PATCH] Fix GCS auth for Workload Identity Federation credentials (#518) * Fix GCS auth for external_account credentials (Workload Identity) obstore's built-in credential parsing only supports service_account and authorized_user JSON types. Use google.auth as a credential_provider callback to support all credential types including external_account (Workload Identity Federation), impersonated credentials, and metadata server credentials. * Hide GOOGLE_APPLICATION_CREDENTIALS during GCSStore construction GCSStore eagerly parses the credential file from env vars even when a custom credential_provider is passed. Temporarily unset the env var during construction so obstore doesn't choke on external_account credential files (Workload Identity Federation). * Support HINDSIGHT_GOOGLE_CREDENTIALS_FILE for GCS auth When GOOGLE_APPLICATION_CREDENTIALS must be unset to prevent obstore from parsing unsupported credential types (e.g. external_account), google.auth can load credentials from HINDSIGHT_GOOGLE_CREDENTIALS_FILE instead. This avoids mutating env vars at runtime. * Simplify GCS credential workaround: hide env var during construction Remove HINDSIGHT_GOOGLE_CREDENTIALS_FILE indirection. Instead, let google.auth.default() load credentials normally via GOOGLE_APPLICATION_CREDENTIALS, then temporarily hide the env var during GCSStore() construction so obstore doesn't try to parse credential types it doesn't support. * Work around obstore bug: hide env var during GCSStore construction obstore always parses credential files from GOOGLE_APPLICATION_CREDENTIALS and the well-known ADC path, even when credential_provider is supplied (contrary to docs). This crashes on external_account credentials from Workload Identity Federation. Temporarily hide the env var during GCSStore() construction. google.auth has already loaded credentials by this point via credential_provider. --- .../hindsight_api/engine/storage/gcs.py | 50 ++++++++++++++++++- 1 file changed, 48 insertions(+), 2 deletions(-) diff --git a/hindsight-api/hindsight_api/engine/storage/gcs.py b/hindsight-api/hindsight_api/engine/storage/gcs.py index f44968bf..6013a193 100644 --- a/hindsight-api/hindsight_api/engine/storage/gcs.py +++ b/hindsight-api/hindsight_api/engine/storage/gcs.py @@ -1,7 +1,8 @@ """Google Cloud Storage backend using obstore.""" import logging -from datetime import timedelta +import os +from datetime import datetime, timedelta, timezone import obstore as obs from obstore.store import GCSStore @@ -11,6 +12,30 @@ from .base import FileStorage logger = logging.getLogger(__name__) +def _make_google_auth_credential_provider(): + """Create a credential provider using google.auth (supports all credential types). + + obstore's built-in credential parsing only supports service_account and + authorized_user JSON types. This provider uses the google-auth library + which additionally handles external_account (Workload Identity Federation), + impersonated credentials, and metadata-server credentials. + """ + import google.auth + import google.auth.transport.requests + + credentials, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"]) + request = google.auth.transport.requests.Request() + + def _provide(): + credentials.refresh(request) + expiry = credentials.expiry + if expiry and expiry.tzinfo is None: + expiry = expiry.replace(tzinfo=timezone.utc) + return {"token": credentials.token, "expires_at": expiry} + + return _provide + + class GCSFileStorage(FileStorage): """ Google Cloud Storage backend. @@ -27,8 +52,29 @@ class GCSFileStorage(FileStorage): kwargs: dict = {} if service_account_key: kwargs["service_account_key"] = service_account_key + else: + # Use google.auth credential provider for broad credential type support + # (service_account, authorized_user, external_account, metadata server, etc.) + try: + kwargs["credential_provider"] = _make_google_auth_credential_provider() + logger.info("Using google.auth credential provider for GCS") + except Exception as e: + logger.warning( + f"Failed to create google.auth credential provider, falling back to obstore defaults: {e}" + ) - self._store = GCSStore(bucket, **kwargs) + # Workaround for https://github.com/developmentseed/obstore/issues/605 + # obstore's Rust layer doesn't support external_account credentials (Workload + # Identity Federation) and eagerly parses GOOGLE_APPLICATION_CREDENTIALS even + # when credential_provider is given. Per the obstore maintainer's guidance, + # remove env vars so the Rust code doesn't try to authenticate itself. + # google.auth (used by credential_provider above) has already loaded credentials. + gac = os.environ.pop("GOOGLE_APPLICATION_CREDENTIALS", None) + try: + self._store = GCSStore(bucket, **kwargs) + finally: + if gac is not None: + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = gac logger.info(f"Initialized GCS file storage: bucket={bucket}") async def store(self, file_data: bytes, key: str, metadata: dict[str, str] | None = None) -> str: