diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py index 6faa32606..b59bc86d3 100644 --- a/google/cloud/bigquery/_helpers.py +++ b/google/cloud/bigquery/_helpers.py @@ -19,6 +19,7 @@ import decimal import math import re +import os from typing import Optional, Union from dateutil import relativedelta @@ -28,8 +29,8 @@ from google.cloud._helpers import _RFC3339_MICROS from google.cloud._helpers import _RFC3339_NO_FRACTION from google.cloud._helpers import _to_bytes -import packaging.version +import packaging.version _RFC3339_MICROS_NO_ZULU = "%Y-%m-%dT%H:%M:%S.%f" _TIMEONLY_WO_MICROS = "%H:%M:%S" @@ -51,6 +52,16 @@ _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = packaging.version.Version("2.6.0") +BIGQUERY_EMULATOR_HOST = "BIGQUERY_EMULATOR_HOST" +"""Environment variable defining host for emulator.""" + +_DEFAULT_HOST = "https://bigquery.googleapis.com" +"""Default host for JSON API.""" + + +def _get_bigquery_host(): + return os.environ.get(BIGQUERY_EMULATOR_HOST, _DEFAULT_HOST) + class BQStorageVersions: """Version comparisons for google-cloud-bigqueyr-storage package.""" diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index b388f1d4c..fb772ea11 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -56,7 +56,6 @@ import google.cloud._helpers # type: ignore from google.cloud import exceptions # pytype: disable=import-error from google.cloud.client import ClientWithProject # type: ignore # pytype: disable=import-error - from google.cloud.bigquery_storage_v1.services.big_query_read.client import ( DEFAULT_CLIENT_INFO as DEFAULT_BQSTORAGE_CLIENT_INFO, ) @@ -67,6 +66,8 @@ from google.cloud.bigquery._helpers import _record_field_to_json from google.cloud.bigquery._helpers import _str_or_none from google.cloud.bigquery._helpers import _verify_job_config_type +from google.cloud.bigquery._helpers import _get_bigquery_host +from google.cloud.bigquery._helpers import _DEFAULT_HOST from google.cloud.bigquery._http import Connection from google.cloud.bigquery import _pandas_helpers from google.cloud.bigquery.dataset import Dataset @@ -230,6 +231,8 @@ def __init__( ) kw_args = {"client_info": client_info} + bq_host = _get_bigquery_host() + kw_args["api_endpoint"] = bq_host if bq_host != _DEFAULT_HOST else None if client_options: if type(client_options) == dict: client_options = google.api_core.client_options.from_dict( diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py index 885e773d3..2e714c707 100644 --- a/tests/unit/test__helpers.py +++ b/tests/unit/test__helpers.py @@ -1288,3 +1288,29 @@ def test_decimal_as_float_api_repr(): "parameterValue": {"value": 42.0}, "name": "x", } + + +class Test__get_bigquery_host(unittest.TestCase): + @staticmethod + def _call_fut(): + from google.cloud.bigquery._helpers import _get_bigquery_host + + return _get_bigquery_host() + + def test_wo_env_var(self): + from google.cloud.bigquery._helpers import _DEFAULT_HOST + + with mock.patch("os.environ", {}): + host = self._call_fut() + + self.assertEqual(host, _DEFAULT_HOST) + + def test_w_env_var(self): + from google.cloud.bigquery._helpers import BIGQUERY_EMULATOR_HOST + + HOST = "https://api.example.com" + + with mock.patch("os.environ", {BIGQUERY_EMULATOR_HOST: HOST}): + host = self._call_fut() + + self.assertEqual(host, HOST)