From 19495eedbe3ba934310f65ee6f320abd12da9007 Mon Sep 17 00:00:00 2001 From: Amir Fathi Date: Tue, 8 Sep 2026 23:24:33 -0600 Subject: [PATCH] fix(bigquery): support Application Default Credentials (#2726) --- core/wren/docs/connections.md | 6 + core/wren/src/wren/connector/bigquery.py | 35 +++-- core/wren/src/wren/model/__init__.py | 13 +- core/wren/src/wren/model/field_registry.py | 4 +- core/wren/tests/unit/test_bigquery_adc.py | 175 +++++++++++++++++++++ 5 files changed, 215 insertions(+), 18 deletions(-) create mode 100644 core/wren/tests/unit/test_bigquery_adc.py diff --git a/core/wren/docs/connections.md b/core/wren/docs/connections.md index 0b6280565..5bce2e1ab 100644 --- a/core/wren/docs/connections.md +++ b/core/wren/docs/connections.md @@ -75,6 +75,12 @@ Both formats are accepted. The CLI auto-flattens the envelope format. } ``` +`credentials` is optional: if omitted, the connector falls back to +Application Default Credentials (`gcloud auth application-default login`, +workload identity, or the GCE/Cloud Run metadata server). ADC obtained via +`gcloud auth application-default login` does not carry the Drive scope; +BigQuery external tables over Drive/Sheets still need a service account. + ## Snowflake ```json diff --git a/core/wren/src/wren/connector/bigquery.py b/core/wren/src/wren/connector/bigquery.py index 4dd5115d6..0b145ee09 100644 --- a/core/wren/src/wren/connector/bigquery.py +++ b/core/wren/src/wren/connector/bigquery.py @@ -21,26 +21,33 @@ def _apply_limit(sql: str, limit: int) -> str: return f"SELECT * FROM ({cleaned}) AS _sub LIMIT {limit}" +_SCOPES = [ + "https://www.googleapis.com/auth/drive", + "https://www.googleapis.com/auth/cloud-platform", +] + + class BigQueryConnector(ConnectorABC): def __init__(self, connection_info): from google.cloud import bigquery # noqa: PLC0415 - from google.oauth2 import service_account # noqa: PLC0415 self.connection_info = connection_info - credits_json = loads( - base64.b64decode(connection_info.credentials.get_secret_value()).decode( - "utf-8" + if connection_info.credentials: + from google.oauth2 import service_account # noqa: PLC0415 + + credits_json = loads( + base64.b64decode(connection_info.credentials.get_secret_value()).decode( + "utf-8" + ) ) - ) - credentials = service_account.Credentials.from_service_account_info( - credits_json - ) - credentials = credentials.with_scopes( - [ - "https://www.googleapis.com/auth/drive", - "https://www.googleapis.com/auth/cloud-platform", - ] - ) + credentials = service_account.Credentials.from_service_account_info( + credits_json + ) + credentials = credentials.with_scopes(_SCOPES) + else: + import google.auth # noqa: PLC0415 + + credentials, _ = google.auth.default(scopes=_SCOPES) client = bigquery.Client( credentials=credentials, project=connection_info.get_billing_project_id(), diff --git a/core/wren/src/wren/model/__init__.py b/core/wren/src/wren/model/__init__.py index 22fb8c47a..b0ba2b4da 100644 --- a/core/wren/src/wren/model/__init__.py +++ b/core/wren/src/wren/model/__init__.py @@ -38,8 +38,17 @@ class BaseConnectionInfo(BaseModel): class BigQueryConnectionInfo(BaseConnectionInfo): - credentials: SecretStr = Field( - description="Base64 encode `credentials.json`", examples=["eyJ..."] + credentials: SecretStr | None = Field( + default=None, + description=( + "Base64 encode `credentials.json`. Omit to use Application Default " + "Credentials (gcloud auth application-default login, workload " + "identity, or the GCE/Cloud Run metadata server). ADC obtained via " + "gcloud auth application-default login does not carry the Drive " + "scope; BigQuery external tables over Drive/Sheets still require a " + "service account." + ), + examples=["eyJ..."], ) job_timeout_ms: int | None = Field(default=None) diff --git a/core/wren/src/wren/model/field_registry.py b/core/wren/src/wren/model/field_registry.py index b3c4596c8..d227beec2 100644 --- a/core/wren/src/wren/model/field_registry.py +++ b/core/wren/src/wren/model/field_registry.py @@ -94,14 +94,14 @@ _MODEL_UI_OVERRIDES: dict[str, dict[str, dict]] = { "credentials": { "input_type": "file_base64", "accept": ".json", - "hint": "Upload your GCP service account credentials.json file. It will be base64-encoded automatically.", + "hint": "Upload your GCP service account credentials.json file (it will be base64-encoded automatically), or leave blank to use Application Default Credentials.", }, }, "BigQueryProjectConnectionInfo": { "credentials": { "input_type": "file_base64", "accept": ".json", - "hint": "Upload your GCP service account credentials.json file. It will be base64-encoded automatically.", + "hint": "Upload your GCP service account credentials.json file (it will be base64-encoded automatically), or leave blank to use Application Default Credentials.", }, }, "SnowflakeConnectionInfo": { diff --git a/core/wren/tests/unit/test_bigquery_adc.py b/core/wren/tests/unit/test_bigquery_adc.py new file mode 100644 index 000000000..abeb4a9c1 --- /dev/null +++ b/core/wren/tests/unit/test_bigquery_adc.py @@ -0,0 +1,175 @@ +"""BigQuery connector: Application Default Credentials fallback. + +Stubs ``google.cloud.bigquery``, ``google.oauth2.service_account`` and +``google.auth`` for the duration of each test, via ``monkeypatch.setitem`` +so a prior real import of ``google`` (e.g. from ``protobuf``/``grpc``, which +also occupy that namespace package) is overridden rather than left in place +by ``setdefault``, and is restored afterward. ``BigQueryConnector.__init__`` +imports these lazily, so the stub only needs to be live while a test runs, +not at collection time. The real ``google-*`` packages are only installed +under the optional ``bigquery`` extra, which CI's unit-test job does not +install. +""" + +from __future__ import annotations + +import base64 +import sys +import types +from unittest.mock import MagicMock + +import pytest + +pytestmark = pytest.mark.unit + +# --------------------------------------------------------------------------- +# stub google.cloud.bigquery / google.oauth2.service_account / google.auth +# --------------------------------------------------------------------------- + +_client_calls: list[dict] = [] +_service_account_calls: list[dict] = [] +_default_calls: list[dict] = [] + +_service_account_credentials = MagicMock(name="service_account_credentials") +_service_account_credentials.with_scopes.return_value = ( + "scoped-service-account-credentials" +) +_adc_credentials = MagicMock(name="adc_credentials") + + +class _FakeQueryJobConfig: + def __init__(self, dry_run=False, use_query_cache=True): + self.dry_run = dry_run + self.use_query_cache = use_query_cache + self.job_timeout_ms = None + + +class _FakeClient: + def __init__(self, **kwargs): + _client_calls.append(kwargs) + self.default_query_job_config = None + + def close(self): + pass + + +def _fake_from_service_account_info(credits_json): + _service_account_calls.append(credits_json) + return _service_account_credentials + + +def _fake_default(scopes=None): + _default_calls.append({"scopes": scopes}) + return _adc_credentials, "detected-project" + + +_google_mod = types.ModuleType("google") +_google_cloud_mod = types.ModuleType("google.cloud") +_google_cloud_bigquery_mod = types.ModuleType("google.cloud.bigquery") +_google_cloud_bigquery_mod.Client = _FakeClient +_google_cloud_bigquery_mod.QueryJobConfig = _FakeQueryJobConfig +_google_oauth2_mod = types.ModuleType("google.oauth2") +_google_oauth2_service_account_mod = types.ModuleType("google.oauth2.service_account") +_google_oauth2_service_account_mod.Credentials = types.SimpleNamespace( + from_service_account_info=_fake_from_service_account_info +) +_google_auth_mod = types.ModuleType("google.auth") +_google_auth_mod.default = _fake_default + +# `import google.auth` resolves `.auth` as an attribute of "google", so each +# submodule also needs wiring onto its parent, not just onto sys.modules. +_google_mod.cloud = _google_cloud_mod +_google_cloud_mod.bigquery = _google_cloud_bigquery_mod +_google_mod.oauth2 = _google_oauth2_mod +_google_oauth2_mod.service_account = _google_oauth2_service_account_mod +_google_mod.auth = _google_auth_mod + +_GOOGLE_STUBS = { + "google": _google_mod, + "google.cloud": _google_cloud_mod, + "google.cloud.bigquery": _google_cloud_bigquery_mod, + "google.oauth2": _google_oauth2_mod, + "google.oauth2.service_account": _google_oauth2_service_account_mod, + "google.auth": _google_auth_mod, +} + +from wren.connector.bigquery import BigQueryConnector # noqa: E402 +from wren.model import BigQueryDatasetConnectionInfo # noqa: E402 + + +@pytest.fixture(autouse=True) +def _reset_calls(monkeypatch): + for name, module in _GOOGLE_STUBS.items(): + monkeypatch.setitem(sys.modules, name, module) + _client_calls.clear() + _service_account_calls.clear() + _default_calls.clear() + _service_account_credentials.reset_mock(return_value=True, side_effect=True) + _service_account_credentials.with_scopes.return_value = ( + "scoped-service-account-credentials" + ) + yield + + +def test_connection_info_allows_missing_credentials() -> None: + """The model must not require a service-account credentials field. + + On unmodified main, ``credentials`` has no default and this raises a + pydantic ValidationError: a user relying on ADC cannot even construct + the connection info, before a connector is ever built. + """ + info = BigQueryDatasetConnectionInfo( + project_id="my-project", dataset_id="my_dataset" + ) + assert info.credentials is None + + +def test_connector_uses_service_account_when_credentials_present() -> None: + creds_json = base64.b64encode(b'{"type": "service_account"}').decode("ascii") + info = BigQueryDatasetConnectionInfo( + project_id="my-project", dataset_id="my_dataset", credentials=creds_json + ) + BigQueryConnector(info) + assert len(_service_account_calls) == 1 + assert _service_account_calls[0] == {"type": "service_account"} + assert not _default_calls + assert _client_calls[0]["credentials"] == "scoped-service-account-credentials" + assert _client_calls[0]["project"] == "my-project" + + +def test_connector_falls_back_to_adc_when_credentials_absent() -> None: + """Root cause: BigQueryConnector.__init__ never called google.auth.default(), + so omitting credentials (now possible per the model fix above) previously + left ``connection_info.credentials.get_secret_value()`` to crash on + ``None``. It must instead use Application Default Credentials. + """ + info = BigQueryDatasetConnectionInfo( + project_id="my-project", dataset_id="my_dataset" + ) + BigQueryConnector(info) + assert not _service_account_calls + assert len(_default_calls) == 1 + assert _default_calls[0]["scopes"] == [ + "https://www.googleapis.com/auth/drive", + "https://www.googleapis.com/auth/cloud-platform", + ] + assert _client_calls[0]["credentials"] is _adc_credentials + assert _client_calls[0]["project"] == "my-project" + + +def test_connector_falls_back_to_adc_when_credentials_is_empty_string() -> None: + """Root cause: ``is not None`` treated a set-but-empty ``credentials`` + (e.g. an expanded env var like ``${SOME_VAR}`` that resolves to "") as + "use a service account", so the value went straight into base64/JSON + decoding and raised ``json.decoder.JSONDecodeError`` instead of falling + back to ADC. ``SecretStr("")`` is falsy, so the fix is the truthy check + on ``connection_info.credentials`` in ``BigQueryConnector.__init__``. + """ + info = BigQueryDatasetConnectionInfo( + project_id="my-project", dataset_id="my_dataset", credentials="" + ) + BigQueryConnector(info) + assert not _service_account_calls + assert len(_default_calls) == 1 + assert _client_calls[0]["credentials"] is _adc_credentials + assert _client_calls[0]["project"] == "my-project"