cloud-sql-python-connector 1.22.0__py3-none-any.whl → 1.24.0__py3-none-any.whl

google/cloud/sql/connector/enums.py

@@ -41,6 +41,7 @@ class IPTypes(Enum):
     PUBLIC = "PRIMARY"
     PRIVATE = "PRIVATE"
     PSC = "PSC"
+    SQL_DATA = "SQL_DATA"
 
     @classmethod
     def _missing_(cls, value: object) -> None:
@@ -54,6 +55,8 @@ class IPTypes(Enum):
         """Convert IP type from a str into IPTypes."""
         if ip_type_str.upper() == "PUBLIC":
             ip_type_str = "PRIMARY"
+        elif ip_type_str.upper() in ("SQLDATA", "SQL_DATA"):
+            ip_type_str = "SQL_DATA"
         return cls(ip_type_str.upper())
 
 

google/cloud/sql/connector/exceptions.py

@@ -14,6 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from __future__ import annotations
+
 
 class ConnectorLoopError(Exception):
     """
@@ -91,3 +93,34 @@ class ClosedConnectorError(Exception):
     Exception to be raised when a Connector is closed and connect method is
     called on it.
     """
+
+
+class CloudSQLConnectionError(Exception):
+    """
+    Exception to be raised when a connection cannot be established to a Cloud SQL instance.
+    """
+
+
+class ResourceExhaustedError(CloudSQLConnectionError):
+    """
+    Exception to be raised when a connection cannot be established because
+    the SQL Data Service is busy / in cooldown due to resource exhaustion.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        connection_name: str | None = None,
+        raw_error: Exception | None = None,
+    ) -> None:
+        self.message = message
+        self.connection_name = connection_name
+        self.raw_error = raw_error
+        super().__init__(
+            f"[{connection_name}] {message}: {raw_error}"
+            if connection_name and raw_error
+            else f"[{connection_name}] {message}"
+            if connection_name
+            else message
+        )
+

google/cloud/sql/connector/instance.py

@@ -17,16 +17,21 @@ limitations under the License.
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Awaitable
+from collections.abc import Callable
 from datetime import datetime
 from datetime import timedelta
-from datetime import timezone
+from datetime import UTC
 import logging
 
 from google.cloud.sql.connector.client import CloudSQLClient
 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.postgres_connection import build_postgres_startup_packet
+from google.cloud.sql.connector.postgres_connection import build_terminate_packet
 from google.cloud.sql.connector.rate_limiter import AsyncRateLimiter
 from google.cloud.sql.connector.refresh_utils import _is_valid
 from google.cloud.sql.connector.refresh_utils import _seconds_until_refresh
@@ -34,6 +39,8 @@ from google.cloud.sql.connector.refresh_utils import _seconds_until_refresh
 logger = logging.getLogger(name=__name__)
 
 APPLICATION_NAME = "cloud-sql-python-connector"
+SERVER_PROXY_PORT = 3307
+DEFAULT_CONNECT_TIMEOUT = 30
 
 
 class RefreshAheadCache(ConnectionInfoCache):
@@ -50,6 +57,15 @@ class RefreshAheadCache(ConnectionInfoCache):
         client: CloudSQLClient,
         keys: asyncio.Future,
         enable_iam_auth: bool = False,
+        ip_type: IPTypes | str = IPTypes.PUBLIC,
+        timeout: int = DEFAULT_CONNECT_TIMEOUT,
+        dial_func: (
+            Callable[
+                [ConnectionInfo],
+                Awaitable[tuple[asyncio.StreamReader, asyncio.StreamWriter]],
+            ]
+            | None
+        ) = None,
     ) -> None:
         """Initializes a RefreshAheadCache instance.
 
@@ -62,10 +78,18 @@ class RefreshAheadCache(ConnectionInfoCache):
             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.
+            dial_func: Optional callable to dial probe connections.
         """
         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._dial_func = dial_func
         self._keys = keys
         self._client = client
         self._refresh_rate_limiter = AsyncRateLimiter(
@@ -73,10 +97,15 @@ class RefreshAheadCache(ConnectionInfoCache):
             rate=1 / 30,
         )
         self._refresh_in_progress = asyncio.locks.Event()
+        self._iam_principals: list[tuple[str, str]] = []
         self._current: asyncio.Task = self._schedule_refresh(0)
         self._next: asyncio.Task = self._current
         self._closed = False
 
+    def record_principal(self, user: str, database: str) -> None:
+        """Records a PostgreSQL (user, database) pair observed on an IAM connection."""
+        _append_iam_principal(self._iam_principals, user, database)
+
     @property
     def conn_name(self) -> ConnectionName:
         return self._conn_name
@@ -119,6 +148,8 @@ class RefreshAheadCache(ConnectionInfoCache):
                 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 +169,18 @@ class RefreshAheadCache(ConnectionInfoCache):
             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."""
