feat: add universe domain support for TPC

google/cloud/alloydbconnector/async_connector.py

@@ -16,6 +16,7 @@
 
 import asyncio
 import logging
+import os
 from types import TracebackType
 from typing import TYPE_CHECKING
 from typing import Any
@@ -40,6 +41,10 @@
 
 logger = logging.getLogger(name=__name__)
 
+_DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
+_DEFAULT_ALLOYDB_API_ENDPOINT = "alloydb.googleapis.com"
+_ALLOYDB_HOST_TEMPLATE = "alloydb.{universe_domain}"
+
 
 class AsyncConnector:
     """A class to configure and create connections to Cloud SQL instances
@@ -63,7 +68,10 @@ class AsyncConnector:
             billing purposes.
             Defaults to None, picking up project from environment.
         alloydb_api_endpoint (str): Base URL to use when calling
-            the AlloyDB API endpoint. Defaults to "alloydb.googleapis.com".
+            the AlloyDB API endpoint. Defaults to "alloydb.googleapis.com",
+            this argument should only be used in development.
+        universe_domain (str): The universe domain for AlloyDB API calls.
+            Default: "googleapis.com".
         enable_iam_auth (bool): Enables automatic IAM database authentication.
         ip_type (str | IPTypes): Default IP type for all AlloyDB connections.
             Defaults to IPTypes.PRIVATE ("PRIVATE") for private IP connections.
@@ -84,11 +92,11 @@ def __init__(
         ip_type: str | IPTypes = IPTypes.PRIVATE,
         user_agent: Optional[str] = None,
         refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
+        universe_domain: Optional[str] = None,
     ) -> None:
         self._cache: dict[str, CacheTypes] = {}
         # initialize default params
         self._quota_project = quota_project
-        self._alloydb_api_endpoint = strip_http_prefix(alloydb_api_endpoint)
         self._enable_iam_auth = enable_iam_auth
         # if ip_type is str, convert to IPTypes enum
         if isinstance(ip_type, str):
@@ -99,13 +107,39 @@ def __init__(
             refresh_strategy = RefreshStrategy(refresh_strategy.upper())
         self._refresh_strategy = refresh_strategy
         self._user_agent = user_agent
+        # check for universe domain arg and then env var
+        if universe_domain:
+            self._universe_domain: Optional[str] = universe_domain
+        else:
+            self._universe_domain = os.environ.get("GOOGLE_CLOUD_UNIVERSE_DOMAIN")
+        # construct service endpoint for AlloyDB API calls
+        # if user has not overridden the endpoint, build it from universe domain
+        if alloydb_api_endpoint == _DEFAULT_ALLOYDB_API_ENDPOINT:
+            self._alloydb_api_endpoint = _ALLOYDB_HOST_TEMPLATE.format(
+                universe_domain=self.universe_domain
+            )
+        else:
+            # user explicitly provided a custom endpoint, use it as-is
+            self._alloydb_api_endpoint = strip_http_prefix(alloydb_api_endpoint)
         # initialize credentials for authenticating with AlloyDB Admin API
         scopes = ["https://www.googleapis.com/auth/cloud-platform"]
         if credentials:
             self._credentials = with_scopes_if_required(credentials, scopes=scopes)
         # otherwise use application default credentials
         else:
             self._credentials, _ = google.auth.default(scopes=scopes)
+
+        # 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."
+            )
+
         # initialize credentials for authenticating with the DB
         if db_credentials:
             self._db_credentials = db_credentials
@@ -133,6 +167,10 @@ def __init__(
         self._client: Optional[AlloyDBClient] = None
         self._closed = False
 
+    @property
+    def universe_domain(self) -> str:
+        return self._universe_domain or _DEFAULT_UNIVERSE_DOMAIN
+
     async def connect(
         self,
         instance_uri: str,

google/cloud/alloydbconnector/connector.py

@@ -20,6 +20,7 @@
 from functools import partial
 import io
 import logging
+import os
 import socket
 import struct
 from threading import Thread
@@ -59,6 +60,10 @@
 # the maximum amount of time to wait before aborting a metadata exchange
 IO_TIMEOUT = 30
 
+_DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
+_DEFAULT_ALLOYDB_API_ENDPOINT = "alloydb.googleapis.com"
+_ALLOYDB_HOST_TEMPLATE = "alloydb.{universe_domain}"
+
 
 class Connector:
     """A class to configure and create connections to Cloud SQL instances.
