|
| 1 | +"""Write the worker's Google Cloud credentials from its environment.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import contextlib |
| 6 | +import json |
| 7 | +import logging |
| 8 | +import os |
| 9 | +import tempfile |
| 10 | +from typing import TYPE_CHECKING |
| 11 | + |
| 12 | +from policyengine_simulation_observability import scopes |
| 13 | +from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage |
| 14 | + |
| 15 | +if TYPE_CHECKING: |
| 16 | + from collections.abc import Iterator |
| 17 | + |
| 18 | + from policyengine_observability import ObservabilityRuntime |
| 19 | + |
| 20 | +# Moved out of simulation_runtime.py; it still logs under that module's name. |
| 21 | +logger = logging.getLogger("policyengine_simulation_executor.simulation_runtime") |
| 22 | + |
| 23 | +# Credential-adjacent environment values are cut to this many characters. |
| 24 | +_ENV_VALUE_PREVIEW_CHARS = 50 |
| 25 | + |
| 26 | + |
| 27 | +def _normalize_credentials_blob(creds_json: str) -> str: |
| 28 | + """Return the raw JSON blob, decoding the outer escape if present. |
| 29 | +
|
| 30 | + The upstream Modal secret sometimes stores the credentials payload |
| 31 | + double-encoded (the entire JSON object is wrapped in quotes with |
| 32 | + backslash-escaped interior quotes). Historically we always attempted |
| 33 | + the unescape as a fallback which could accidentally parse an already |
| 34 | + clean blob. Only unwrap when the payload looks wrapped. |
| 35 | + """ |
| 36 | + try: |
| 37 | + json.loads(creds_json) |
| 38 | + except json.JSONDecodeError: |
| 39 | + looks_escaped = creds_json.lstrip().startswith('"') or '\\"' in creds_json |
| 40 | + if looks_escaped: |
| 41 | + return json.loads(f'"{creds_json}"') |
| 42 | + raise |
| 43 | + return creds_json |
| 44 | + |
| 45 | + |
| 46 | +@contextlib.contextmanager |
| 47 | +def setup_gcp_credentials( |
| 48 | + runtime: ObservabilityRuntime | None = None, |
| 49 | +) -> Iterator[None]: |
| 50 | + """Set up GCP credentials from environment variable. |
| 51 | +
|
| 52 | + Modal secrets are injected as environment variables. The GCP library |
| 53 | + expects GOOGLE_APPLICATION_CREDENTIALS to point to a file path. If |
| 54 | + credentials JSON is provided, write it to a temp file that's deleted |
| 55 | + on exit. This runs as a context manager to guarantee cleanup even if |
| 56 | + the caller raises mid-simulation; the previous fire-and-forget |
| 57 | + ``tempfile.mkstemp`` path leaked credential material on disk every |
| 58 | + time a container served a request. |
| 59 | + """ |
| 60 | + previous = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") |
| 61 | + creds_file = None |
| 62 | + with contextlib.ExitStack() as credentials_file_stack: |
| 63 | + try: |
| 64 | + credential_span = ( |
| 65 | + scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP)) |
| 66 | + if runtime is not None |
| 67 | + else contextlib.nullcontext() |
| 68 | + ) |
| 69 | + with credential_span: |
| 70 | + # Log available GCP-related env vars for debugging. |
| 71 | + gcp_vars = { |
| 72 | + k: v[:_ENV_VALUE_PREVIEW_CHARS] + "..." |
| 73 | + if len(v) > _ENV_VALUE_PREVIEW_CHARS |
| 74 | + else v |
| 75 | + for k, v in os.environ.items() |
| 76 | + if "GOOGLE" in k or "GCP" in k or "CREDENTIAL" in k |
| 77 | + } |
| 78 | + logger.info("GCP-related env vars: %s", list(gcp_vars.keys())) |
| 79 | + |
| 80 | + if previous: |
| 81 | + logger.info("GOOGLE_APPLICATION_CREDENTIALS already set") |
| 82 | + else: |
| 83 | + creds_json = ( |
| 84 | + os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON") |
| 85 | + or os.environ.get("GCP_CREDENTIALS_JSON") |
| 86 | + or os.environ.get("GOOGLE_CREDENTIALS") |
| 87 | + or os.environ.get("SERVICE_ACCOUNT_JSON") |
| 88 | + ) |
| 89 | + |
| 90 | + if not creds_json: |
| 91 | + logger.warning( |
| 92 | + "No GCP credentials found in environment variables" |
| 93 | + ) |
| 94 | + else: |
| 95 | + normalized = _normalize_credentials_blob(creds_json) |
| 96 | + # The exit stack keeps the file open across the yield, |
| 97 | + # and closing it (``delete=True``) removes the file. |
| 98 | + # The finally block first restores any prior value of |
| 99 | + # ``GOOGLE_APPLICATION_CREDENTIALS`` so a retry in the same |
| 100 | + # container doesn't silently pick up a stale path. |
| 101 | + creds_file = credentials_file_stack.enter_context( |
| 102 | + tempfile.NamedTemporaryFile( |
| 103 | + mode="w", |
| 104 | + suffix=".json", |
| 105 | + delete=True, |
| 106 | + ) |
| 107 | + ) |
| 108 | + creds_file.write(normalized) |
| 109 | + creds_file.flush() |
| 110 | + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = creds_file.name |
| 111 | + logger.info("GCP credentials written to %s", creds_file.name) |
| 112 | + |
| 113 | + yield |
| 114 | + finally: |
| 115 | + if creds_file is not None: |
| 116 | + if previous is None: |
| 117 | + os.environ.pop("GOOGLE_APPLICATION_CREDENTIALS", None) |
| 118 | + else: |
| 119 | + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = previous |
0 commit comments