diff --git a/.gitignore b/.gitignore index 06827b584..6fed8e218 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,7 @@ dist/ sponge_log.xml .envrc *.iml +.mypy_cache/ +.nox/ +.pytest_cache/ +.ruff_cache/ diff --git a/google/cloud/sql/connector/__init__.py b/google/cloud/sql/connector/__init__.py index 6913337d3..81462847b 100644 --- a/google/cloud/sql/connector/__init__.py +++ b/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", ] diff --git a/google/cloud/sql/connector/client.py b/google/cloud/sql/connector/client.py index 11508ce17..ebf823e2d 100644 --- a/google/cloud/sql/connector/client.py +++ b/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,20 +156,22 @@ 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, @@ -177,6 +179,47 @@ async def _get_metadata( "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( diff --git a/google/cloud/sql/connector/connection_info.py b/google/cloud/sql/connector/connection_info.py index 1bd26a4e2..49404f491 100644 --- a/google/cloud/sql/connector/connection_info.py +++ b/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( diff --git a/google/cloud/sql/connector/connection_name.py b/google/cloud/sql/connector/connection_name.py index ad5dc40fb..ee980cdcf 100644 --- a/google/cloud/sql/connector/connection_name.py +++ b/google/cloud/sql/connector/connection_name.py @@ -18,7 +18,7 @@ # Instance connection name is the format :: # 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: diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py old mode 100755 new mode 100644 index 798969c2c..3a1df0ea0 --- a/google/cloud/sql/connector/connector.py +++ b/google/cloud/sql/connector/connector.py @@ -23,13 +23,16 @@ import socket from threading import Thread from types import TracebackType -from typing import Any, Callable, Optional, Union +from typing import Any, Callable import google.auth from google.auth.credentials import Credentials from google.auth.credentials import with_scopes_if_required -import google.cloud.sql.connector.asyncpg as asyncpg +from google.cloud.sql.connector import asyncpg +from google.cloud.sql.connector import pg8000 +from google.cloud.sql.connector import pymysql +from google.cloud.sql.connector import pytds from google.cloud.sql.connector.client import CloudSQLClient from google.cloud.sql.connector.enums import DriverMapping from google.cloud.sql.connector.enums import IPTypes @@ -39,9 +42,6 @@ from google.cloud.sql.connector.instance import RefreshAheadCache from google.cloud.sql.connector.lazy import LazyRefreshCache from google.cloud.sql.connector.monitored_cache import MonitoredCache -import google.cloud.sql.connector.pg8000 as pg8000 -import google.cloud.sql.connector.pymysql as pymysql -import google.cloud.sql.connector.pytds as pytds from google.cloud.sql.connector.resolver import DefaultResolver from google.cloud.sql.connector.resolver import DnsResolver from google.cloud.sql.connector.utils import format_database_user @@ -64,14 +64,14 @@ def __init__( ip_type: str | IPTypes = IPTypes.PUBLIC, enable_iam_auth: bool = False, timeout: int = 30, - credentials: Optional[Credentials] = None, - loop: Optional[asyncio.AbstractEventLoop] = None, - quota_project: Optional[str] = None, - sqladmin_api_endpoint: Optional[str] = None, - user_agent: Optional[str] = None, - universe_domain: Optional[str] = None, + credentials: Credentials | None = None, + loop: asyncio.AbstractEventLoop | None = None, + quota_project: str | None = None, + sqladmin_api_endpoint: str | None = None, + user_agent: str | None = None, + universe_domain: str | None = None, refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND, - resolver: type[DefaultResolver] | type[DnsResolver] = DefaultResolver, + resolver: type[DefaultResolver | DnsResolver] = DefaultResolver, failover_period: int = 30, ) -> None: """Initializes a Connector instance. @@ -133,10 +133,10 @@ def __init__( # if event loop is given, use for background tasks if loop: self._loop: asyncio.AbstractEventLoop = loop - self._thread: Optional[Thread] = None + self._thread: Thread | None = None # if lazy refresh is specified we should lazy init keys if self._refresh_strategy == RefreshStrategy.LAZY: - self._keys: Optional[asyncio.Future] = None + self._keys: asyncio.Future | None = None else: self._keys = loop.create_task(generate_keys()) # if no event loop is given, spin up new loop in background thread @@ -155,7 +155,7 @@ def __init__( # initialize dict to store caches, key is a tuple consisting of instance # connection name string and enable_iam_auth boolean flag self._cache: dict[tuple[str, bool], MonitoredCache] = {} - self._client: Optional[CloudSQLClient] = None + self._client: CloudSQLClient | None = None self._closed: bool = False # initialize credentials @@ -176,7 +176,8 @@ def __init__( self._timeout = timeout self._enable_iam_auth = enable_iam_auth self._user_agent = user_agent - self._resolver = resolver() + self._resolver_cls = resolver + self._resolver: DefaultResolver | DnsResolver | None = None self._failover_period = failover_period # if ip_type is str, convert to IPTypes enum if isinstance(ip_type, str): @@ -316,6 +317,8 @@ async def connect_async( user_agent=self._user_agent, driver=driver, ) + if self._resolver is None: + self._resolver = self._resolver_cls(client=self._client) enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth) conn_name = await self._resolver.resolve(instance_connection_string) @@ -329,7 +332,7 @@ async def connect_async( logger.debug( f"['{conn_name}']: Refresh strategy is set to lazy refresh" ) - cache: Union[LazyRefreshCache, RefreshAheadCache] = LazyRefreshCache( + cache: LazyRefreshCache | RefreshAheadCache = LazyRefreshCache( conn_name, self._client, self._keys, @@ -384,13 +387,14 @@ async def connect_async( conn_info = await monitored_cache.connect_info() # validate driver matches intended database engine DriverMapping.validate_engine(driver, conn_info.database_version) - ip_address = conn_info.get_preferred_ip(ip_type) + preferred_ips = conn_info.get_preferred_ips(ip_type) except Exception: # with an error from Cloud SQL Admin API call or IP type, invalidate # the cache and re-raise the error await self._remove_cached(str(conn_name), enable_iam_auth) raise + targets = [] # If the connector is configured with a custom DNS name, attempt to use # that DNS name to connect to the instance. Fall back to the metadata IP # address if the DNS name does not resolve to an IP address. @@ -398,26 +402,29 @@ async def connect_async( try: ips = await self._resolver.resolve_a_record(conn_info.conn_name.domain_name) if ips: - ip_address = ips[0] + targets.extend(ips) logger.debug( f"['{instance_connection_string}']: Custom DNS name " - f"'{conn_info.conn_name.domain_name}' resolved to '{ip_address}', " + f"'{conn_info.conn_name.domain_name}' resolved to '{ips}', " "using it to connect" ) else: logger.debug( f"['{instance_connection_string}']: Custom DNS name " f"'{conn_info.conn_name.domain_name}' resolved but returned no " - f"entries, using '{ip_address}' from instance metadata" + f"entries, using '{preferred_ips}' from instance metadata" ) - except Exception as e: + targets.extend(preferred_ips) + except Exception as e: # noqa: BLE001 logger.debug( f"['{instance_connection_string}']: Custom DNS name " f"'{conn_info.conn_name.domain_name}' did not resolve to an IP " - f"address: {e}, using '{ip_address}' from instance metadata" + f"address: {e}, using '{preferred_ips}' from instance metadata" ) + targets.extend(preferred_ips) + else: + targets.extend(preferred_ips) - logger.debug(f"['{conn_info.conn_name}']: Connecting to {ip_address}:3307") # format `user` param for automatic IAM database authn if enable_iam_auth: formatted_user = format_database_user( @@ -428,32 +435,56 @@ async def connect_async( f"['{instance_connection_string}']: Truncated IAM database username from {kwargs['user']} to {formatted_user}" ) kwargs["user"] = formatted_user + try: - # async drivers are unblocking and can be awaited directly - if driver in ASYNC_DRIVERS: - return await connector( - ip_address, - await conn_info.create_ssl_context(enable_iam_auth), - **kwargs, - ) - # Create socket with SSLContext for sync drivers - ctx = await conn_info.create_ssl_context(enable_iam_auth) - sock = ctx.wrap_socket( - socket.create_connection((ip_address, SERVER_PROXY_PORT)), - server_hostname=ip_address, - ) - # If this connection was opened using a domain name, then store it - # for later in case we need to forcibly close it on failover. - if conn_info.conn_name.domain_name: - monitored_cache.sockets.append(sock) - # Synchronous drivers are blocking and run using executor - connect_partial = partial( - connector, - ip_address, - sock, - **kwargs, - ) - return await self._loop.run_in_executor(None, connect_partial) + last_ex = None + for target_ip in targets: + logger.debug(f"['{conn_info.conn_name}']: Connecting to {target_ip}:3307") + try: + # async drivers are unblocking and can be awaited directly + if driver in ASYNC_DRIVERS: + conn = await connector( + target_ip, + await conn_info.create_ssl_context(enable_iam_auth), + **kwargs, + ) + last_ex = None + return conn + + # Create socket with SSLContext for sync drivers + ctx = await conn_info.create_ssl_context(enable_iam_auth) + raw_sock = socket.create_connection((target_ip, SERVER_PROXY_PORT)) + try: + sock = ctx.wrap_socket( + raw_sock, + server_hostname=target_ip, + ) + except Exception: + raw_sock.close() + raise + + # If this connection was opened using a domain name, then store it + # for later in case we need to forcibly close it on failover. + if conn_info.conn_name.domain_name: + monitored_cache.sockets.append(sock) + # Synchronous drivers are blocking and run using executor + connect_partial = partial( + connector, + target_ip, + sock, + **kwargs, + ) + conn = await self._loop.run_in_executor(None, connect_partial) + last_ex = None + return conn + except Exception as e: # noqa: BLE001 + logger.debug( + f"['{conn_info.conn_name}']: Connection to {target_ip} failed: {e}" + ) + last_ex = e + + if last_ex: + raise last_ex except Exception: # with any exception, we attempt a force refresh, then throw the error @@ -479,9 +510,9 @@ def __enter__(self) -> Any: def __exit__( self, - exc_type: Optional[type[BaseException]], - exc_val: Optional[BaseException], - exc_tb: Optional[TracebackType], + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, ) -> None: """Exit context manager by closing Connector""" self.close() @@ -492,9 +523,9 @@ async def __aenter__(self) -> Any: async def __aexit__( self, - exc_type: Optional[type[BaseException]], - exc_val: Optional[BaseException], - exc_tb: Optional[TracebackType], + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, ) -> None: """Exit async context manager by closing Connector""" await self.close_async() @@ -528,14 +559,14 @@ async def create_async_connector( ip_type: str | IPTypes = IPTypes.PUBLIC, enable_iam_auth: bool = False, timeout: int = 30, - credentials: Optional[Credentials] = None, - loop: Optional[asyncio.AbstractEventLoop] = None, - quota_project: Optional[str] = None, - sqladmin_api_endpoint: Optional[str] = None, - user_agent: Optional[str] = None, - universe_domain: Optional[str] = None, + credentials: Credentials | None = None, + loop: asyncio.AbstractEventLoop | None = None, + quota_project: str | None = None, + sqladmin_api_endpoint: str | None = None, + user_agent: str | None = None, + universe_domain: str | None = None, refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND, - resolver: type[DefaultResolver] | type[DnsResolver] = DefaultResolver, + resolver: type[DefaultResolver | DnsResolver] = DefaultResolver, failover_period: int = 30, ) -> Connector: """Helper function to create Connector object for asyncio connections. diff --git a/google/cloud/sql/connector/enums.py b/google/cloud/sql/connector/enums.py index e6b56af0e..88b5bf47e 100644 --- a/google/cloud/sql/connector/enums.py +++ b/google/cloud/sql/connector/enums.py @@ -61,7 +61,7 @@ class DriverMapping(Enum): """Maps a given database driver to it's corresponding database engine.""" ASYNCPG = "POSTGRES" - PG8000 = "POSTGRES" + PG8000 = "POSTGRES" # noqa: PIE796 PYMYSQL = "MYSQL" PYTDS = "SQLSERVER" diff --git a/google/cloud/sql/connector/exceptions.py b/google/cloud/sql/connector/exceptions.py index 1f15ced47..7399f7109 100644 --- a/google/cloud/sql/connector/exceptions.py +++ b/google/cloud/sql/connector/exceptions.py @@ -21,7 +21,6 @@ class ConnectorLoopError(Exception): in an invalid state (event loop in current thread). """ - pass class TLSVersionError(Exception): @@ -29,7 +28,6 @@ class TLSVersionError(Exception): Raised when the required TLS protocol version is not supported. """ - pass class CloudSQLIPTypeError(Exception): @@ -37,7 +35,6 @@ class CloudSQLIPTypeError(Exception): Raised when IP address for the preferred IP type is not found. """ - pass class PlatformNotSupportedError(Exception): @@ -45,7 +42,6 @@ class PlatformNotSupportedError(Exception): Raised when a feature is not supported on the current platform. """ - pass class AutoIAMAuthNotSupported(Exception): @@ -54,7 +50,6 @@ class AutoIAMAuthNotSupported(Exception): supported with database engine version. """ - pass class RefreshNotValidError(Exception): @@ -62,7 +57,6 @@ class RefreshNotValidError(Exception): Exception to be raised when the task returned from refresh is not valid. """ - pass class IncompatibleDriverError(Exception): diff --git a/google/cloud/sql/connector/instance.py b/google/cloud/sql/connector/instance.py index fb8711309..28ab54e40 100644 --- a/google/cloud/sql/connector/instance.py +++ b/google/cloud/sql/connector/instance.py @@ -130,7 +130,7 @@ async def _perform_refresh(self) -> ConnectionInfo: except Exception as e: logger.debug( f"['{self._conn_name}']: Connection info " - f"refresh operation failed: {str(e)}" + f"refresh operation failed: {e!s}" ) raise diff --git a/google/cloud/sql/connector/lazy.py b/google/cloud/sql/connector/lazy.py index c75d07e52..4e3de018a 100644 --- a/google/cloud/sql/connector/lazy.py +++ b/google/cloud/sql/connector/lazy.py @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import asyncio from datetime import datetime from datetime import timedelta from datetime import timezone import logging -from typing import Optional from google.cloud.sql.connector.client import CloudSQLClient from google.cloud.sql.connector.connection_info import ConnectionInfo @@ -61,7 +62,7 @@ def __init__( self._keys = keys self._client = client self._lock = asyncio.Lock() - self._cached: Optional[ConnectionInfo] = None + self._cached: ConnectionInfo | None = None self._needs_refresh = False self._closed = False @@ -112,7 +113,7 @@ async def connect_info(self) -> ConnectionInfo: except Exception as e: logger.debug( f"['{self._conn_name}']: Connection info " - f"refresh operation failed: {str(e)}" + f"refresh operation failed: {e!s}" ) raise logger.debug( @@ -121,7 +122,7 @@ async def connect_info(self) -> ConnectionInfo: ) logger.debug( f"['{self._conn_name}']: Current certificate " - f"expiration = {str(conn_info.expiration)}" + f"expiration = {conn_info.expiration!s}" ) self._cached = conn_info self._needs_refresh = False @@ -132,4 +133,3 @@ async def close(self) -> None: other cache types. """ self._closed = True - return diff --git a/google/cloud/sql/connector/monitored_cache.py b/google/cloud/sql/connector/monitored_cache.py index 0c3fc4d03..fb56cdac9 100644 --- a/google/cloud/sql/connector/monitored_cache.py +++ b/google/cloud/sql/connector/monitored_cache.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import asyncio import logging import ssl -from typing import Any, Callable, Optional, Union +from typing import Any, Callable from google.cloud.sql.connector.connection_info import ConnectionInfo from google.cloud.sql.connector.connection_info import ConnectionInfoCache @@ -31,13 +33,13 @@ class MonitoredCache(ConnectionInfoCache): def __init__( self, - cache: Union[RefreshAheadCache, LazyRefreshCache], + cache: RefreshAheadCache | LazyRefreshCache, failover_period: int, - resolver: Union[DefaultResolver, DnsResolver], + resolver: DefaultResolver | DnsResolver, ) -> None: self.resolver = resolver self.cache = cache - self.domain_name_ticker: Optional[asyncio.Task] = None + self.domain_name_ticker: asyncio.Task | None = None self.sockets: list[ssl.SSLSocket] = [] # If domain name is configured for instance and failover period is set, @@ -87,7 +89,7 @@ async def _check_domain_name(self) -> None: ) await self.close() - except Exception as e: + except Exception as e: # noqa: BLE001 # Domain name checks should not be fatal, log error and continue. logger.debug( f"['{self.cache.conn_name}']: Unable to check domain name, " diff --git a/google/cloud/sql/connector/pymysql.py b/google/cloud/sql/connector/pymysql.py index e01cfed08..c087e9bc1 100644 --- a/google/cloud/sql/connector/pymysql.py +++ b/google/cloud/sql/connector/pymysql.py @@ -47,7 +47,7 @@ def connect( ) # allow automatic IAM database authentication to not require password - kwargs["password"] = kwargs["password"] if "password" in kwargs else None + kwargs["password"] = kwargs.get("password", None) # pop timeout as timeout arg is called 'connect_timeout' for pymysql timeout = kwargs.pop("timeout") diff --git a/google/cloud/sql/connector/rate_limiter.py b/google/cloud/sql/connector/rate_limiter.py index be9c68c64..c5243e2df 100644 --- a/google/cloud/sql/connector/rate_limiter.py +++ b/google/cloud/sql/connector/rate_limiter.py @@ -13,12 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. """ +from __future__ import annotations import asyncio -from typing import Optional -class AsyncRateLimiter(object): +class AsyncRateLimiter: """An asyncio-compatible rate limiter which uses the Token Bucket algorithm (https://en.wikipedia.org/wiki/Token_bucket) to limit the number of function calls over a time interval using an event queue. @@ -35,7 +35,7 @@ def __init__( self, max_capacity: int = 1, rate: float = 1 / 60, - loop: Optional[asyncio.AbstractEventLoop] = None, + loop: asyncio.AbstractEventLoop | None = None, ) -> None: self.rate = rate self.max_capacity = max_capacity diff --git a/google/cloud/sql/connector/refresh_utils.py b/google/cloud/sql/connector/refresh_utils.py index 898f0f7a9..0edf5c345 100644 --- a/google/cloud/sql/connector/refresh_utils.py +++ b/google/cloud/sql/connector/refresh_utils.py @@ -72,7 +72,7 @@ async def _is_valid(task: asyncio.Task) -> bool: # only valid if now is before the cert expires if datetime.datetime.now(datetime.timezone.utc) < metadata.expiration: return True - except Exception: + except Exception: # noqa: BLE001 # supress any errors from task logger.debug("Current instance metadata is invalid.") return False @@ -80,7 +80,7 @@ async def _is_valid(task: asyncio.Task) -> bool: def _downscope_credentials( credentials: Credentials, - scopes: list[str] = ["https://www.googleapis.com/auth/sqlservice.login"], + scopes: list[str] | None = None, ) -> Credentials: """Generate a down-scoped credential. @@ -95,6 +95,8 @@ def _downscope_credentials( """ # credentials sourced from a service account or metadata are children of # Scoped class and are capable of being re-scoped + if scopes is None: + scopes = ["https://www.googleapis.com/auth/sqlservice.login"] if isinstance(credentials, Scoped): scoped_creds = credentials.with_scopes(scopes=scopes) # authenticated user credentials can not be re-scoped diff --git a/google/cloud/sql/connector/resolver.py b/google/cloud/sql/connector/resolver.py index e255f328a..7800f4956 100644 --- a/google/cloud/sql/connector/resolver.py +++ b/google/cloud/sql/connector/resolver.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List +from __future__ import annotations + +import re +from typing import Any import dns.asyncresolver @@ -24,10 +27,17 @@ from google.cloud.sql.connector.connection_name import ConnectionName from google.cloud.sql.connector.exceptions import DnsResolutionError +PSC_DNS_PATTERN = re.compile( + r"^([a-f0-9]{12})\.([^.]+)\.([a-z0-9]+-[a-z0-9]+)\.(sql|sql-psa|sql-psc)\.goog\.?$" +) + class DefaultResolver: """DefaultResolver simply validates and parses instance connection name.""" + def __init__(self, *args: Any, client: Any | None = None, **kwargs: Any) -> None: + pass + async def resolve(self, connection_name: str) -> ConnectionName: return _parse_connection_name(connection_name) @@ -38,54 +48,113 @@ class DnsResolver(dns.asyncresolver.Resolver): TXT records in DNS. """ + def __init__(self, *args: Any, client: Any | None = None, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._client = client + async def resolve(self, dns: str) -> ConnectionName: # type: ignore - try: - conn_name = _parse_connection_name(dns) - except ValueError: - # The connection name was not project:region:instance format. - # Check if connection name is a valid DNS domain name - if _is_valid_domain(dns): - # Attempt to query a TXT record to get connection name. - conn_name = await self.query_dns(dns) - else: + current = dns + visited = {current} + + # Max 10 iterations to prevent infinite CNAME loops + for _ in range(10): + try: + domain_val = dns if current != dns else "" + conn_name = _parse_connection_name_with_domain_name( + current, domain_val + ) + return conn_name + except ValueError: + pass + + dns_normalized = current.rstrip(".") + match = PSC_DNS_PATTERN.match(dns_normalized.lower()) + if match: + region = match.group(3) + if region != "global": + if self._client is None: + raise ValueError( + "SQLAdmin client is not configured in the resolver." + ) + + dns_name_with_dot = dns_normalized + "." + resp = await self._client.resolve_connect_settings( + dns_name_with_dot, region + ) + resolved_conn_name = resp["connectionName"] + return _parse_connection_name_with_domain_name( + resolved_conn_name, dns + ) + + if not _is_valid_domain(current): raise ValueError( "Arg `instance_connection_string` must have " "format: PROJECT:REGION:INSTANCE or be a valid DNS domain " f"name, got {dns}." ) - return conn_name - async def resolve_a_record(self, dns: str) -> List[str]: - try: - # Attempt to query the A records. - records = await super().resolve(dns, "A", raise_on_no_answer=True) - # return IP addresses as strings - return [record.to_text() for record in records] - except Exception: - # On any error, return empty list - return [] + cname_found = False + try: + cname = await self.resolve_cname(current) + cname_found = True + except DnsResolutionError: + pass + + if cname_found: + if cname in visited: + raise DnsResolutionError(f"CNAME loop detected for `{dns}`") + visited.add(cname) + current = cname + continue + + try: + rdata = await self.resolve_txt(current) + except DnsResolutionError as e: + raise DnsResolutionError( + f"Unable to resolve TXT record for `{dns}`" + ) from e - async def query_dns(self, dns: str) -> ConnectionName: - try: - # Attempt to query the TXT records. - records = await super().resolve(dns, "TXT", raise_on_no_answer=True) - # Sort the TXT record values alphabetically, strip quotes as record - # values can be returned as raw strings - rdata = [record.to_text().strip('"') for record in records] rdata.sort() - # Attempt to parse records, returning the first valid record. for record in rdata: try: - conn_name = _parse_connection_name_with_domain_name(record, dns) + conn_name = _parse_connection_name_with_domain_name( + record, dns + ) return conn_name - except Exception: + except Exception: # noqa: BLE001, S112 continue - # If all records failed to parse, throw error + raise DnsResolutionError( - f"Unable to parse TXT record for `{dns}` -> `{rdata[0]}`" + f"Unable to parse TXT record for `{current}` -> `{rdata[0]}`" + if rdata + else f"Unable to resolve TXT record for `{current}`" ) - # Don't override above DnsResolutionError - except DnsResolutionError: - raise + + raise DnsResolutionError( + f"CNAME loop detected or max resolution depth reached for `{dns}`" + ) + + async def resolve_cname(self, dns: str) -> str: + try: + answers = await super().resolve(dns, "CNAME", raise_on_no_answer=True) + return str(answers[0].target).rstrip(".") + except Exception as e: + raise DnsResolutionError( + f"Unable to resolve CNAME record for `{dns}`" + ) from e + + async def resolve_txt(self, dns: str) -> list[str]: + try: + answers = await super().resolve(dns, "TXT", raise_on_no_answer=True) + return [record.to_text().strip('"') for record in answers] except Exception as e: - raise DnsResolutionError(f"Unable to resolve TXT record for `{dns}`") from e + raise DnsResolutionError( + f"Unable to resolve TXT record for `{dns}`" + ) from e + + async def resolve_a_record(self, dns: str) -> list[str]: + try: + records = await super().resolve(dns, "A", raise_on_no_answer=True) + return [record.to_text() for record in records] + except Exception: # noqa: BLE001 + return [] diff --git a/google/cloud/sql/connector/utils.py b/google/cloud/sql/connector/utils.py old mode 100755 new mode 100644 index 2c355d965..e6d7ff4c2 --- a/google/cloud/sql/connector/utils.py +++ b/google/cloud/sql/connector/utils.py @@ -123,7 +123,7 @@ def format_database_user(database_version: str, user: str) -> str: # remove suffix for Postgres service accounts if database_version.startswith("POSTGRES"): suffix = ".gserviceaccount.com" - user = user[: -len(suffix)] if user.endswith(suffix) else user + user = user.removesuffix(suffix) return user # remove everything after and including the @ for MySQL diff --git a/noxfile.py b/noxfile.py index eb1b1b347..469d61ec8 100644 --- a/noxfile.py +++ b/noxfile.py @@ -14,7 +14,6 @@ limitations under the License. """ -from __future__ import absolute_import import os diff --git a/tests/conftest.py b/tests/conftest.py index 1f92a099e..883834b8a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,18 +15,19 @@ """ import asyncio +from collections.abc import AsyncGenerator import inspect import os import socket import ssl from threading import Thread -from typing import Any, AsyncGenerator +from typing import Any from unittest.mock import Mock import aiohttp from aiohttp import web from cryptography.hazmat.primitives import serialization -import pytest # noqa F401 Needed to run the tests +import pytest from unit.mocks import create_ssl_context # type: ignore from unit.mocks import FakeCredentials # type: ignore from unit.mocks import FakeCSQLInstance # type: ignore diff --git a/tests/system/test_asyncpg_connection.py b/tests/system/test_asyncpg_connection.py index 1aae1f7e6..00589942e 100644 --- a/tests/system/test_asyncpg_connection.py +++ b/tests/system/test_asyncpg_connection.py @@ -13,10 +13,11 @@ See the License for the specific language governing permissions and limitations under the License. """ +from __future__ import annotations import asyncio import os -from typing import Any, Union +from typing import Any import asyncpg import sqlalchemy @@ -34,7 +35,7 @@ async def create_sqlalchemy_engine( db: str, ip_type: str = "public", refresh_strategy: str = "background", - resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver, + resolver: type[DefaultResolver | DnsResolver] = DefaultResolver, **kwargs: Any, ) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, Connector]: """Creates a connection pool for a Cloud SQL instance and returns the pool diff --git a/tests/system/test_connector_object.py b/tests/system/test_connector_object.py index 3dd22fda5..5b05319f0 100644 --- a/tests/system/test_connector_object.py +++ b/tests/system/test_connector_object.py @@ -81,8 +81,6 @@ def test_multiple_connectors() -> None: (instance_connection_string, second_connector._enable_iam_auth) ] ) - except Exception: - raise finally: # close connectors first_connector.close() @@ -136,9 +134,10 @@ def test_connector_sqlserver_iam_auth_error() -> None: Test that connecting with enable_iam_auth set to True for SQL Server raises exception. """ - with pytest.raises(AutoIAMAuthNotSupported): - with Connector(enable_iam_auth=True) as connector: - connector.connect( + with pytest.raises(AutoIAMAuthNotSupported), Connector( + enable_iam_auth=True + ) as connector: + connector.connect( os.environ["SQLSERVER_CONNECTION_NAME"], "pytds", user="my-user", diff --git a/tests/system/test_pg8000_connection.py b/tests/system/test_pg8000_connection.py index c7074b0cf..0fef1b965 100644 --- a/tests/system/test_pg8000_connection.py +++ b/tests/system/test_pg8000_connection.py @@ -13,13 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. """ +from __future__ import annotations from datetime import datetime import os # [START cloud_sql_connector_postgres_pg8000] -from typing import Union - import sqlalchemy from google.cloud.sql.connector import Connector @@ -34,7 +33,7 @@ def create_sqlalchemy_engine( db: str, ip_type: str = "public", refresh_strategy: str = "background", - resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver, + resolver: type[DefaultResolver | DnsResolver] = DefaultResolver, ) -> tuple[sqlalchemy.engine.Engine, Connector]: """Creates a connection pool for a Cloud SQL instance and returns the pool and the connector. Callers are responsible for closing the pool and the diff --git a/tests/unit/mocks.py b/tests/unit/mocks.py index 42c5cc666..bba3a1d71 100644 --- a/tests/unit/mocks.py +++ b/tests/unit/mocks.py @@ -21,7 +21,7 @@ import datetime import json import ssl -from typing import Any, Callable, Literal, Optional +from typing import Any, Callable, Literal from aiohttp import web from cryptography import x509 @@ -42,7 +42,7 @@ class FakeCredentials: def __init__( - self, token: Optional[str] = None, expiry: Optional[datetime.datetime] = None + self, token: str | None = None, expiry: datetime.datetime | None = None ) -> None: self.token = token self.expiry = expiry @@ -70,9 +70,7 @@ def expired(self) -> bool: """ if self.expiry is None: return False - if self.expiry > datetime.datetime.now(datetime.timezone.utc): - return False - return True + return not self.expiry > datetime.datetime.now(datetime.timezone.utc) @property def universe_domain(self) -> str: @@ -112,13 +110,16 @@ def token_state( def generate_cert( project: str, name: str, - cert_before: datetime.datetime = datetime.datetime.now(datetime.timezone.utc), - cert_after: datetime.datetime = datetime.datetime.now(datetime.timezone.utc) - + datetime.timedelta(hours=1), + cert_before: datetime.datetime | None = None, + cert_after: datetime.datetime | None = None, ) -> tuple[x509.CertificateBuilder, rsa.RSAPrivateKey]: """ Generate a private key and cert object to be used in testing. """ + if cert_before is None: + cert_before = datetime.datetime.now(datetime.timezone.utc) + if cert_after is None: + cert_after = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=1) # generate private key key = rsa.generate_private_key(public_exponent=65537, key_size=2048) common_name = f"{project}:{name}" @@ -160,13 +161,16 @@ def client_key_signed_cert( cert: x509.CertificateBuilder, priv_key: rsa.RSAPrivateKey, client_key: rsa.RSAPublicKey, - cert_before: datetime.datetime = datetime.datetime.now(datetime.timezone.utc), - cert_expiration: datetime.datetime = datetime.datetime.now(datetime.timezone.utc) - + datetime.timedelta(hours=1), + cert_before: datetime.datetime | None = None, + cert_expiration: datetime.datetime | None = None, ) -> str: """ Create a PEM encoded certificate that is signed by given public key. """ + if cert_before is None: + cert_before = datetime.datetime.now(datetime.timezone.utc) + if cert_expiration is None: + cert_expiration = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=1) # configure cert subject subject = issuer = x509.Name( [ @@ -223,20 +227,26 @@ def __init__( region: str = "test-region", name: str = "test-instance", db_version: str = "POSTGRES_15", - ip_addrs: dict = { - "PRIMARY": "127.0.0.1", - "PRIVATE": "10.0.0.1", - }, + ip_addrs: dict | None = None, + dns_names: list | None = None, legacy_dns_name: bool = False, - cert_before: datetime = datetime.datetime.now(datetime.timezone.utc), - cert_expiration: datetime = datetime.datetime.now(datetime.timezone.utc) - + datetime.timedelta(hours=1), + cert_before: datetime.datetime | None = None, + cert_expiration: datetime.datetime | None = None, ) -> None: + if cert_before is None: + cert_before = datetime.datetime.now(datetime.timezone.utc) + if cert_expiration is None: + cert_expiration = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=1) + if dns_names is None: + dns_names = ["abcde.12345.us-central1.sql.goog"] + if ip_addrs is None: + ip_addrs = {"PRIMARY": "127.0.0.1", "PRIVATE": "10.0.0.1"} self.project = project self.region = region self.name = name self.db_version = db_version self.ip_addrs = ip_addrs + self.dns_names = dns_names self.psc_enabled = False self.cert_before = cert_before self.cert_expiration = cert_expiration @@ -265,14 +275,15 @@ async def connect_settings(self, request: Any) -> web.Response: "databaseVersion": self.db_version, } if self.legacy_dns_name: - response["dnsName"] = "abcde.12345.us-central1.sql.goog" + response["dnsName"] = self.dns_names[0] if self.dns_names else None else: response["dnsNames"] = [ { - "name": "abcde.12345.us-central1.sql.goog", + "name": name, "connectionType": "PRIVATE_SERVICE_CONNECT", "dnsScope": "INSTANCE", } + for name in self.dns_names ] return web.Response(content_type="application/json", body=json.dumps(response)) diff --git a/tests/unit/test_asyncpg.py b/tests/unit/test_asyncpg.py index b29df8841..ee1da1ecf 100644 --- a/tests/unit/test_asyncpg.py +++ b/tests/unit/test_asyncpg.py @@ -16,9 +16,9 @@ import ssl from typing import Any +from unittest.mock import AsyncMock +from unittest.mock import patch -from mock import AsyncMock -from mock import patch import pytest from google.cloud.sql.connector.asyncpg import connect diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index cfe509470..e9f7eb88b 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -12,8 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import datetime -from typing import Optional from aiohttp import ClientResponseError from aioresponses import aioresponses @@ -38,8 +39,8 @@ async def test_get_metadata_no_psc(fake_client: CloudSQLClient) -> None: ) assert resp["database_version"] == "POSTGRES_15" assert resp["ip_addresses"] == { - "PRIMARY": "127.0.0.1", - "PRIVATE": "10.0.0.1", + "PRIMARY": ["127.0.0.1"], + "PRIVATE": ["10.0.0.1"], } assert isinstance(resp["server_ca_cert"], str) @@ -58,9 +59,9 @@ async def test_get_metadata_with_psc(fake_client: CloudSQLClient) -> None: ) assert resp["database_version"] == "POSTGRES_15" assert resp["ip_addresses"] == { - "PRIMARY": "127.0.0.1", - "PRIVATE": "10.0.0.1", - "PSC": "abcde.12345.us-central1.sql.goog", + "PRIMARY": ["127.0.0.1"], + "PRIVATE": ["10.0.0.1"], + "PSC": ["abcde.12345.us-central1.sql.goog"], } assert isinstance(resp["server_ca_cert"], str) @@ -80,9 +81,9 @@ async def test_get_metadata_legacy_dns_with_psc(fake_client: CloudSQLClient) -> ) assert resp["database_version"] == "POSTGRES_15" assert resp["ip_addresses"] == { - "PRIMARY": "127.0.0.1", - "PRIVATE": "10.0.0.1", - "PSC": "abcde.12345.us-central1.sql.goog", + "PRIMARY": ["127.0.0.1"], + "PRIVATE": ["10.0.0.1"], + "PSC": ["abcde.12345.us-central1.sql.goog"], } assert isinstance(resp["server_ca_cert"], str) @@ -148,7 +149,7 @@ async def test_CloudSQLClient_init_custom_user_agent( ) @pytest.mark.asyncio async def test_CloudSQLClient_user_agent( - driver: Optional[str], fake_credentials: FakeCredentials + driver: str | None, fake_credentials: FakeCredentials ) -> None: """ Test to check whether the __init__ method of CloudSQLClient @@ -290,3 +291,35 @@ async def test_get_ephemeral_error_parsing_json( assert exc_info.value.status == 404 assert exc_info.value.message == "Not Found" await client.close() + + +@pytest.mark.asyncio +async def test_get_metadata_multiple_psc_dns_sorted(fake_client: CloudSQLClient) -> None: + """ + Test _get_metadata returns successfully with multiple PSC IP types sorted. + """ + fake_client.instance.psc_enabled = True + fake_client.instance.legacy_dns_name = False + fake_client.instance.dns_names = [ + "dns1.sql.goog", + "dns2.sql-psc.goog", + "dns3.sql.goog", + ] + try: + resp = await fake_client._get_metadata( + "test-project", + "test-region", + "test-instance", + ) + assert resp["database_version"] == "POSTGRES_15" + assert resp["ip_addresses"] == { + "PRIMARY": ["127.0.0.1"], + "PRIVATE": ["10.0.0.1"], + "PSC": ["dns2.sql-psc.goog", "dns1.sql.goog", "dns3.sql.goog"], + } + assert isinstance(resp["server_ca_cert"], str) + finally: + fake_client.instance.psc_enabled = False + fake_client.instance.legacy_dns_name = False + fake_client.instance.dns_names = ["abcde.12345.us-central1.sql.goog"] + diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py index a09b5b72f..8ea8c88c5 100644 --- a/tests/unit/test_connector.py +++ b/tests/unit/test_connector.py @@ -13,16 +13,16 @@ See the License for the specific language governing permissions and limitations under the License. """ +from __future__ import annotations import asyncio import os from threading import Thread -from typing import Union +from unittest.mock import patch from aiohttp import ClientResponseError from google.auth.credentials import Credentials -from mock import patch -import pytest # noqa F401 Needed to run the tests +import pytest from google.cloud.sql.connector import Connector from google.cloud.sql.connector import create_async_connector @@ -215,7 +215,7 @@ async def test_Connector_Init_async_context_manager( ], ) def test_Connector_init_ip_type( - ip_type: Union[str, IPTypes], expected: IPTypes, fake_credentials: Credentials + ip_type: str | IPTypes, expected: IPTypes, fake_credentials: Credentials ) -> None: """ Test to check whether the __init__ method of Connector @@ -659,3 +659,66 @@ async def test_Connector_connect_async_custom_dns_resolver_fallback( fake_client.instance.ip_addrs = original_ips +@pytest.mark.asyncio +async def test_Connector_connect_async_custom_dns_resolver_no_fallback_psc_to_private_ip( + fake_credentials: Credentials, fake_client: CloudSQLClient +) -> None: + """Test that Connector.connect_async does not fall back to Private IP if CNAME/PSC DNS resolution fails.""" + + with patch( + "google.cloud.sql.connector.resolver.DnsResolver.resolve_a_record" + ) as mock_resolve_a: + # DNS resolution fails + mock_resolve_a.return_value = [] + + with patch( + "google.cloud.sql.connector.resolver.DnsResolver.resolve" + ) as mock_resolve: + conn_name_with_domain = ConnectionName( + "test-project", "test-region", "test-instance", "db.example.com" + ) + mock_resolve.return_value = conn_name_with_domain + + async with Connector( + credentials=fake_credentials, + loop=asyncio.get_running_loop(), + resolver=DnsResolver, + ip_type="PSC", # Use PSC IP type + ) as connector: + connector._client = fake_client + + original_ips = fake_client.instance.ip_addrs + original_dns_names = fake_client.instance.dns_names + # Configure instance to be PSC enabled, but also have a PRIVATE IP! + fake_client.instance.psc_enabled = True + fake_client.instance.dns_names = ["1ad3b5d73f10.3oxon2yfo9tob.us-east1.sql.goog"] + fake_client.instance.ip_addrs = { + "PSC": "1ad3b5d73f10.3oxon2yfo9tob.us-east1.sql.goog", + "PRIVATE": "10.0.0.1", + } + + try: + with patch( + "google.cloud.sql.connector.asyncpg.connect" + ) as mock_connect: + mock_connect.return_value = True + + connection = await connector.connect_async( + "db.example.com", + "asyncpg", + user="my-user", + password="my-pass", + db="my-db", + ) + + # Verify mock_connect used PSC DNS instead of falling back to PRIVATE IP "10.0.0.1"! + args, _ = mock_connect.call_args + assert args[0] == "1ad3b5d73f10.3oxon2yfo9tob.us-east1.sql.goog" + assert connection is True + finally: + # Restore original IPs and DNS names + fake_client.instance.ip_addrs = original_ips + fake_client.instance.dns_names = original_dns_names + fake_client.instance.psc_enabled = False + + diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py index 3699ddc2d..3dbba59b8 100644 --- a/tests/unit/test_instance.py +++ b/tests/unit/test_instance.py @@ -16,10 +16,10 @@ import asyncio import datetime +from unittest.mock import patch -from mock import patch import mocks -import pytest # noqa F401 Needed to run the tests +import pytest from google.cloud.sql.connector import IPTypes from google.cloud.sql.connector.client import CloudSQLClient @@ -109,7 +109,7 @@ async def test_schedule_refresh_replaces_invalid_result( a valid one """ # allow more frequent refreshes for tests - setattr(cache, "_refresh_rate_limiter", test_rate_limiter) + cache._refresh_rate_limiter = test_rate_limiter # set certificate to be expired cache._client.instance.cert_expiration = datetime.datetime.now( datetime.timezone.utc @@ -147,7 +147,7 @@ async def test_force_refresh_cancels_pending_refresh( Test that force_refresh cancels pending task if refresh_in_progress event is not set. """ # allow more frequent refreshes for tests - setattr(cache, "_refresh_rate_limiter", test_rate_limiter) + cache._refresh_rate_limiter = test_rate_limiter # make sure initial refresh is finished await cache._current # since the pending refresh isn't for another 55 min, the refresh_in_progress event @@ -215,7 +215,7 @@ async def test_perform_refresh_expiration( minutes=1 ) credentials = mocks.FakeCredentials(token="my-token", expiry=expiration) - setattr(cache, "_enable_iam_auth", True) + cache._enable_iam_auth = True # set downscoped credential to mock with patch("google.cloud.sql.connector.client._downscope_credentials") as mock_auth: mock_auth.return_value = credentials @@ -249,13 +249,13 @@ async def test_get_preferred_ip_CloudSQLIPTypeError(cache: RefreshAheadCache) -> when missing Public or Private IP addresses. """ instance_metadata: ConnectionInfo = await cache._current - instance_metadata.ip_addrs = {"PRIVATE": "1.1.1.1"} + instance_metadata.ip_addrs = {"PRIVATE": ["1.1.1.1"]} # test error when Public IP is missing with pytest.raises(CloudSQLIPTypeError): instance_metadata.get_preferred_ip(IPTypes.PUBLIC) # test error when Private IP is missing - instance_metadata.ip_addrs = {"PRIMARY": "0.0.0.0"} + instance_metadata.ip_addrs = {"PRIMARY": ["0.0.0.0"]} with pytest.raises(CloudSQLIPTypeError): instance_metadata.get_preferred_ip(IPTypes.PRIVATE) @@ -284,7 +284,7 @@ async def test_AutoIAMAuthNotSupportedError(fake_client: CloudSQLClient) -> None async def test_ConnectionInfo_caches_sslcontext() -> None: info = ConnectionInfo( - "", "cert", "cert", "key".encode(), {}, "POSTGRES", datetime.datetime.now() + "", "cert", "cert", b"key", {}, "POSTGRES", datetime.datetime.now(datetime.timezone.utc) ) # context should default to None assert info.context is None diff --git a/tests/unit/test_monitored_cache.py b/tests/unit/test_monitored_cache.py index 1c1f1df86..1460ba1dd 100644 --- a/tests/unit/test_monitored_cache.py +++ b/tests/unit/test_monitored_cache.py @@ -15,12 +15,12 @@ import asyncio import socket import ssl +from unittest.mock import patch import dns.message import dns.rdataclass import dns.rdatatype import dns.resolver -from mock import patch import pytest from google.cloud.sql.connector.client import CloudSQLClient diff --git a/tests/unit/test_pg8000.py b/tests/unit/test_pg8000.py index 2c003b8a9..c291d6f8a 100644 --- a/tests/unit/test_pg8000.py +++ b/tests/unit/test_pg8000.py @@ -17,8 +17,8 @@ import socket import ssl from typing import Any +from unittest.mock import patch -from mock import patch import pytest from google.cloud.sql.connector.pg8000 import connect diff --git a/tests/unit/test_pymysql.py b/tests/unit/test_pymysql.py index 13cd8e98a..ef4b23b4b 100644 --- a/tests/unit/test_pymysql.py +++ b/tests/unit/test_pymysql.py @@ -17,8 +17,8 @@ import socket import ssl from typing import Any +from unittest.mock import patch -from mock import patch import pytest from google.cloud.sql.connector.pymysql import connect as pymysql_connect diff --git a/tests/unit/test_pytds.py b/tests/unit/test_pytds.py index faa20ad8c..0d9ca98c3 100644 --- a/tests/unit/test_pytds.py +++ b/tests/unit/test_pytds.py @@ -18,8 +18,8 @@ import socket import ssl from typing import Any +from unittest.mock import patch -from mock import patch import pytest from google.cloud.sql.connector.exceptions import PlatformNotSupportedError @@ -58,7 +58,7 @@ async def test_pytds_platform_error(context: ssl.SSLContext, kwargs: Any) -> Non """Test to verify that pytds.connect throws proper PlatformNotSupportedError.""" ip_addr = "127.0.0.1" # stub operating system to Linux - setattr(platform, "system", stub_platform_linux) + platform.system = stub_platform_linux assert platform.system() == "Linux" sock = context.wrap_socket( socket.create_connection((ip_addr, 3307)), @@ -81,7 +81,7 @@ async def test_pytds_windows_active_directory_auth( """ ip_addr = "127.0.0.1" # stub operating system to Windows - setattr(platform, "system", stub_platform_windows) + platform.system = stub_platform_windows assert platform.system() == "Windows" sock = context.wrap_socket( socket.create_connection((ip_addr, 3307)), diff --git a/tests/unit/test_rate_limiter.py b/tests/unit/test_rate_limiter.py index 8ef586b58..de4418ef9 100644 --- a/tests/unit/test_rate_limiter.py +++ b/tests/unit/test_rate_limiter.py @@ -16,7 +16,7 @@ import asyncio -import pytest # noqa F401 Needed to run the tests +import pytest from google.cloud.sql.connector.rate_limiter import AsyncRateLimiter diff --git a/tests/unit/test_refresh_utils.py b/tests/unit/test_refresh_utils.py index 119e92c7a..0db71f18a 100644 --- a/tests/unit/test_refresh_utils.py +++ b/tests/unit/test_refresh_utils.py @@ -18,15 +18,15 @@ import asyncio import datetime +from unittest.mock import Mock +from unittest.mock import patch from conftest import SCOPES # type: ignore import google.auth from google.auth.credentials import Credentials from google.auth.credentials import TokenState import google.oauth2.credentials -from mock import Mock -from mock import patch -import pytest # noqa F401 Needed to run the tests +import pytest from google.cloud.sql.connector.refresh_utils import _downscope_credentials from google.cloud.sql.connector.refresh_utils import _exponential_backoff diff --git a/tests/unit/test_resolver.py b/tests/unit/test_resolver.py index c649e8e58..a7c52ed01 100644 --- a/tests/unit/test_resolver.py +++ b/tests/unit/test_resolver.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import AsyncMock +from unittest.mock import patch + import dns.message import dns.rdataclass import dns.rdatatype import dns.resolver -from mock import patch import pytest from google.cloud.sql.connector.connection_name import ConnectionName @@ -165,4 +167,130 @@ async def test_DnsResolver_resolve_a_record_empty() -> None: mock_resolve.side_effect = Exception("DNS Error") resolver = DnsResolver() result = await resolver.resolve_a_record("db.example.com") - assert result == [] \ No newline at end of file + assert result == [] + + + +async def test_DnsResolver_with_direct_psc_dns_name() -> None: + """Test DnsResolver resolves direct PSC DNS name using resolve_connect_settings.""" + dns_name = "0123456789ab.fedcba9876543.europe-north2.sql-psc.goog" + real_conn_name = ConnectionName( + "my-project", "europe-north2", "my-instance", dns_name + ) + + mock_client = AsyncMock() + mock_client.resolve_connect_settings.return_value = { + "connectionName": "my-project:europe-north2:my-instance" + } + + resolver = DnsResolver(client=mock_client) + + result = await resolver.resolve(dns_name) + + assert result == real_conn_name + # Verify mock_client was called with correct trailing dot DNS name! + mock_client.resolve_connect_settings.assert_awaited_once_with( + dns_name + ".", "europe-north2" + ) + + +async def test_DnsResolver_with_cname_resolving_to_psc_dns_name() -> None: + """Test DnsResolver resolves CNAME to PSC DNS and returns proper connection name.""" + dns_name = "db.example.com" + cname_target = "0123456789ab.fedcba9876543.europe-north2.sql-psc.goog" + real_conn_name = ConnectionName( + "my-project", "europe-north2", "my-instance", dns_name + ) + + mock_client = AsyncMock() + mock_client.resolve_connect_settings.return_value = { + "connectionName": "my-project:europe-north2:my-instance" + } + + resolver = DnsResolver(client=mock_client) + + # Patch resolver CNAME and TXT methods + with patch.object( + resolver, "resolve_cname", AsyncMock(return_value=cname_target) + ), patch.object( + resolver, "resolve_txt", AsyncMock(side_effect=Exception("No TXT")) + ): + + result = await resolver.resolve(dns_name) + + assert result == real_conn_name + mock_client.resolve_connect_settings.assert_awaited_once_with( + cname_target + ".", "europe-north2" + ) + + +async def test_DnsResolver_with_recursive_cnames_to_psc_dns_name() -> None: + """Test DnsResolver resolves recursive CNAMEs to PSC DNS successfully.""" + dns_name = "name1.example.com" + cname2 = "name2.example.com" + cname_target = "0123456789ab.fedcba9876543.europe-north2.sql-psc.goog" + real_conn_name = ConnectionName( + "my-project", "europe-north2", "my-instance", dns_name + ) + + mock_client = AsyncMock() + mock_client.resolve_connect_settings.return_value = { + "connectionName": "my-project:europe-north2:my-instance" + } + + resolver = DnsResolver(client=mock_client) + + # Mock Lookup CNAME sequence + cname_mock = AsyncMock( + side_effect=lambda name: cname2 if name == dns_name else cname_target + ) + + with patch.object(resolver, "resolve_cname", cname_mock), patch.object( + resolver, "resolve_txt", AsyncMock(side_effect=Exception("No TXT")) + ): + + result = await resolver.resolve(dns_name) + + assert result == real_conn_name + mock_client.resolve_connect_settings.assert_awaited_once_with( + cname_target + ".", "europe-north2" + ) + + +async def test_DnsResolver_cname_loop_throws_error() -> None: + """Test DnsResolver throws error if a CNAME loop is detected.""" + dns_name = "name1.example.com" + cname2 = "name2.example.com" + + resolver = DnsResolver() + + cname_mock = AsyncMock( + side_effect=lambda name: cname2 if name == dns_name else dns_name + ) + + with patch.object(resolver, "resolve_cname", cname_mock), patch.object( + resolver, "resolve_txt", AsyncMock(side_effect=Exception("No TXT")) + ): + + with pytest.raises(DnsResolutionError) as exc_info: + await resolver.resolve(dns_name) + assert "CNAME loop detected" in str(exc_info.value) + + +async def test_DnsResolver_global_region_skips_direct_resolution() -> None: + """Test DnsResolver skips direct resolution if the PSC DNS name has a global region.""" + dns_name = "0123456789ab.fedcba9876543.global.sql-psc.goog" + + mock_client = AsyncMock() + resolver = DnsResolver(client=mock_client) + + # Patch CNAME and TXT to fail, so it eventually raises DnsResolutionError + with patch.object( + resolver, "resolve_cname", AsyncMock(side_effect=DnsResolutionError("No CNAME")) + ), patch.object( + resolver, "resolve_txt", AsyncMock(side_effect=DnsResolutionError("No TXT")) + ), pytest.raises(DnsResolutionError): + await resolver.resolve(dns_name) + + # Verify mock_client was NOT called because direct resolution was skipped + mock_client.resolve_connect_settings.assert_not_called() \ No newline at end of file diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index fe4e90955..1c2ae4fdf 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -14,7 +14,7 @@ limitations under the License. """ -import pytest # noqa F401 Needed to run the tests +import pytest from google.cloud.sql.connector import utils