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 28ab54e4..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( @@ -119,6 +130,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 +151,62 @@ 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: + 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( + f"['{self._conn_name}']: Proactive IAM token refresh probe skipped: no target IP addresses" + ) + return + + port = SERVER_PROXY_PORT + 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=float(self._timeout), + ) + 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.