+        await _probe_instance_connection(
+            self._conn_name,
+            conn_info,
+            self._enable_iam_auth,
+            self._ip_type,
+            self._timeout,
+            self._iam_principals,
+            self._dial_func,
+        )
+
     def _schedule_refresh(self, delay: int) -> asyncio.Task:
         """
         Schedule task to sleep and then perform refresh to get ConnectionInfo.
@@ -192,7 +235,7 @@ class RefreshAheadCache(ConnectionInfoCache):
             logger.debug(
                 f"['{self._conn_name}']: Connection info refresh"
                 " operation scheduled for "
-                f"{(datetime.now(timezone.utc) + timedelta(seconds=delay)).isoformat(timespec='seconds')} "
+                f"{(datetime.now(UTC) + timedelta(seconds=delay)).isoformat(timespec='seconds')} "
                 f"(now + {timedelta(seconds=delay)})"
             )
             self._next = self._schedule_refresh(delay)
@@ -223,3 +266,196 @@ class RefreshAheadCache(ConnectionInfoCache):
         tasks = asyncio.gather(self._current, self._next, return_exceptions=True)
         await asyncio.wait_for(tasks, timeout=2.0)
         self._closed = True
+
+
+def _append_iam_principal(
+    principals: list[tuple[str, str]], user: str, database: str
+) -> None:
+    """Appends (user, database) to principals if not already present."""
+    if not user:
+        return
+    if not database:
+        database = user
+    pair = (user, database)
+    if pair not in principals:
+        principals.append(pair)
+
+
+async def _probe_instance_connection(
+    conn_name: ConnectionName,
+    conn_info: ConnectionInfo,
+    enable_iam_auth: bool,
+    ip_type: IPTypes,
+    timeout: int,
+    iam_principals: list[tuple[str, str]],
+    dial_func: (
+        Callable[
+            [ConnectionInfo],
+            Awaitable[tuple[asyncio.StreamReader, asyncio.StreamWriter]],
+        ]
+        | None
+    ) = None,
+) -> None:
+    """Proactively probes the database to refresh IAM tokens on server-side MCP."""
+    if not enable_iam_auth:
+        return
+
+    # Proactive IAM refresh probe is PostgreSQL-only.
+    if not conn_info.database_version.startswith("POSTGRES"):
+        logger.debug(
+            f"['{conn_name}']: Proactive IAM token refresh probe skipped: instance engine '{conn_info.database_version}' is not PostgreSQL"
+        )
+        return
+
+    # Do not probe if no IAM principals have connected yet.
+    if not iam_principals:
+        logger.debug(
+            f"['{conn_name}']: Proactive IAM token refresh probe skipped: no IAM principals connected"
+        )
+        return
+
+    targets: list[str] = []
+    ssl_context = None
+    if dial_func is None:
+        if conn_name.domain_name:
+            targets.append(conn_name.domain_name)
+        if ip_type.value in conn_info.ip_addrs:
+            targets.extend(conn_info.ip_addrs[ip_type.value])
+
+        if not targets:
+            logger.debug(
+                f"['{conn_name}']: Proactive IAM token refresh probe skipped: no target IP addresses"
+            )
+            return
+
+        try:
+            ssl_context = await conn_info.create_ssl_context(enable_iam_auth)
+        except Exception as e:  # noqa: BLE001
+            logger.debug(
+                f"['{conn_name}']: Failed to create SSL context for probe: {e!s}"
+            )
+            return
+
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + float(max(1, timeout))
+
+    all_succeeded = True
+    for user, database in list(iam_principals):
+        remaining = deadline - loop.time()
+        if remaining <= 0:
+            all_succeeded = False
+            break
+
+        probed_principal = False
+        if dial_func is not None:
+            try:
+                reader, writer = await asyncio.wait_for(
+                    dial_func(conn_info),
+                    timeout=remaining,
+                )
+                probed_principal = await _send_startup_probe(
+                    reader,
+                    writer,
+                    user,
+                    database,
+                    max(0.0, deadline - loop.time()),
+                    conn_name,
+                )
+            except Exception as e:  # noqa: BLE001
+                logger.debug(
+                    f"['{conn_name}']: Probing IAM token refresh for ({user}, {database}) failed: {e!s}"
+                )
+        else:
+            for target in targets:
+                remaining = deadline - loop.time()
+                if remaining <= 0:
+                    break
+                try:
+                    logger.debug(
+                        f"['{conn_name}']: Probing IAM token refresh on {target}:{SERVER_PROXY_PORT}"
+                    )
+                    reader, writer = await asyncio.wait_for(
+                        asyncio.open_connection(
+                            host=target,
+                            port=SERVER_PROXY_PORT,
+                            ssl=ssl_context,
+                            server_hostname=(
+                                conn_name.domain_name
+                                if conn_name.domain_name
+                                else None
+                            ),
+                        ),
+                        timeout=remaining,
+                    )
+                    probed_principal = await _send_startup_probe(
+                        reader,
+                        writer,
+                        user,
+                        database,
+                        max(0.0, deadline - loop.time()),
+                        conn_name,
+                    )
+                    break
+                except Exception as e:  # noqa: BLE001
+                    logger.debug(
+                        f"['{conn_name}']: Probing IAM token refresh on {target}:{SERVER_PROXY_PORT} failed: {e!s}"
+                    )
+        if not probed_principal:
+            all_succeeded = False
+
+    if all_succeeded:
+        logger.debug(
+            f"['{conn_name}']: Proactive IAM token refresh probe successful"
+        )
+    else:
+        logger.debug(
+            f"['{conn_name}']: Proactive IAM token refresh probe encountered error across all targets"
+        )
+
+
+async def _send_startup_probe(
+    reader: asyncio.StreamReader,
+    writer: asyncio.StreamWriter,
+    user: str,
+    database: str,
+    timeout: float,
+    conn_name: ConnectionName,
+) -> bool:
+    """Sends a PostgreSQL StartupMessage probe and gracefully terminates."""
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + max(0.0, float(timeout))
+    try:
+        remaining = deadline - loop.time()
+        if remaining <= 0:
+            return False
+        writer.write(build_postgres_startup_packet(user, database))
+        await asyncio.wait_for(writer.drain(), timeout=remaining)
+        await asyncio.wait_for(
+            reader.read(1024),
+            timeout=max(0.0, deadline - loop.time()),
+        )
+        writer.write(build_terminate_packet())
+        await asyncio.wait_for(
+            writer.drain(),
+            timeout=max(0.0, deadline - loop.time()),
+        )
+        eof_timeout = min(max(0.0, deadline - loop.time()), 1.0)
+        if eof_timeout > 0:
+            try:
+                await asyncio.wait_for(reader.read(1024), timeout=eof_timeout)
+            except Exception as e:  # noqa: BLE001
+                logger.debug(
+                    f"['{conn_name}']: Ignoring error while waiting for server EOF: {e!s}"
+                )
+        return True
+    finally:
+        writer.close()
+        try:
+            await asyncio.wait_for(writer.wait_closed(), timeout=2.0)
+        except Exception as e:  # noqa: BLE001
+            logger.debug(
+                f"['{conn_name}']: Ignoring error while waiting for socket to close: {e!s}"
+            )
+            transport = getattr(writer, "transport", None)
+            if transport is not None and hasattr(transport, "abort"):
+                transport.abort()

google/cloud/sql/connector/lazy.py

@@ -15,15 +15,21 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Awaitable
+from collections.abc import Callable
 from datetime import datetime
 from datetime import timedelta
-from datetime import timezone
+from datetime import UTC
 import logging
 
 from google.cloud.sql.connector.client import CloudSQLClient
 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.instance import _append_iam_principal
+from google.cloud.sql.connector.instance import _probe_instance_connection
+from google.cloud.sql.connector.instance import DEFAULT_CONNECT_TIMEOUT
 from google.cloud.sql.connector.refresh_utils import _refresh_buffer
 
 logger = logging.getLogger(name=__name__)
@@ -44,6 +50,15 @@ class LazyRefreshCache(ConnectionInfoCache):
         client: CloudSQLClient,
         keys: asyncio.Future,
         enable_iam_auth: bool = False,
+        ip_type: IPTypes | str = IPTypes.PUBLIC,
+        timeout: int = DEFAULT_CONNECT_TIMEOUT,
+        dial_func: (
+            Callable[
+                [ConnectionInfo],
+                Awaitable[tuple[asyncio.StreamReader, asyncio.StreamWriter]],
+            ]
+            | None
+        ) = None,
     ) -> None:
         """Initializes a LazyRefreshCache instance.
 
