From 90afe68c74d42fde6db913c04d268fc6f86cf44c Mon Sep 17 00:00:00 2001 From: Jonathan Hess Date: Mon, 17 Aug 2026 23:52:52 +0000 Subject: [PATCH 1/2] feat: proactively probe database on Auto-IAM refresh to update MCP tokens --- google/cloud/sql/connector/instance.py | 59 ++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/google/cloud/sql/connector/instance.py b/google/cloud/sql/connector/instance.py index 28ab54e4..a50c6c9d 100644 --- a/google/cloud/sql/connector/instance.py +++ b/google/cloud/sql/connector/instance.py @@ -119,6 +119,8 @@ async def _perform_refresh(self) -> ConnectionInfo: self._keys, self._enable_iam_auth, ) + if self._enable_iam_auth: + await self._probe_connection(connection_info) logger.debug( f"['{self._conn_name}']: Connection info refresh operation complete" ) @@ -138,6 +140,63 @@ async def _perform_refresh(self) -> ConnectionInfo: self._refresh_in_progress.clear() return connection_info + async def _probe_connection(self, conn_info: ConnectionInfo) -> None: + """Proactively probes the database to refresh IAM tokens on server-side MCP.""" + targets: list[str] = [] + if self._conn_name.domain_name: + targets.append(self._conn_name.domain_name) + else: + for ip_type in ("PSC", "PRIVATE", "PUBLIC"): + if ip_type in conn_info.ip_addrs: + targets.extend(conn_info.ip_addrs[ip_type]) + + if not targets: + logger.debug( + f"['{self._conn_name}']: Proactive IAM token refresh probe skipped: no target IP addresses" + ) + return + + port = 3307 + try: + ssl_context = await conn_info.create_ssl_context(self._enable_iam_auth) + except Exception as e: # noqa: BLE001 + logger.debug( + f"['{self._conn_name}']: Failed to create SSL context for probe: {e!s}" + ) + return + + for target in targets: + try: + logger.debug( + f"['{self._conn_name}']: Probing IAM token refresh on {target}:{port}" + ) + _, writer = await asyncio.wait_for( + asyncio.open_connection( + host=target, + port=port, + ssl=ssl_context, + server_hostname=( + self._conn_name.domain_name + if self._conn_name.domain_name + else None + ), + ), + timeout=15.0, + ) + writer.close() + await writer.wait_closed() + logger.debug( + f"['{self._conn_name}']: Proactive IAM token refresh probe successful" + ) + return + except Exception as e: # noqa: BLE001 + logger.debug( + f"['{self._conn_name}']: Probing IAM token refresh on {target}:{port} failed: {e!s}" + ) + logger.debug( + f"['{self._conn_name}']: Proactive IAM token refresh probe encountered error across all targets" + ) + def _schedule_refresh(self, delay: int) -> asyncio.Task: """ Schedule task to sleep and then perform refresh to get ConnectionInfo. From 36bf1a102c1b270f1d59bc6bd911e52281c9978a Mon Sep 17 00:00:00 2001 From: Jonathan Hess Date: Tue, 18 Aug 2026 02:10:10 +0000 Subject: [PATCH 2/2] fix: address code review comments on proactive IAM token probe Code review comments addressed: - Respect the IP settings in connection configuration (self._ip_type) rather than probing all available IPs - Use SERVER_PROXY_PORT constant (3307) - Use DEFAULT_CONNECT_TIMEOUT constant (30) / configured timeout --- google/cloud/sql/connector/connector.py | 2 ++ google/cloud/sql/connector/instance.py | 20 +++++++++++++++----- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py index 7a9964ca..ee362e5d 100644 --- a/google/cloud/sql/connector/connector.py +++ b/google/cloud/sql/connector/connector.py @@ -348,6 +348,8 @@ async def connect_async( self._client, self._keys, enable_iam_auth, + ip_type=self._ip_type, + timeout=self._timeout, ) # wrap cache as a MonitoredCache monitored_cache = MonitoredCache( diff --git a/google/cloud/sql/connector/instance.py b/google/cloud/sql/connector/instance.py index a50c6c9d..b226049f 100644 --- a/google/cloud/sql/connector/instance.py +++ b/google/cloud/sql/connector/instance.py @@ -26,6 +26,7 @@ from google.cloud.sql.connector.connection_info import ConnectionInfo from google.cloud.sql.connector.connection_info import ConnectionInfoCache from google.cloud.sql.connector.connection_name import ConnectionName +from google.cloud.sql.connector.enums import IPTypes from google.cloud.sql.connector.exceptions import RefreshNotValidError from google.cloud.sql.connector.rate_limiter import AsyncRateLimiter from google.cloud.sql.connector.refresh_utils import _is_valid @@ -34,6 +35,8 @@ logger = logging.getLogger(name=__name__) APPLICATION_NAME = "cloud-sql-python-connector" +SERVER_PROXY_PORT = 3307 +DEFAULT_CONNECT_TIMEOUT = 30 class RefreshAheadCache(ConnectionInfoCache): @@ -50,6 +53,8 @@ def __init__( client: CloudSQLClient, keys: asyncio.Future, enable_iam_auth: bool = False, + ip_type: IPTypes | str = IPTypes.PUBLIC, + timeout: int = DEFAULT_CONNECT_TIMEOUT, ) -> None: """Initializes a RefreshAheadCache instance. @@ -62,10 +67,16 @@ def __init__( enable_iam_auth (bool): Enables automatic IAM database authentication (Postgres and MySQL) as the default authentication method for all connections. + ip_type (IPTypes | str): Preferred IP type used to connect to the instance. + timeout (int): Connect timeout in seconds. """ self._conn_name = conn_name self._enable_iam_auth = enable_iam_auth + if isinstance(ip_type, str): + ip_type = IPTypes._from_str(ip_type) + self._ip_type = ip_type + self._timeout = timeout self._keys = keys self._client = client self._refresh_rate_limiter = AsyncRateLimiter( @@ -146,9 +157,8 @@ async def _probe_connection(self, conn_info: ConnectionInfo) -> None: if self._conn_name.domain_name: targets.append(self._conn_name.domain_name) else: - for ip_type in ("PSC", "PRIVATE", "PUBLIC"): - if ip_type in conn_info.ip_addrs: - targets.extend(conn_info.ip_addrs[ip_type]) + if self._ip_type.value in conn_info.ip_addrs: + targets.extend(conn_info.ip_addrs[self._ip_type.value]) if not targets: logger.debug( @@ -156,7 +166,7 @@ async def _probe_connection(self, conn_info: ConnectionInfo) -> None: ) return - port = 3307 + port = SERVER_PROXY_PORT try: ssl_context = await conn_info.create_ssl_context(self._enable_iam_auth) except Exception as e: # noqa: BLE001 @@ -181,7 +191,7 @@ async def _probe_connection(self, conn_info: ConnectionInfo) -> None: else None ), ), - timeout=15.0, + timeout=float(self._timeout), ) writer.close() await writer.wait_closed()