diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py index ada14cbd..5160513f 100755 --- a/google/cloud/sql/connector/connector.py +++ b/google/cloud/sql/connector/connector.py @@ -19,6 +19,7 @@ import asyncio from functools import partial import logging +import os from threading import Thread from types import TracebackType from typing import Any, Optional, Union @@ -171,7 +172,11 @@ def __init__( if isinstance(ip_type, str): ip_type = IPTypes._from_str(ip_type) self._ip_type = ip_type - self._universe_domain = universe_domain + # check for universe domain arg and then env var + if universe_domain: + self._universe_domain = universe_domain + else: + self._universe_domain = os.environ.get("GOOGLE_CLOUD_UNIVERSE_DOMAIN") # type: ignore # construct service endpoint for Cloud SQL Admin API calls if not sqladmin_api_endpoint: self._sqladmin_api_endpoint = ( diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py index d4f53ed5..eaf188c3 100644 --- a/tests/unit/test_connector.py +++ b/tests/unit/test_connector.py @@ -15,6 +15,7 @@ """ import asyncio +import os from typing import Union from aiohttp import ClientResponseError @@ -428,3 +429,25 @@ def test_configured_universe_domain_mismatched_credentials( "is the default." ) assert exc_info.value.args[0] == err_msg + + +def test_configured_universe_domain_env_var( + fake_credentials: Credentials, +) -> None: + """Test that configured universe domain succeeds with universe + domain set via GOOGLE_CLOUD_UNIVERSE_DOMAIN env var. + """ + universe_domain = "test-universe.test" + # set fake credentials to be configured for the universe domain + fake_credentials._universe_domain = universe_domain + # set environment variable + os.environ["GOOGLE_CLOUD_UNIVERSE_DOMAIN"] = universe_domain + # Note: we are not passing universe_domain arg, env var should set it + with Connector(credentials=fake_credentials) as connector: + # test universe domain was configured + assert connector._universe_domain == universe_domain + # test property and service endpoint construction + assert connector.universe_domain == universe_domain + assert connector._sqladmin_api_endpoint == f"https://sqladmin.{universe_domain}" + # unset env var + del os.environ["GOOGLE_CLOUD_UNIVERSE_DOMAIN"]