refactor: Instantiate resolver with API client in Python connector

- Add __init__ to DefaultResolver and DnsResolver to accept client in constructor.
- Remove set_client method from DnsResolver.
- Lazy-initialize self._resolver in Connector.connect_async passing the initialized self._client.
- Update tests to use constructor injection instead of set_client.
- Fix lint (ruff) and type checking (mypy) errors.

Author: hessjcg

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()
@@ -195,7 +195,7 @@ async def resolve_connect_settings(
             A dictionary containing the resolve response (e.g. connectionName).
         """
         # 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())
 
         headers = {
@@ -213,7 +213,7 @@ async def resolve_connect_settings(
                 message = ret_dict.get("error", {}).get("message")
                 if message:
                     resp.reason = message
-        except Exception:
+        except Exception:  # noqa: BLE001, S110
             pass
         finally:
             resp.raise_for_status()
@@ -266,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()
@@ -287,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(
@@ -317,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.

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:

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,8 +317,8 @@ async def connect_async(
                 user_agent=self._user_agent,
                 driver=driver,
             )
-            if hasattr(self._resolver, "set_client"):
-                self._resolver.set_client(self._client)
+        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)
@@ -331,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,
@@ -414,7 +415,7 @@ async def connect_async(
                         f"entries, using '{preferred_ips}' from instance metadata"
                     )
                     targets.extend(preferred_ips)
-            except Exception as e:
+            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 "
@@ -476,7 +477,7 @@ async def connect_async(
                     conn = await self._loop.run_in_executor(None, connect_partial)
                     last_ex = None
                     return conn
-                except Exception as e:
+                except Exception as e:  # noqa: BLE001
                     logger.debug(
                         f"['{conn_info.conn_name}']: Connection to {target_ip} failed: {e}"
                     )
@@ -509,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()
@@ -522,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()
@@ -558,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.

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"
 

google/cloud/sql/connector/exceptions.py

@@ -21,31 +21,27 @@ class ConnectorLoopError(Exception):
     in an invalid state (event loop in current thread).
     """
 
-    pass
 
 
 class TLSVersionError(Exception):
     """
     Raised when the required TLS protocol version is not supported.
     """
 
-    pass
 
 
 class CloudSQLIPTypeError(Exception):
     """
     Raised when IP address for the preferred IP type is not found.
     """
 
-    pass
 
 
 class PlatformNotSupportedError(Exception):
     """
     Raised when a feature is not supported on the current platform.
     """
 
-    pass
 
 
 class AutoIAMAuthNotSupported(Exception):
@@ -54,15 +50,13 @@ class AutoIAMAuthNotSupported(Exception):
     supported with database engine version.
     """
 
-    pass
 
 
 class RefreshNotValidError(Exception):
     """
     Exception to be raised when the task returned from refresh is not valid.
     """
 
-    pass
 
 
 class IncompatibleDriverError(Exception):

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