feat: Add PSC DNS and Global Write Endpoint support to Python Connector

.gitignore

@@ -10,3 +10,7 @@ dist/
 sponge_log.xml
 .envrc
 *.iml
+.mypy_cache/
+.nox/
+.pytest_cache/
+.ruff_cache/

google/cloud/sql/connector/__init__.py

@@ -23,11 +23,11 @@
 from google.cloud.sql.connector.version import __version__
 
 __all__ = [
-    "__version__",
-    "create_async_connector",
     "Connector",
     "DefaultResolver",
     "DnsResolver",
     "IPTypes",
     "RefreshStrategy",
+    "__version__",
+    "create_async_connector",
 ]

google/cloud/sql/connector/client.py

@@ -17,7 +17,7 @@
 import asyncio
 import datetime
 import logging
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 import aiohttp
 from cryptography.hazmat.backends import default_backend
@@ -42,7 +42,7 @@
 logger = logging.getLogger(name=__name__)
 
 
-def _format_user_agent(driver: Optional[str], custom: Optional[str]) -> str:
+def _format_user_agent(driver: str | None, custom: str | None) -> str:
     agent = f"{USER_AGENT}+{driver}" if driver else USER_AGENT
     if custom and isinstance(custom, str):
         agent = f"{agent} {custom}"
@@ -52,12 +52,12 @@ def _format_user_agent(driver: Optional[str], custom: Optional[str]) -> str:
 class CloudSQLClient:
     def __init__(
         self,
-        sqladmin_api_endpoint: Optional[str],
-        quota_project: Optional[str],
+        sqladmin_api_endpoint: str | None,
+        quota_project: str | None,
         credentials: Credentials,
-        client: Optional[aiohttp.ClientSession] = None,
-        driver: Optional[str] = None,
-        user_agent: Optional[str] = None,
+        client: aiohttp.ClientSession | None = None,
+        driver: str | None = None,
+        user_agent: str | None = None,
     ) -> None:
         """Establishes the client to be used for Cloud SQL Admin API requests.
 
@@ -135,7 +135,7 @@ async def _get_metadata(
                 if message:
                     resp.reason = message
         # skip, raise_for_status will catch all errors in finally block
-        except Exception:
+        except Exception:  # noqa: BLE001, S110
             pass
         finally:
             resp.raise_for_status()
@@ -146,7 +146,7 @@ async def _get_metadata(
             )
 
         ip_addresses = (
-            {ip["type"]: ip["ipAddress"] for ip in ret_dict["ipAddresses"]}
+            {ip["type"]: [ip["ipAddress"]] for ip in ret_dict["ipAddresses"]}
             if "ipAddresses" in ret_dict
             else {}
         )
@@ -156,27 +156,70 @@ async def _get_metadata(
         if ret_dict.get("pscEnabled"):
             # Find PSC instance DNS name in the dns_names field
             psc_dns_names = [
-                d["name"]
+                d["name"].rstrip(".")
                 for d in ret_dict.get("dnsNames", [])
                 if d["connectionType"] == "PRIVATE_SERVICE_CONNECT"
                 and d["dnsScope"] == "INSTANCE"
             ]
-            dns_name = psc_dns_names[0] if psc_dns_names else None
+            # Sort: .sql-psc.goog first
+            psc_dns_names.sort(key=lambda x: x.endswith(".sql-psc.goog"), reverse=True)
 
             # Fall back do dns_name field if dns_names is not set
-            if dns_name is None:
+            if not psc_dns_names:
                 dns_name = ret_dict.get("dnsName", None)
+                if dns_name:
+                    psc_dns_names = [dns_name.rstrip(".")]
 
-            # Remove trailing period from DNS name. Required for SSL in Python
-            if dns_name:
-                ip_addresses["PSC"] = dns_name.rstrip(".")
+            if psc_dns_names:
+                ip_addresses["PSC"] = psc_dns_names
 
         return {
             "ip_addresses": ip_addresses,
             "server_ca_cert": ret_dict["serverCaCert"]["cert"],
             "database_version": ret_dict["databaseVersion"],
         }
 
+    async def resolve_connect_settings(
+        self,
+        dns_name: str,
+        location: str,
+    ) -> dict[str, Any]:
+        """Asynchronously calls the resolveConnectSettings endpoint to resolve a
+        PSC DNS name to a connection name.
+
+        Args:
+            dns_name (str): The DNS name of the Cloud SQL instance.
+            location (str): The region/location of the instance.
+
+        Returns:
+            A dictionary containing the resolve response (e.g. connectionName).
+        """
+        # before making Cloud SQL Admin API calls, refresh creds if required
+        if self._credentials.token_state != TokenState.FRESH:
+            self._credentials.refresh(requests.Request())
+
+        headers = {
+            "Authorization": f"Bearer {self._credentials.token}",
+        }
+
+        url = f"{self._sqladmin_api_endpoint}/sql/{API_VERSION}/locations/{location}/dns/{dns_name}:resolveConnectSettings"
+
+        resp = await self._client.get(url, headers=headers)
+        if resp.status >= 500:
+            resp = await retry_50x(self._client.get, url, headers=headers)
+        try:
+            ret_dict = await resp.json()
+            if resp.status >= 400:
+                message = ret_dict.get("error", {}).get("message")
+                if message:
+                    resp.reason = message
+        except Exception:  # noqa: BLE001, S110
+            pass
+        finally:
+            resp.raise_for_status()
+
+        return ret_dict
+
     async def _get_ephemeral(
         self,
         project: str,
@@ -223,7 +266,7 @@ async def _get_ephemeral(
                 if message:
                     resp.reason = message
         # skip, raise_for_status will catch all errors in finally block
-        except Exception:
+        except Exception:  # noqa: BLE001, S110
             pass
         finally:
             resp.raise_for_status()
@@ -244,8 +287,7 @@ async def _get_ephemeral(
             # Ref: https://github.com/googleapis/google-auth-library-python/blob/49a5ff7411a2ae4d32a7d11700f9f961c55406a9/google/auth/_helpers.py#L93-L99
             token_expiration = token_expiration.replace(tzinfo=datetime.timezone.utc)
 
-            if expiration > token_expiration:
-                expiration = token_expiration
+            expiration = min(expiration, token_expiration)
         return ephemeral_cert, expiration
 
     async def get_connection_info(
@@ -274,7 +316,7 @@ async def get_connection_info(
         """
         priv_key, pub_key = await keys
         # before making Cloud SQL Admin API calls, refresh creds if required
