diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py index 0626ee35..a45c2964 100644 --- a/google/cloud/sql/connector/connector.py +++ b/google/cloud/sql/connector/connector.py @@ -123,7 +123,7 @@ def __init__( refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND, resolver: type[DefaultResolver | DnsResolver] = DefaultResolver, failover_period: int = 30, - sql_data_endpoint: str = "sqladmin.googleapis.com", + sql_data_endpoint: str | None = None, sql_data_stream_timeout: int = 7200, resource_exhausted_cooldown_period: float = 5.0, ) -> None: @@ -180,7 +180,8 @@ def __init__( Must be used with `resolver=DnsResolver` to have any effect. Default: 30 - sql_data_endpoint (str): Endpoint host for SQL Data Service calls. + sql_data_endpoint (str | None): Endpoint host for SQL Data Service + calls. If not specified, derived from `sqladmin_api_endpoint`. Default: "sqladmin.googleapis.com". sql_data_stream_timeout (int): Timeout in seconds for the SQL Data @@ -276,7 +277,12 @@ def __init__( "configured the universe domain explicitly, `googleapis.com` " "is the default." ) - self._sql_data_endpoint = sql_data_endpoint + if sql_data_endpoint: + self._sql_data_endpoint = sql_data_endpoint + else: + self._sql_data_endpoint = ( + self._sqladmin_api_endpoint.split("://", 1)[-1].split("/", 1)[0] + ) self._sql_data_stream_timeout = sql_data_stream_timeout self._resource_exhausted_cooldown_period = ( resource_exhausted_cooldown_period @@ -284,6 +290,7 @@ def __init__( self._sql_data_fallback_cache: set[str] = set() self._sql_data_conn_state: dict[str, SqlDataConnState] = {} self._sqldata_clients: set[Any] = set() + self._sqldata_client: SqlDataClient | None = None @@ -499,16 +506,20 @@ async def connect_async( ) kwargs["user"] = formatted_user - sqldata_client = SqlDataClient( - endpoint=self._sql_data_endpoint, - credentials=self._credentials, - quota_project=self._quota_project, - timeout=self._sql_data_stream_timeout, - ) - self._sqldata_clients.add(sqldata_client) - sqldata_client._on_close_callbacks.append( - lambda: self._sqldata_clients.discard(sqldata_client) - ) + if self._sqldata_client is None: + sqldata_client = SqlDataClient( + endpoint=self._sql_data_endpoint, + credentials=self._credentials, + quota_project=self._quota_project, + timeout=self._sql_data_stream_timeout, + ) + self._sqldata_client = sqldata_client + self._sqldata_clients.add(sqldata_client) + sqldata_client._on_close_callbacks.append( + lambda: self._sqldata_clients.discard(sqldata_client) + ) + else: + sqldata_client = self._sqldata_client def on_resource_exhausted(err: Exception) -> None: backoff = state.record_exhausted( @@ -736,6 +747,7 @@ async def close_async(self) -> None: """Helper function to cancel the cache's tasks and close aiohttp.ClientSession.""" self._closed = True + self._sqldata_client = None if self._client: await self._client.close() await asyncio.gather( @@ -758,7 +770,7 @@ async def create_async_connector( refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND, resolver: type[DefaultResolver | DnsResolver] = DefaultResolver, failover_period: int = 30, - sql_data_endpoint: str = "sqladmin.googleapis.com", + sql_data_endpoint: str | None = None, sql_data_stream_timeout: int = 7200, resource_exhausted_cooldown_period: float = 5.0, ) -> Connector: @@ -818,7 +830,8 @@ async def create_async_connector( Must be used with `resolver=DnsResolver` to have any effect. Default: 30 - sql_data_endpoint (str): Endpoint host for SQL Data Service calls. + sql_data_endpoint (str | None): Endpoint host for SQL Data Service + calls. If not specified, derived from `sqladmin_api_endpoint`. Default: "sqladmin.googleapis.com". sql_data_stream_timeout (int): Timeout in seconds for the SQL Data diff --git a/google/cloud/sql/connector/psycopg.py b/google/cloud/sql/connector/psycopg.py index ad71c6b9..7d00d12f 100644 --- a/google/cloud/sql/connector/psycopg.py +++ b/google/cloud/sql/connector/psycopg.py @@ -28,8 +28,46 @@ logger = logging.getLogger(name=__name__) +def _proxy_sqldata(local: socket.socket, remote: socket.socket) -> None: + """Bidirectional thread-based proxy for queue-backed SqlDataSocket.""" + + def forward_local_to_remote() -> None: + try: + while data := local.recv(8192): + remote.sendall(data) + except OSError: + pass + finally: + remote.close() + + threading.Thread(target=forward_local_to_remote, daemon=True).start() + try: + while True: + try: + if not (data := remote.recv(8192)): + break + local.sendall(data) + except socket.timeout: + if not getattr(remote, "_first_read_done", False) or getattr( + remote, "_closed", False + ): + break + except OSError: + pass + finally: + try: + local.shutdown(socket.SHUT_RDWR) + except OSError: + pass + local.close() + remote.close() + + def _proxy(local: socket.socket, remote: "ssl.SSLSocket") -> None: """Single-threaded selectors-based proxy to avoid SSLSocket thread-safety issues.""" + if hasattr(remote, "_read_queue"): + _proxy_sqldata(local, remote) + return sel = selectors.DefaultSelector() sel.register(local, selectors.EVENT_READ, data="local") sel.register(remote, selectors.EVENT_READ, data="remote") diff --git a/google/cloud/sql/connector/sqldata_client.py b/google/cloud/sql/connector/sqldata_client.py index acf6287f..6361b022 100644 --- a/google/cloud/sql/connector/sqldata_client.py +++ b/google/cloud/sql/connector/sqldata_client.py @@ -14,6 +14,7 @@ from __future__ import annotations +import asyncio import errno import io import logging @@ -23,6 +24,7 @@ from typing import Any, Callable from google.api_core.client_options import ClientOptions +from google.api_core.exceptions import FailedPrecondition from google.api_core.exceptions import ResourceExhausted from google.auth.credentials import Credentials import grpc @@ -38,26 +40,44 @@ logger = logging.getLogger(__name__) -def is_resource_exhausted_error(err: Exception) -> bool: - """Checks whether an exception represents a RESOURCE_EXHAUSTED error.""" - if isinstance(err, ResourceExhausted): +def _is_grpc_status_error( + err: Exception, + exc_cls: type[Exception], + status_code: grpc.StatusCode, +) -> bool: + """Checks whether an exception or its cause matches the given gRPC status.""" + if isinstance(err, exc_cls): return True if isinstance(err, grpc.RpcError): try: - return err.code() == grpc.StatusCode.RESOURCE_EXHAUSTED + return err.code() == status_code except Exception: # noqa: BLE001, S110 pass if hasattr(err, "code") and callable(err.code): try: - return err.code() == grpc.StatusCode.RESOURCE_EXHAUSTED + return err.code() == status_code except Exception: # noqa: BLE001, S110 pass cause = getattr(err, "__cause__", None) or getattr(err, "__context__", None) if isinstance(cause, Exception) and cause is not err: - return is_resource_exhausted_error(cause) + return _is_grpc_status_error(cause, exc_cls, status_code) return False +def is_resource_exhausted_error(err: Exception) -> bool: + """Checks whether an exception represents a RESOURCE_EXHAUSTED error.""" + return _is_grpc_status_error( + err, ResourceExhausted, grpc.StatusCode.RESOURCE_EXHAUSTED + ) + + +def is_failed_precondition_error(err: Exception) -> bool: + """Checks whether an exception represents a FAILED_PRECONDITION error.""" + return _is_grpc_status_error( + err, FailedPrecondition, grpc.StatusCode.FAILED_PRECONDITION + ) + + class _RequestQueue: """Thread-safe queue iterator feeding requests to synchronous gRPC stream.""" @@ -167,10 +187,6 @@ def _reader_loop(self) -> None: for resp in self._response_stream: if self._closed: break - if not self._first_read_done: - self._first_read_done = True - if self._on_success: - self._on_success() raw_pb = ( sqladmin_v1beta4.StreamSqlDataResponse.pb(resp) if hasattr(sqladmin_v1beta4.StreamSqlDataResponse, "pb") @@ -182,6 +198,18 @@ def _reader_loop(self) -> None: if hasattr(raw_pb, "WhichOneof") else None ) + if which == "terminate_session" or ( + which is None + and hasattr(resp, "terminate_session") + and resp.terminate_session + ): + logger.debug("Received TerminateSession from server") + self._closed = True + break + if not self._first_read_done: + self._first_read_done = True + if self._on_success: + self._on_success() if which == "data" or (which is None and hasattr(resp, "data") and resp.data): data = resp.data.data if data: @@ -190,12 +218,6 @@ def _reader_loop(self) -> None: which is None and hasattr(resp, "session_metadata") and resp.session_metadata ): logger.debug("Received SessionMetadata") - elif which == "terminate_session" or ( - which is None and hasattr(resp, "terminate_session") and resp.terminate_session - ): - logger.debug("Received TerminateSession from server") - self._closed = True - break except Exception as e: # noqa: BLE001 if not self._closed: logger.debug(f"gRPC sync stream reader encountered: {e}") @@ -262,7 +284,7 @@ def recv(self, bufsize: int, flags: int = 0) -> bytes: if ( not self._first_read_done and self._fallback_fn is not None - and not is_resource_exhausted_error(self._error) + and is_failed_precondition_error(self._error) ): logger.info( f"SQL Data Service returned error before first read: {self._error}. " @@ -270,6 +292,7 @@ def recv(self, bufsize: int, flags: int = 0) -> bytes: ) try: self._direct_sock = self._fallback_fn() + self._direct_sock.settimeout(self._timeout) if self._write_buffer: self._direct_sock.sendall(bytes(self._write_buffer)) self._first_read_done = True @@ -446,6 +469,7 @@ def __init__( self._credentials = credentials self._quota_project = quota_project self._timeout = timeout + self._grpc_client: sqladmin_v1beta4.SqlDataServiceClient | None = None self._active_sockets: set[SqlDataSocket] = set() self._on_close_callbacks: list[Callable[[], None]] = [] @@ -491,6 +515,7 @@ def create_direct_sock() -> socket.socket: ssl_sock = ssl_context.wrap_socket( raw_sock, server_hostname=target_ip ) + ssl_sock.settimeout(None) return ssl_sock except Exception as e: # noqa: BLE001 logger.debug(f"Direct TLS connection to {target_ip} failed: {e}") @@ -508,16 +533,18 @@ def create_direct_sock() -> socket.socket: create_direct = await get_direct_socket_sync_factory() return create_direct() - # Connect synchronous gRPC stream via GAPIC client - endpoint = self._endpoint.removeprefix("https://").removeprefix("http://") - client_options = ClientOptions( - api_endpoint=endpoint, - quota_project_id=self._quota_project, - ) - grpc_client = sqladmin_v1beta4.SqlDataServiceClient( - credentials=self._credentials, - client_options=client_options, - ) + # Connect synchronous gRPC stream via shared GAPIC client + if self._grpc_client is None: + endpoint = self._endpoint.removeprefix("https://").removeprefix("http://") + client_options = ClientOptions( + api_endpoint=endpoint, + quota_project_id=self._quota_project, + ) + self._grpc_client = sqladmin_v1beta4.SqlDataServiceClient( + credentials=self._credentials, + client_options=client_options, + ) + grpc_client = self._grpc_client instance_id = ( f"projects/{project}/instances/{instance_connection_name.split(':')[-1]}" @@ -531,16 +558,19 @@ def create_direct_sock() -> socket.socket: ) ] - create_direct_sock_fn: Callable[[], socket.socket] | None = None - try: - create_direct_sock_fn = await get_direct_socket_sync_factory() - except Exception as e: # noqa: BLE001 - logger.debug(f"Could not prepare direct fallback factory: {e}") + loop = asyncio.get_running_loop() def fallback_fn() -> socket.socket: on_fallback(instance_connection_name) - if create_direct_sock_fn is None: - raise ValueError("No direct fallback connection factory available.") + try: + create_direct_sock_fn = asyncio.run_coroutine_threadsafe( + get_direct_socket_sync_factory(), loop + ).result(timeout=connect_timeout) + except Exception as e: + logger.debug(f"Could not prepare direct fallback factory: {e}") + raise ValueError( + "No direct fallback connection factory available." + ) from e return create_direct_sock_fn() try: @@ -570,7 +600,7 @@ def fallback_fn() -> socket.socket: sock = SqlDataSocket( request_queue=request_queue, response_stream=response_stream, - grpc_client=grpc_client, + grpc_client=None, timeout=connect_timeout, on_close=lambda: self._active_sockets.discard(sock), on_success=on_success, @@ -582,29 +612,23 @@ def fallback_fn() -> socket.socket: except Exception as e: logger.debug(f"Sync gRPC connection attempt failed: {e}") - try: - if hasattr(grpc_client, "transport") and hasattr(grpc_client.transport, "close"): - grpc_client.transport.close() - except Exception: # noqa: BLE001, S110 - pass - if is_resource_exhausted_error(e): if on_resource_exhausted: on_resource_exhausted(e) raise - # Fallback to direct TLS on connection failure + if not is_failed_precondition_error(e): + raise + + # Fallback to direct TLS on FAILED_PRECONDITION logger.info( f"SQL Data Service connection failed for {instance_connection_name}. " "Falling back to direct TLS connection." ) on_fallback(instance_connection_name) - if create_direct_sock_fn: - return create_direct_sock_fn() create_direct = await get_direct_socket_sync_factory() return create_direct() - async def close(self) -> None: """Closes all active sockets created by this client.""" for sock in list(self._active_sockets): @@ -613,6 +637,15 @@ async def close(self) -> None: except Exception: # noqa: BLE001, S110 pass self._active_sockets.clear() + if self._grpc_client is not None: + try: + if hasattr(self._grpc_client, "transport") and hasattr( + self._grpc_client.transport, "close" + ): + self._grpc_client.transport.close() + except Exception: # noqa: BLE001, S110 + pass + self._grpc_client = None for cb in self._on_close_callbacks: try: cb() diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py index 4e49643d..d526fb19 100644 --- a/tests/unit/test_connector.py +++ b/tests/unit/test_connector.py @@ -764,6 +764,16 @@ def test_Connector_init_sqladmin_endpoint(fake_credentials: Credentials) -> None sqladmin_api_endpoint="https://my-custom-endpoint.com", ) assert connector._sqladmin_api_endpoint == "https://my-custom-endpoint.com" + assert connector._sql_data_endpoint == "my-custom-endpoint.com" + connector.close() + + connector_override = Connector( + credentials=fake_credentials, + sqladmin_api_endpoint="https://my-custom-endpoint.com", + sql_data_endpoint="explicit-sqldata-endpoint.com", + ) + assert connector_override._sql_data_endpoint == "explicit-sqldata-endpoint.com" + connector_override.close() @pytest.mark.asyncio @@ -1516,13 +1526,13 @@ async def fake_connect(**kwargs): on_success = kwargs["on_success"] get_conn_info = kwargs["get_conn_info"] + if is_fallback_cached("test-project:test-region:test-instance"): + return mock_sock + # Call get_conn_info conn_info = await get_conn_info() assert conn_info is not None - # Verify initial fallback cache state - assert not is_fallback_cached("test-project:test-region:test-instance") - # Trigger fallback on_fallback("test-project:test-region:test-instance") assert is_fallback_cached("test-project:test-region:test-instance") @@ -1552,6 +1562,14 @@ async def fake_connect(**kwargs): db="my-db", ) assert conn is True + await connector.connect_async( + "test-project:test-region:test-instance", + "pg8000", + user="my-user", + password="my-pass", + db="my-db", + ) + assert len(connector._sqldata_clients) == 1 @@ -1567,5 +1585,3 @@ def test_Connector_close_handles_exception(fake_credentials: Credentials) -> Non mock_run.return_value = mock_future # Should log and not raise exception connector.close() - - diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py index e419d3cf..9443c7f4 100644 --- a/tests/unit/test_psycopg.py +++ b/tests/unit/test_psycopg.py @@ -13,9 +13,11 @@ # limitations under the License. import os +import queue import socket import ssl import threading +import time from typing import Any from unittest.mock import MagicMock from unittest.mock import patch @@ -813,3 +815,51 @@ def socket_side_effect(family, type, proto=0, fileno=None): mock_local_sock.close.assert_called_once() assert mock_remote_sock.close.call_count >= 1 + + +def test_proxy_sqldata_bidirectional() -> None: + """Test that _proxy forwards traffic for queue-backed SqlDataSocket.""" + local_client, local_server = socket.socketpair() + read_queue: queue.Queue = queue.Queue() + remote = MagicMock() + remote._read_queue = read_queue + remote._first_read_done = True + remote._closed = False + remote.sendall.side_effect = lambda payload: read_queue.put( + b"echo:" + payload + ) + + def fake_recv(bufsize: int) -> bytes: + try: + return read_queue.get(timeout=0.01) or b"" + except queue.Empty: + if remote._closed: + return b"" + raise socket.timeout("timed out") from None + + remote.recv.side_effect = fake_recv + remote.close.side_effect = lambda: ( + setattr(remote, "_closed", True), + read_queue.put(b""), + ) + + proxy_thread = threading.Thread( + target=_proxy, args=(local_server, remote), daemon=True + ) + proxy_thread.start() + local_client.sendall(b"SELECT 1") + assert local_client.recv(1024) == b"echo:SELECT 1" + time.sleep(0.03) + local_client.close() + proxy_thread.join(timeout=1.0) + + # Also cover OSError / initial timeout exit path + client_sock, server_sock = mockable_socketpair() + server_sock.recv = MagicMock(side_effect=OSError("err")) + remote._first_read_done = False + remote._closed = False + remote.recv.side_effect = [socket.timeout("handshake"), OSError("err")] + _proxy(server_sock, remote) + remote.recv.side_effect = OSError("err") + _proxy(server_sock, remote) + client_sock.close() diff --git a/tests/unit/test_sqldata_client.py b/tests/unit/test_sqldata_client.py index 4a8ac9c1..a5ad23ac 100644 --- a/tests/unit/test_sqldata_client.py +++ b/tests/unit/test_sqldata_client.py @@ -14,6 +14,7 @@ from __future__ import annotations +import asyncio import queue import socket import time @@ -405,7 +406,7 @@ def test_sqldata_socket_fallback_error_propagation(): def stream_err(): time.sleep(0.01) - raise MockRpcError(grpc.StatusCode.UNAVAILABLE) + raise MockRpcError(grpc.StatusCode.FAILED_PRECONDITION) yield mock_response_stream.__iter__.side_effect = stream_err @@ -520,7 +521,9 @@ async def test_sqldata_client_direct_socket_factory_no_ips(): mock_grpc_client = MagicMock() mock_grpc_client.transport.stream_sql_data = mock_grpc_client.stream_sql_data - mock_grpc_client.stream_sql_data.side_effect = MockRpcError(grpc.StatusCode.UNAVAILABLE) + mock_grpc_client.stream_sql_data.side_effect = MockRpcError( + grpc.StatusCode.FAILED_PRECONDITION + ) with ( patch( @@ -561,7 +564,9 @@ async def test_sqldata_client_direct_socket_factory_ip_retry_and_exhausted(): mock_grpc_client = MagicMock() mock_grpc_client.transport.stream_sql_data = mock_grpc_client.stream_sql_data - mock_grpc_client.stream_sql_data.side_effect = MockRpcError(grpc.StatusCode.UNAVAILABLE) + mock_grpc_client.stream_sql_data.side_effect = MockRpcError( + grpc.StatusCode.FAILED_PRECONDITION + ) # First IP fails, second IP succeeds with patch( @@ -733,7 +738,7 @@ async def test_sqldata_client_fallback(): mock_grpc_client = MagicMock() mock_grpc_client.transport.stream_sql_data = mock_grpc_client.stream_sql_data - rpc_err = MockRpcError(grpc.StatusCode.UNAVAILABLE) + rpc_err = MockRpcError(grpc.StatusCode.FAILED_PRECONDITION) mock_grpc_client.stream_sql_data.side_effect = rpc_err mock_conn_info = MagicMock() @@ -948,8 +953,11 @@ def stream_failing(*args, **kwargs): on_fallback=on_fallback, is_fallback_cached=lambda _: False, ) - resp = sock.recv(1024) + # Verify get_conn_info is not called eagerly before fallback is triggered + get_conn_info.assert_not_called() + resp = await asyncio.to_thread(sock.recv, 1024) assert resp == b"direct response" + get_conn_info.assert_called_once() on_fallback.assert_called_once_with("proj:region:inst") await client.close() @@ -992,7 +1000,7 @@ def stream_failing(*args, **kwargs): is_fallback_cached=lambda _: False, ) with pytest.raises(ValueError, match="No direct fallback connection factory available"): - sock.recv(1024) + await asyncio.to_thread(sock.recv, 1024) await client.close() @@ -1010,13 +1018,13 @@ async def test_sqldata_client_connect_close_transport_exception(): mock_conn_info = MagicMock() mock_conn_info.get_preferred_ips.return_value = ["1.2.3.4"] - mock_conn_info.create_ssl_context = AsyncMock( - side_effect=[Exception("init prep fail"), mock_ssl_ctx] - ) + mock_conn_info.create_ssl_context = AsyncMock(return_value=mock_ssl_ctx) get_conn_info = AsyncMock(return_value=mock_conn_info) mock_grpc_client = MagicMock() - mock_grpc_client.transport.stream_sql_data.side_effect = Exception("stream fail") + mock_grpc_client.transport.stream_sql_data.side_effect = MockRpcError( + grpc.StatusCode.FAILED_PRECONDITION + ) mock_grpc_client.transport.close.side_effect = Exception("close fail") with patch( @@ -1038,4 +1046,32 @@ async def test_sqldata_client_connect_close_transport_exception(): await client.close() - +@pytest.mark.asyncio +async def test_sqldata_client_unavailable_does_not_fallback(): + creds = MagicMock(spec=Credentials) + client = SqlDataClient(endpoint="sqladmin.googleapis.com", credentials=creds) + get_conn_info = AsyncMock() + on_fallback = MagicMock() + mock_grpc_client = MagicMock() + mock_grpc_client.transport.stream_sql_data.side_effect = MockRpcError( + grpc.StatusCode.UNAVAILABLE + ) + with ( + patch( + "google.cloud.sqladmin_v1beta4.SqlDataServiceClient", + return_value=mock_grpc_client, + ), + pytest.raises(MockRpcError), + ): + await client.connect( + instance_connection_name="proj:region:inst", + region="region", + project="proj", + get_conn_info=get_conn_info, + enable_iam_auth=False, + on_fallback=on_fallback, + is_fallback_cached=lambda _: False, + ) + on_fallback.assert_not_called() + get_conn_info.assert_not_called() + await client.close()