@@ -80,7 +85,10 @@ class Connector:
             billing purposes.
             Defaults to None, picking up project from environment.
         alloydb_api_endpoint (str): Base URL to use when calling
-            the AlloyDB API endpoint. Defaults to "alloydb.googleapis.com".
+            the AlloyDB API endpoint. Defaults to "alloydb.googleapis.com",
+            this argument should only be used in development.
+        universe_domain (str): The universe domain for AlloyDB API calls.
+            Default: "googleapis.com".
         enable_iam_auth (bool): Enables automatic IAM database authentication.
         ip_type (str | IPTypes): Default IP type for all AlloyDB connections.
             Defaults to IPTypes.PRIVATE ("PRIVATE") for private IP connections.
@@ -108,6 +116,7 @@ def __init__(
         user_agent: Optional[str] = None,
         refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
         static_conn_info: Optional[io.TextIOBase] = None,
+        universe_domain: Optional[str] = None,
     ) -> None:
         # create event loop and start it in background thread
         self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
@@ -116,7 +125,6 @@ def __init__(
         self._cache: dict[str, CacheTypes] = {}
         # initialize default params
         self._quota_project = quota_project
-        self._alloydb_api_endpoint = strip_http_prefix(alloydb_api_endpoint)
         self._enable_iam_auth = enable_iam_auth
         # if ip_type is str, convert to IPTypes enum
         if isinstance(ip_type, str):
@@ -127,13 +135,39 @@ def __init__(
             refresh_strategy = RefreshStrategy(refresh_strategy.upper())
         self._refresh_strategy = refresh_strategy
         self._user_agent = user_agent
+        # check for universe domain arg and then env var
+        if universe_domain:
+            self._universe_domain: Optional[str] = universe_domain
+        else:
+            self._universe_domain = os.environ.get("GOOGLE_CLOUD_UNIVERSE_DOMAIN")
+        # construct service endpoint for AlloyDB API calls
+        # if user has not overridden the endpoint, build it from universe domain
+        if alloydb_api_endpoint == _DEFAULT_ALLOYDB_API_ENDPOINT:
+            self._alloydb_api_endpoint = _ALLOYDB_HOST_TEMPLATE.format(
+                universe_domain=self.universe_domain
+            )
+        else:
+            # user explicitly provided a custom endpoint, use it as-is
+            self._alloydb_api_endpoint = strip_http_prefix(alloydb_api_endpoint)
         # initialize credentials for authenticating with AlloyDB Admin API
         scopes = ["https://www.googleapis.com/auth/cloud-platform"]
         if credentials:
             self._credentials = with_scopes_if_required(credentials, scopes=scopes)
         # otherwise use application default credentials
         else:
             self._credentials, _ = default(scopes=scopes)
+
+        # 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."
+            )
+
         # initialize credentials for authenticating with the DB
         if db_credentials:
             self._db_credentials = db_credentials
@@ -152,6 +186,10 @@ def __init__(
         self._static_conn_info = static_conn_info
         self._closed = False
 
+    @property
+    def universe_domain(self) -> str:
+        return self._universe_domain or _DEFAULT_UNIVERSE_DOMAIN
+
     def connect(self, instance_uri: str, driver: str, **kwargs: Any) -> Any:
         """
         Prepares and returns a database DBAPI connection object.

tests/unit/mocks.py

@@ -40,12 +40,14 @@
 from google.cloud import alloydb_v1beta
 import google.cloud.alloydb_connectors_v1.proto.resources_pb2 as connectorspb
 from google.cloud.alloydbconnector.connection_info import ConnectionInfo
+from google.cloud.alloydbconnector.connector import _DEFAULT_UNIVERSE_DOMAIN
 
 
 class FakeCredentials:
     def __init__(self) -> None:
         self.token: Optional[str] = None
         self.expiry: Optional[datetime] = None
+        self._universe_domain = _DEFAULT_UNIVERSE_DOMAIN
 
     def refresh(self, _: Callable) -> None:
         """Refreshes the access token."""
@@ -67,6 +69,11 @@ def valid(self) -> bool:
         """Checks if the credentials are valid."""
         return self.token is not None and not self.expired
 
+    @property
+    def universe_domain(self) -> str:
+        """The universe domain value."""
+        return self._universe_domain
+
     @property
     def token_state(
         self,
@@ -98,6 +105,13 @@ def token_state(
 
 
 class FakeCredentialsRequiresScopes(Scoped):
+    _universe_domain = _DEFAULT_UNIVERSE_DOMAIN
+
+    @property
+    def universe_domain(self) -> str:
+        """The universe domain value."""
+        return self._universe_domain
+
     def requires_scopes(self) -> bool:
         """
         Overrides the requires_scopes() method of the Scoped class to require

tests/unit/test_async_connector.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import os
 from typing import Any
 from typing import Union
 
@@ -491,3 +492,84 @@ async def test_connect_when_closed(credentials: FakeCredentials) -> None:
         exc_info.value.args[0]
         == "Connection attempt failed because the connector has already been closed."
     )
+
+
+def test_async_default_universe_domain() -> None:
+    """Test that default universe domain and constructed service endpoint are
+    formatted correctly.
+    """
+    credentials = FakeCredentials()
+    credentials.token = "test-token"
+    credentials.expiry = None
+    connector = AsyncConnector(credentials=credentials)
+    # 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._alloydb_api_endpoint == "alloydb.googleapis.com"
+
+
+def test_async_configured_universe_domain_matches_GDU() -> None:
+    """Test that configured universe domain succeeds with matched GDU credentials."""
+    credentials = FakeCredentials()
+    credentials.token = "test-token"
+    credentials.expiry = None
+    universe_domain = "googleapis.com"
+    connector = AsyncConnector(credentials=credentials, universe_domain=universe_domain)
+    assert connector._universe_domain == universe_domain
+    assert connector.universe_domain == universe_domain
+    assert connector._alloydb_api_endpoint == f"alloydb.{universe_domain}"
+
+
+def test_async_configured_universe_domain_matches_credentials() -> None:
+    """Test that configured universe domain succeeds with matching universe
+    domain credentials.
+    """
+    universe_domain = "test-universe.test"
+    credentials = FakeCredentials()
+    credentials.token = "test-token"
+    credentials.expiry = None
+    credentials._universe_domain = universe_domain
+    connector = AsyncConnector(credentials=credentials, universe_domain=universe_domain)
+    assert connector._universe_domain == universe_domain
+    assert connector.universe_domain == universe_domain
+    assert connector._alloydb_api_endpoint == f"alloydb.{universe_domain}"
+
+
+def test_async_configured_universe_domain_mismatched_credentials() -> None:
+    """Test that configured universe domain errors with mismatched universe
+    domain credentials.
+    """
+    universe_domain = "test-universe.test"
+    credentials = FakeCredentials()
+    credentials.token = "test-token"
+    credentials.expiry = None
+    with pytest.raises(ValueError) as exc_info:
+        AsyncConnector(credentials=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"({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
+
+
+def test_async_configured_universe_domain_env_var() -> None:
+    """Test that configured universe domain succeeds with universe
+    domain set via GOOGLE_CLOUD_UNIVERSE_DOMAIN env var.
+    """
+    universe_domain = "test-universe.test"
+    credentials = FakeCredentials()
+    credentials.token = "test-token"
+    credentials.expiry = None
+    credentials._universe_domain = universe_domain
+    os.environ["GOOGLE_CLOUD_UNIVERSE_DOMAIN"] = universe_domain
+    try:
+        connector = AsyncConnector(credentials=credentials)
+        assert connector._universe_domain == universe_domain
+        assert connector.universe_domain == universe_domain
+        assert connector._alloydb_api_endpoint == f"alloydb.{universe_domain}"
+    finally:
+        del os.environ["GOOGLE_CLOUD_UNIVERSE_DOMAIN"]