-        if not self._credentials.token_state == TokenState.FRESH:
+        if self._credentials.token_state != TokenState.FRESH:
             self._credentials.refresh(requests.Request())
 
         metadata_task = asyncio.create_task(

google/cloud/sql/connector/connection_info.py

@@ -18,7 +18,7 @@
 from dataclasses import dataclass
 import logging
 import ssl
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 from google.cloud.sql.connector.connection_name import ConnectionName
 from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError
@@ -67,7 +67,7 @@ class ConnectionInfo:
     ip_addrs: dict[str, Any]
     database_version: str
     expiration: datetime.datetime
-    context: Optional[ssl.SSLContext] = None
+    context: ssl.SSLContext | None = None
 
     async def create_ssl_context(self, enable_iam_auth: bool = False) -> ssl.SSLContext:
         """Constructs a SSL/TLS context for the given connection info.
@@ -121,6 +121,17 @@ def get_preferred_ip(self, ip_type: IPTypes) -> str:
         """Returns the first IP address for the instance, according to the preference
         supplied by ip_type. If no IP addressess with the given preference are found,
         an error is raised."""
+        if ip_type.value in self.ip_addrs:
+            return self.ip_addrs[ip_type.value][0]
+        raise CloudSQLIPTypeError(
+            "Cloud SQL instance does not have any IP addresses matching "
+            f"preference: {ip_type.value}"
+        )
+
+    def get_preferred_ips(self, ip_type: IPTypes) -> list[str]:
+        """Returns all IP addresses for the instance, according to the preference
+        supplied by ip_type. If no IP addressess with the given preference are found,
+        an error is raised."""
         if ip_type.value in self.ip_addrs:
             return self.ip_addrs[ip_type.value]
         raise CloudSQLIPTypeError(

google/cloud/sql/connector/connection_name.py

@@ -18,7 +18,7 @@
 # Instance connection name is the format <PROJECT>:<REGION>:<INSTANCE_NAME>
 # Additionally, we have to support legacy "domain-scoped" projects
 # (e.g. "google.com:PROJECT")
-CONN_NAME_REGEX = re.compile(("([^:]+(:[^:]+)?):([^:]+):([^:]+)"))
+CONN_NAME_REGEX = re.compile("([^:]+(:[^:]+)?):([^:]+):([^:]+)")
 # The domain name pattern in accordance with RFC 1035, RFC 1123 and RFC 2181.
 DOMAIN_NAME_REGEX = re.compile(
     r"^(?:[_a-z0-9](?:[_a-z0-9-]{0,61}[a-z0-9])?\.)+(?:[a-z](?:[a-z0-9-]{0,61}[a-z0-9])?)?$"
@@ -48,9 +48,7 @@ def get_connection_string(self) -> str:
 
 
 def _is_valid_domain(domain_name: str) -> bool:
-    if DOMAIN_NAME_REGEX.fullmatch(domain_name) is None:
-        return False
-    return True
+    return DOMAIN_NAME_REGEX.fullmatch(domain_name) is not None
 
 
 def _parse_connection_name(connection_name: str) -> ConnectionName: