feat(sqldata): prefer PRIVATE, PSC, PUBLIC IP order on direct fallback

Author: hessjcg

google/cloud/sql/connector/sqldata_client.py

@@ -142,6 +142,20 @@ async def close(self) -> None:
             except Exception:  # noqa: BLE001, S110
                 pass
 
+    async def _open_direct_connection(
+        self,
+        target_ip: str,
+        port: int,
+        ssl_context: Any,
+        connect_timeout: float,
+    ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+        return await asyncio.wait_for(
+            asyncio.open_connection(
+                target_ip, port, ssl=ssl_context, server_hostname=target_ip
+            ),
+            timeout=connect_timeout,
+        )
+
     async def _handle_tunnel(
         self,
         client_reader: asyncio.StreamReader,
@@ -225,9 +239,9 @@ async def connect_grpc() -> tuple[grpc.aio.Channel, Any]:
         async def connect_direct() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
             logger.debug("Fallback triggered, fetching connection info...")
             conn_info = await get_conn_info()
-            # Find a fallback IP address, prioritizing PUBLIC for direct fallback connectivity
+            # Find a fallback IP address, prioritizing PRIVATE, PSC, PUBLIC
             targets: list[str] = []
-            for t in [IPTypes.PUBLIC, IPTypes.PSC, IPTypes.PRIVATE]:
+            for t in [IPTypes.PRIVATE, IPTypes.PSC, IPTypes.PUBLIC]:
                 try:
                     targets.extend(conn_info.get_preferred_ips(t))
                 except CloudSQLIPTypeError as e:
@@ -240,11 +254,11 @@ async def connect_direct() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
             for target_ip in targets:
                 logger.debug(f"Connecting directly to {target_ip}:{SERVER_PROXY_PORT}")
                 try:
-                    r, w = await asyncio.wait_for(
-                        asyncio.open_connection(
-                            target_ip, SERVER_PROXY_PORT, ssl=ssl_context, server_hostname=target_ip
-                        ),
-                        timeout=connect_timeout,
+                    r, w = await self._open_direct_connection(
+                        target_ip,
+                        SERVER_PROXY_PORT,
+                        ssl_context,
+                        connect_timeout,
                     )
                     self._active_writers.add(w)
                     return r, w

tests/unit/test_connector.py

@@ -1080,5 +1080,55 @@ async def mock_connect_tunnel(**kwargs):
             assert state.last_err is None
 
 
+@pytest.mark.asyncio
+async def test_sqldata_fallback_ip_order(fake_credentials: Credentials) -> None:
+    """Test that direct fallback queries IP addresses in PRIVATE, PSC, PUBLIC order."""
+    client = SqlDataClient(
+        endpoint="sqladmin.googleapis.com",
+        credentials=fake_credentials,
+    )
+    mock_conn_info = MagicMock()
+    queried_ip_types: list[IPTypes] = []
+
+    def mock_get_preferred_ips(ip_type: IPTypes):
+        queried_ip_types.append(ip_type)
+        if ip_type == IPTypes.PUBLIC:
+            return ["1.2.3.4"]
+        from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError
+
+        raise CloudSQLIPTypeError(f"{ip_type} not available")
+
+    mock_conn_info.get_preferred_ips.side_effect = mock_get_preferred_ips
+    mock_conn_info.create_ssl_context = AsyncMock(return_value=None)
+    get_conn_info = AsyncMock(return_value=mock_conn_info)
+
+    mock_reader = AsyncMock()
+    mock_reader.read = AsyncMock(return_value=b"")
+    mock_writer = MagicMock()
+    mock_writer.wait_closed = AsyncMock()
+    client._open_direct_connection = AsyncMock(
+        return_value=(mock_reader, mock_writer)
+    )
+
+    port = await client.connect_tunnel(
+        instance_connection_name="proj:reg:inst",
+        region="reg",
+        project="proj",
+        get_conn_info=get_conn_info,
+        enable_iam_auth=False,
+        on_fallback=MagicMock(),
+        is_fallback_cached=MagicMock(return_value=True),
+    )
+
+    # Trigger client connection to tunnel
+    _r, w = await asyncio.open_connection("127.0.0.1", port)
+    await asyncio.sleep(0.1)
+    w.close()
+    await w.wait_closed()
+
+    assert queried_ip_types == [IPTypes.PRIVATE, IPTypes.PSC, IPTypes.PUBLIC]
+    await client.close()
+
+