feat: add universe domain support to Connector (TPC) (#1045)

Author: jackwotherspoon

google/cloud/sql/connector/connector.py

@@ -43,6 +43,9 @@
 logger = logging.getLogger(name=__name__)
 
 ASYNC_DRIVERS = ["asyncpg"]
+_DEFAULT_SCHEME = "https://"
+_DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
+_SQLADMIN_HOST_TEMPLATE = "sqladmin.{universe_domain}"
 
 
 class Connector:
@@ -58,6 +61,7 @@ def __init__(
  quota_project: Optional[str] = None,
  sqladmin_api_endpoint: Optional[str] = None,
  user_agent: Optional[str] = None,
+ universe_domain: Optional[str] = None,
  ) -> None:
  """Initializes a Connector instance.
 
@@ -90,6 +94,10 @@ def __init__(
  sqladmin_api_endpoint (str): Base URL to use when calling the Cloud SQL
  Admin API endpoint. Defaults to "https://sqladmin.googleapis.com",
  this argument should only be used in development.
+
+ universe_domain (str): The universe domain for Cloud SQL API calls.
+ Default: "googleapis.com".
+
  """
  # if event loop is given, use for background tasks
  if loop:
@@ -126,12 +134,36 @@ def __init__(
  self._timeout = timeout
  self._enable_iam_auth = enable_iam_auth
  self._quota_project = quota_project
- self._sqladmin_api_endpoint = sqladmin_api_endpoint
  self._user_agent = user_agent
  # if ip_type is str, convert to IPTypes enum
  if isinstance(ip_type, str):
  ip_type = IPTypes._from_str(ip_type)
  self._ip_type = ip_type
+ self._universe_domain = universe_domain
+ # construct service endpoint for Cloud SQL Admin API calls
+ if not sqladmin_api_endpoint:
+ self._sqladmin_api_endpoint = (
+ _DEFAULT_SCHEME
+ + _SQLADMIN_HOST_TEMPLATE.format(universe_domain=self.universe_domain)
+ )
+ # otherwise if endpoint override is passed in use it
+ else:
+ self._sqladmin_api_endpoint = sqladmin_api_endpoint
+
+ # validate that the universe domain of the credentials matches the
+ # universe domain of the service endpoint
+ if self._credentials.universe_domain != self.universe_domain:
+ raise ValueError(
+ f"The configured universe domain ({self.universe_domain}) does "
+ "not match the universe domain found in the credentials "
+ f"({self._credentials.universe_domain}). If you haven't "
+ "configured the universe domain explicitly, `googleapis.com` "
+ "is the default."
+ )
+
+ @property
+ def universe_domain(self) -> str:
+ return self._universe_domain or _DEFAULT_UNIVERSE_DOMAIN
 
  def connect(
  self, instance_connection_string: str, driver: str, **kwargs: Any
@@ -371,6 +403,7 @@ async def create_async_connector(
  quota_project: Optional[str] = None,
  sqladmin_api_endpoint: Optional[str] = None,
  user_agent: Optional[str] = None,
+ universe_domain: Optional[str] = None,
 ) -> Connector:
  """Helper function to create Connector object for asyncio connections.
 

tests/unit/mocks.py

@@ -30,6 +30,7 @@
 from cryptography.x509.oid import NameOID
 from google.auth.credentials import Credentials
 
+from google.cloud.sql.connector.connector import _DEFAULT_UNIVERSE_DOMAIN
 from google.cloud.sql.connector.instance import ConnectionInfo
 from google.cloud.sql.connector.utils import generate_keys
 from google.cloud.sql.connector.utils import write_to_file
@@ -41,6 +42,7 @@ def __init__(
  ) -> None:
  self.token = token
  self.expiry = expiry
+ self._universe_domain = _DEFAULT_UNIVERSE_DOMAIN
 
  @property
  def __class__(self) -> Credentials:
@@ -68,6 +70,11 @@ def expired(self) -> bool:
  return False
  return True
 
+ @property
+ def universe_domain(self) -> str:
+ """The universe domain value."""
+ return self._universe_domain
+
  @property
  def valid(self) -> bool:
  """Checks the validity of the credentials.

tests/unit/test_client.py

@@ -77,21 +77,6 @@ async def test_CloudSQLClient_init_(fake_credentials: FakeCredentials) -> None:
  await client.close()
 
 
-async def test_CloudSQLClient_init_default_service_endpoint(
- fake_credentials: FakeCredentials,
-) -> None:
- """
- Test to check whether the __init__ method of CloudSQLClient
- can correctly initialize the default service endpoint.
- """
- driver = "pg8000"
- client = CloudSQLClient(None, "my-quota-project", fake_credentials, driver=driver)
- # verify base endpoint is set to proper default
- assert client._sqladmin_api_endpoint == "https://sqladmin.googleapis.com"
- # close client
- await client.close()
-
-
 @pytest.mark.asyncio
 async def test_CloudSQLClient_init_custom_user_agent(
  fake_credentials: FakeCredentials,

tests/unit/test_connector.py

@@ -273,3 +273,67 @@ def test_Connector_close_called_multiple_times(fake_credentials: Credentials) ->
  assert connector._thread.is_alive() is False
  # call connector.close a second time
  connector.close()
+
+
+def test_default_universe_domain(fake_credentials: Credentials) -> None:
+ """Test that default universe domain and constructed service endpoint are
+ formatted correctly.
+ """
+ with Connector(credentials=fake_credentials) as connector:
+ # test universe domain was not configured
+ assert connector._universe_domain is None
+ # test property and service endpoint construction
+ assert connector.universe_domain == "googleapis.com"
+ assert connector._sqladmin_api_endpoint == "https://sqladmin.googleapis.com"
+
+
+def test_configured_universe_domain_matches_GDU(fake_credentials: Credentials) -> None:
+ """Test that configured universe domain succeeds with matched GDU credentials."""
+ universe_domain = "googleapis.com"
+ with Connector(
+ credentials=fake_credentials, universe_domain=universe_domain
+ ) 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}"
+
+
+def test_configured_universe_domain_matches_credentials(
+ fake_credentials: Credentials,
+) -> None:
+ """Test that configured universe domain succeeds with matching universe
+ domain credentials.
+ """
+ universe_domain = "test-universe.test"
+ # set fake credentials to be configured for the universe domain
+ fake_credentials._universe_domain = universe_domain
+ with Connector(
+ credentials=fake_credentials, universe_domain=universe_domain
+ ) 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}"
+
+
+def test_configured_universe_domain_mismatched_credentials(
+ fake_credentials: Credentials,
+) -> None:
+ """Test that configured universe domain errors with mismatched universe
+ domain credentials.
+ """
+ universe_domain = "test-universe.test"
+ # credentials have GDU domain ("googleapis.com")
+ with pytest.raises(ValueError) as exc_info:
+ Connector(credentials=fake_credentials, universe_domain=universe_domain)
+ err_msg = (
+ f"The configured universe domain ({universe_domain}) does "
+ "not match the universe domain found in the credentials "
+ f"({fake_credentials.universe_domain}). If you haven't "
+ "configured the universe domain explicitly, `googleapis.com` "
+ "is the default."
+ )
+ assert exc_info.value.args[0] == err_msg