@@ -56,16 +71,29 @@ class LazyRefreshCache(ConnectionInfoCache):
             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.
+            dial_func: Optional callable to dial probe connections.
         """
         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._dial_func = dial_func
         self._keys = keys
         self._client = client
         self._lock = asyncio.Lock()
+        self._iam_principals: list[tuple[str, str]] = []
         self._cached: ConnectionInfo | None = None
         self._needs_refresh = False
         self._closed = False
 
+    def record_principal(self, user: str, database: str) -> None:
+        """Records a PostgreSQL (user, database) pair observed on an IAM connection."""
+        _append_iam_principal(self._iam_principals, user, database)
+
     @property
     def conn_name(self) -> ConnectionName:
         return self._conn_name
@@ -82,6 +110,18 @@ class LazyRefreshCache(ConnectionInfoCache):
         async with self._lock:
             self._needs_refresh = True
 
+    async def _probe_connection(self, conn_info: ConnectionInfo) -> None:
+        """Proactively probes the database to refresh IAM tokens on server-side MCP."""
+        await _probe_instance_connection(
+            self._conn_name,
+            conn_info,
+            self._enable_iam_auth,
+            self._ip_type,
+            self._timeout,
+            self._iam_principals,
+            self._dial_func,
+        )
+
     async def connect_info(self) -> ConnectionInfo:
         """Retrieves ConnectionInfo instance for establishing a secure
         connection to the Cloud SQL instance.
@@ -93,7 +133,7 @@ class LazyRefreshCache(ConnectionInfoCache):
             if (
                 self._cached
                 and not self._needs_refresh
-                and datetime.now(timezone.utc)
+                and datetime.now(UTC)
                 < (self._cached.expiration - timedelta(seconds=_refresh_buffer))
             ):
                 logger.debug(
@@ -126,7 +166,11 @@ class LazyRefreshCache(ConnectionInfoCache):
             )
             self._cached = conn_info
             self._needs_refresh = False
-            return conn_info
+
+        if self._enable_iam_auth:
+            await self._probe_connection(conn_info)
+
+        return conn_info
 
     async def close(self) -> None:
         """Close is a no-op and provided purely for a consistent interface with

google/cloud/sql/connector/monitored_cache.py

@@ -15,9 +15,11 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Callable
 import logging
+import socket
 import ssl
-from typing import Any, Callable
+from typing import Any
 
 import aiohttp
 
@@ -27,6 +29,7 @@ from google.cloud.sql.connector.exceptions import CacheClosedError
 from google.cloud.sql.connector.exceptions import DnsLoopError
 from google.cloud.sql.connector.instance import RefreshAheadCache
 from google.cloud.sql.connector.lazy import LazyRefreshCache
+from google.cloud.sql.connector.postgres_connection import PostgresConnection
 from google.cloud.sql.connector.resolver import DefaultResolver
 from google.cloud.sql.connector.resolver import DnsResolver
 from google.cloud.sql.connector.resolver import is_instance_dns_name
@@ -44,7 +47,7 @@ class MonitoredCache(ConnectionInfoCache):
         self.resolver = resolver
         self.cache = cache
         self.domain_name_ticker: asyncio.Task | None = None
-        self.sockets: list[ssl.SSLSocket] = []
+        self.sockets: list[socket.socket | ssl.SSLSocket | PostgresConnection] = []
 
         # If domain name is configured for instance and failover period is set,
         # poll for DNS record changes.
@@ -77,11 +80,11 @@ class MonitoredCache(ConnectionInfoCache):
         list of sockets.
         """
         open_sockets = []
-        for socket in self.sockets:
+        for sock in self.sockets:
             # Check fileno for if socket is closed. Will return
             # -1 on failure, which will be used to signal socket closed.
-            if socket.fileno() != -1:
-                open_sockets.append(socket)
+            if sock.fileno() != -1:
+                open_sockets.append(sock)
         self.sockets = open_sockets
 
     async def _check_domain_name(self) -> None:
@@ -139,6 +142,9 @@ class MonitoredCache(ConnectionInfoCache):
             return
         return await self.cache.force_refresh()
 
+    def record_principal(self, user: str, database: str) -> None:
+        self.cache.record_principal(user, database)
+
     async def close(self) -> None:
         # Cancel domain name ticker task.
         if self.domain_name_ticker:
@@ -159,11 +165,11 @@ class MonitoredCache(ConnectionInfoCache):
         await self.cache.close()
 
         # Close any still open sockets
-        for socket in self.sockets:
+        for sock in self.sockets:
             # Check fileno for if socket is closed. Will return
             # -1 on failure, which will be used to signal socket closed.
-            if socket.fileno() != -1:
-                socket.close()
+            if sock.fileno() != -1:
+                sock.close()
 
 
 async def ticker(interval: int, function: Callable, *args: Any, **kwargs: Any) -> None: