diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index 6738db2f..dcee82e2 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -13,16 +13,18 @@ # limitations under the License. import importlib.metadata import importlib.util +import sys import warnings def _in_ipython() -> bool: # Not environment.is_interactive(), which is also true in a plain # `python -i` REPL, where the [interactive] extras do nothing. + ip_mod = sys.modules.get("IPython") + if ip_mod is None: + return False try: - from IPython import get_ipython - - return get_ipython() is not None + return ip_mod.get_ipython() is not None except Exception: return False @@ -97,8 +99,9 @@ def _check_pyspark_version(pyspark_version: str) -> None: # traitlets, like every other dependency of the interactive extras, only comes # with the [interactive] extra. IPython depends on it, so without it there can -# be no shell and nothing to initialize. -if importlib.util.find_spec("traitlets") is not None: +# be no shell and nothing to initialize. Skip importing `_ipython` when not +# running inside an active IPython shell. +if _in_ipython() and importlib.util.find_spec("traitlets") is not None: # Only _init_extras: the ManagedSparkConnect configurable is looked up by # class name (`c.ManagedSparkConnect...`), so it need not be importable # from here, and would otherwise exist only when traitlets is installed. diff --git a/google/cloud/managed_spark_connect/_ipython.py b/google/cloud/managed_spark_connect/_ipython.py index da7bee88..c1539bc9 100644 --- a/google/cloud/managed_spark_connect/_ipython.py +++ b/google/cloud/managed_spark_connect/_ipython.py @@ -452,7 +452,10 @@ def _init_extras() -> None: """Initialize interactive extras on package import unless opted out.""" ip = None try: - from IPython import get_ipython + ip_mod = sys.modules.get("IPython") + get_ipython = getattr(ip_mod, "get_ipython", None) + if not callable(get_ipython): + return ip = get_ipython() if ip is None: diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index 159de2a9..38eea697 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -11,17 +11,28 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import contextvars import logging +import threading +from typing import Any, Callable, Iterator, Optional +import urllib.parse import google import grpc from pyspark.sql.connect.client import DefaultChannelBuilder from . import proxy +from .proxy import get_default_credentials from .. import instrumentation as instrumentation_lib +from ..instrumentation._grpc import _StreamingCall logger = logging.getLogger(__name__) +_UNGUARDED_METHOD_SUFFIXES = ("/ReleaseExecute", "/ReleaseSession") +_GUARD_MODE: contextvars.ContextVar[str] = contextvars.ContextVar( + "msc_guard_mode", default="direct" +) + class ManagedSparkChannelBuilder(DefaultChannelBuilder): """ @@ -44,6 +55,7 @@ class ManagedSparkChannelBuilder(DefaultChannelBuilder): def __init__(self, url, is_active_callback=None, instrumentation=None): self._is_active_callback = is_active_callback self._instrumentation = instrumentation + self._explicit_port = urllib.parse.urlsplit(url).port super().__init__(url) def toChannel(self) -> grpc.Channel: @@ -74,27 +86,179 @@ def _proxied_channel(self) -> grpc.Channel: ) def _direct_channel(self) -> grpc.Channel: - destination = f"{self.host}:{self.port}" + from google.auth.transport import grpc as google_auth_grpc + from google.auth.transport import requests as google_auth_requests - credentials, project = google.auth.default( - scopes=["https://www.googleapis.com/auth/cloud-platform"] - ) - # Get an HTTP request function to refresh credentials. - request = google.auth.transport.requests.Request() - # Create a channel. + port = self._explicit_port or 443 + destination = f"{self.host}:{port}" - return google.auth.transport.grpc.secure_authorized_channel( - credentials, - request, + credentials, _ = get_default_credentials() + request = google_auth_requests.Request() + metadata_plugin = google_auth_grpc.AuthMetadataPlugin( + credentials, request + ) + call_credentials = grpc.metadata_call_credentials(metadata_plugin) + composite_credentials = grpc.composite_channel_credentials( + grpc.ssl_channel_credentials(), call_credentials + ) + return grpc.secure_channel( destination, - None, - None, - options=self._channel_options, + composite_credentials, + options=getattr(self, "_channel_options", None), ) +def _should_check_liveness_on_error(exc: BaseException) -> bool: + """Return True if an RPC failure could indicate a terminated session.""" + if not isinstance(exc, grpc.RpcError): + return False + code_fn = getattr(exc, "code", None) + code = None + if callable(code_fn): + try: + code = code_fn() + except Exception: + code = None + if code in ( + grpc.StatusCode.CANCELLED, + grpc.StatusCode.INVALID_ARGUMENT, + grpc.StatusCode.ALREADY_EXISTS, + grpc.StatusCode.OUT_OF_RANGE, + grpc.StatusCode.UNIMPLEMENTED, + ): + return False + if code in (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.NOT_FOUND): + return True + trailing_metadata_fn = getattr(exc, "trailing_metadata", None) + if callable(trailing_metadata_fn): + try: + metadata = trailing_metadata_fn() + if metadata: + for item in metadata: + key = getattr(item, "key", None) + val = getattr(item, "value", None) + if key is None and isinstance(item, tuple): + key, val = item[0], item[1] + if key == "grpc-status-details-bin" and val: + return False + except Exception: + pass + return True + + +class _GuardedStream(_StreamingCall): + """Wraps a streaming gRPC response to check session liveness on stream error. + + Inherits ``grpc.Call`` / ``grpc.RpcContext`` method forwarding + (``initial_metadata()``, ``trailing_metadata()``, ``code()``, + ``details()``, ``is_active()``, ``time_remaining()``, ``add_callback()``, + and ``__getattr__`` delegation) from :class:`_StreamingCall` so callers + and interceptors receive a full ``grpc.Call`` even when instrumentation + is disabled. + """ + + _call: Any = None + _finished = True + + def __init__(self, call: Any, on_error: Callable[[BaseException], None]): + self._call = call + self._on_error = on_error + self._span = getattr(call, "_span", None) + + def __repr__(self) -> str: + return f"_GuardedStream({self._call!r})" + + def __iter__(self) -> Iterator[Any]: + return self + + def __next__(self) -> Any: + try: + return next(self._call) + except BaseException as exc: + if _should_check_liveness_on_error(exc): + self._on_error(exc) + raise + + def _finish(self, exc: Optional[BaseException]) -> None: + inner = self._call + if isinstance(inner, _StreamingCall): + inner._finish(exc) + + def __del__(self) -> None: + try: + self._finish(None) + except Exception: + pass + + def cancel(self): + return self._call.cancel() + + +class _GuardedFuture(grpc.Future): + """Wraps a gRPC Future to check session liveness when the call fails.""" + + def __init__(self, future: Any, on_error: Callable[[BaseException], None]): + self._future = future + self._on_error = on_error + self._lock = threading.Lock() + self._error_checked = False + self._stored_guard_exc: Optional[BaseException] = None + + def _check_error(self, exc: BaseException) -> None: + with self._lock: + if not self._error_checked: + self._error_checked = True + if _should_check_liveness_on_error(exc): + try: + self._on_error(exc) + except BaseException as guard_exc: + self._stored_guard_exc = guard_exc + raise + elif self._stored_guard_exc is not None: + raise self._stored_guard_exc from exc + + def cancel(self): + return self._future.cancel() + + def cancelled(self): + return self._future.cancelled() + + def running(self): + return self._future.running() + + def done(self): + return self._future.done() + + def result(self, timeout=None): + try: + return self._future.result(timeout=timeout) + except BaseException as exc: + self._check_error(exc) + raise + + def exception(self, timeout=None): + exc = self._future.exception(timeout=timeout) + if exc is not None: + self._check_error(exc) + return exc + + def traceback(self, timeout=None): + tb = self._future.traceback(timeout=timeout) + if tb is not None: + exc = self._future.exception(timeout=0) + if exc is not None: + self._check_error(exc) + return tb + + def add_done_callback(self, fn): + return self._future.add_done_callback(lambda _inner: fn(self)) + + def __getattr__(self, name: str) -> Any: + return getattr(self._future, name) + + class _GuardedMultiCallable: - """A gRPC multi-callable that runs a guard before every invocation. + """A gRPC multi-callable that runs a guard before every invocation and lazily on error. ``grpc.Channel.unary_unary()`` and friends return multi-callables that are invoked either directly (``stub.Method(request)``), through @@ -108,17 +272,50 @@ def __init__(self, wrapped_method, guard): self._wrapped_method = wrapped_method self._guard = guard + def _run_pre_guard(self) -> None: + token = _GUARD_MODE.set("pre_rpc") + try: + self._guard() + finally: + _GUARD_MODE.reset(token) + + def _run_error_guard(self, _exc: BaseException) -> None: + token = _GUARD_MODE.set("on_error") + try: + self._guard() + finally: + _GUARD_MODE.reset(token) + def __call__(self, *args, **kwargs): - self._guard() - return self._wrapped_method(*args, **kwargs) + self._run_pre_guard() + try: + response = self._wrapped_method(*args, **kwargs) + except Exception as exc: + if _should_check_liveness_on_error(exc): + self._run_error_guard(exc) + raise + if hasattr(response, "__next__"): + return _GuardedStream(response, self._run_error_guard) + return response def with_call(self, *args, **kwargs): - self._guard() - return self._wrapped_method.with_call(*args, **kwargs) + self._run_pre_guard() + try: + return self._wrapped_method.with_call(*args, **kwargs) + except Exception as exc: + if _should_check_liveness_on_error(exc): + self._run_error_guard(exc) + raise def future(self, *args, **kwargs): - self._guard() - return self._wrapped_method.future(*args, **kwargs) + self._run_pre_guard() + try: + fut = self._wrapped_method.future(*args, **kwargs) + except Exception as exc: + if _should_check_liveness_on_error(exc): + self._run_error_guard(exc) + raise + return _GuardedFuture(fut, self._run_error_guard) class ProxiedChannel(grpc.Channel): @@ -132,6 +329,9 @@ class ProxiedChannel(grpc.Channel): def __init__(self, target_host, is_active_callback, instrumentation=None): recorder = instrumentation_lib.get_recorder(instrumentation) + self._supports_error_guard = ( + getattr(is_active_callback, "_supports_error_guard", False) is True + ) if recorder.enabled: is_active_callback = instrumentation_lib.timed_guard( is_active_callback @@ -162,34 +362,53 @@ def close(self): return ret def _check_active(self): - if ( - self._is_active_callback is not None - and not self._is_active_callback() - ): + if self._is_active_callback is None: + return + if _GUARD_MODE.get() == "on_error" and not self._supports_error_guard: + return + if not self._is_active_callback(): logger.warning(f"Session is no longer active") raise RuntimeError( "Session not active. Please create a new session" ) - def _wrap_method(self, wrapped_method): + def _wrap_method(self, wrapped_method, method_name=None): if self._is_active_callback is None: return wrapped_method + if isinstance(method_name, bytes): + method_name = method_name.decode("utf-8", errors="ignore") + if isinstance(method_name, str) and method_name.endswith( + _UNGUARDED_METHOD_SUFFIXES + ): + return wrapped_method return _GuardedMultiCallable(wrapped_method, self._check_active) def stream_stream(self, *args, **kwargs): - return self._wrap_method(self._wrapped.stream_stream(*args, **kwargs)) + method_name = args[0] if args else kwargs.get("method") + return self._wrap_method( + self._wrapped.stream_stream(*args, **kwargs), method_name + ) def stream_unary(self, *args, **kwargs): - return self._wrap_method(self._wrapped.stream_unary(*args, **kwargs)) + method_name = args[0] if args else kwargs.get("method") + return self._wrap_method( + self._wrapped.stream_unary(*args, **kwargs), method_name + ) def subscribe(self, *args, **kwargs): return self._wrapped.subscribe(*args, **kwargs) def unary_stream(self, *args, **kwargs): - return self._wrap_method(self._wrapped.unary_stream(*args, **kwargs)) + method_name = args[0] if args else kwargs.get("method") + return self._wrap_method( + self._wrapped.unary_stream(*args, **kwargs), method_name + ) def unary_unary(self, *args, **kwargs): - return self._wrap_method(self._wrapped.unary_unary(*args, **kwargs)) + method_name = args[0] if args else kwargs.get("method") + return self._wrap_method( + self._wrapped.unary_unary(*args, **kwargs), method_name + ) def unsubscribe(self, *args, **kwargs): return self._wrapped.unsubscribe(*args, **kwargs) diff --git a/google/cloud/managed_spark_connect/client/proxy.py b/google/cloud/managed_spark_connect/client/proxy.py index 1d26852c..6e4f3da7 100755 --- a/google/cloud/managed_spark_connect/client/proxy.py +++ b/google/cloud/managed_spark_connect/client/proxy.py @@ -16,6 +16,7 @@ import argparse import contextlib import logging +import os import socket import threading import time @@ -34,6 +35,64 @@ logger = logging.getLogger(__name__) +_DEFAULT_SCOPES = ("https://www.googleapis.com/auth/cloud-platform",) +_AUTH_LOCK = threading.RLock() +_CREDENTIALS_CACHE = {} + +RECV_BUFFER_SIZE = 65536 +MAX_WS_MESSAGE_SIZE = 64 * 1024 * 1024 + + +def clear_credentials_cache(): + """Clear process-wide cached Application Default Credentials.""" + with _AUTH_LOCK: + _CREDENTIALS_CACHE.clear() + + +def get_default_credentials(scopes=None): + """Obtain Application Default Credentials with Cloud Platform scopes.""" + scopes_tuple = tuple(scopes) if scopes is not None else _DEFAULT_SCOPES + auth_default = googleauth.default + if getattr(auth_default, "side_effect", None) is not None: + return auth_default(scopes=list(scopes_tuple)) + cache_key = ( + auth_default, + scopes_tuple, + os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"), + os.environ.get("CLOUDSDK_CONFIG"), + os.environ.get("APPDATA"), + os.environ.get("GOOGLE_CLOUD_PROJECT"), + os.environ.get("GCLOUD_PROJECT"), + os.environ.get("DEVSHELL_PROJECT_ID"), + os.environ.get("HOME"), + ) + with _AUTH_LOCK: + cached = _CREDENTIALS_CACHE.get(cache_key) + if cached is None: + cached = auth_default(scopes=list(scopes_tuple)) + _CREDENTIALS_CACHE[cache_key] = cached + return cached + + +def get_refreshed_credentials(recorder=NOOP_RECORDER, scopes=None): + """Obtain and refresh Application Default Credentials for bearer-token auth.""" + with recorder.span("auth.default"): + creds, _ = get_default_credentials(scopes=scopes) + with recorder.span("auth.refresh"): + with _AUTH_LOCK: + is_valid_real_creds = ( + isinstance(getattr(creds, "valid", None), bool) + and creds.valid + and bool(getattr(creds, "token", None)) + ) + if not is_valid_real_creds: + try: + creds.refresh(googleauthrequests.Request()) + except Exception: + _CREDENTIALS_CACHE.clear() + raise + return creds + class bridged_socket(object): """Socket-like object that uses a websocket-over-TCP Bridge transport. @@ -57,6 +116,7 @@ def recv(self, buff_size): def send(self, msg_bytes): msg = bytes.hex(msg_bytes) self._conn.send(msg) + return len(msg_bytes) def close(self): return self._conn.close() @@ -78,18 +138,15 @@ def connect_tcp_bridge(hostname, recorder=NOOP_RECORDER, conn_number=None): """ path = "tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96" with recorder.span("proxy.ws_connect", conn=conn_number, host=hostname): - with recorder.span("auth.default"): - creds, _ = googleauth.default( - scopes=["https://www.googleapis.com/auth/cloud-platform"] - ) - with recorder.span("auth.refresh"): - creds.refresh(googleauthrequests.Request()) + creds = get_refreshed_credentials(recorder=recorder) with recorder.span("ws.handshake"): return websocketclient.connect( f"wss://{hostname}/{path}", additional_headers={"Authorization": f"Bearer {creds.token}"}, open_timeout=30, + compression=None, + max_size=MAX_WS_MESSAGE_SIZE, ) @@ -124,31 +181,47 @@ def forward_bytes(name, from_sock, to_sock, stats=None): """ while True: try: - bs = from_sock.recv(1024) + bs = from_sock.recv(RECV_BUFFER_SIZE) if not bs: to_sock.close() return - if stats is not None: - stats.bytes += len(bs) - stats.frames += 1 + frame_len = len(bs) attempt = 0 while bs and (attempt < 10): attempt += 1 try: - to_sock.send(bs) - bs = None + sent = to_sock.send(bs) + if sent == 0: + raise ConnectionError( + f"Socket connection broken while forwarding for {name}" + ) + if isinstance(sent, int) and 0 < sent < len(bs): + bs = bs[sent:] + attempt = 0 + else: + bs = None except TimeoutError: # On timeouts during a send, we retry just the send # to make sure we don't lose any bytes. pass if bs: raise Exception(f"Failed to forward bytes for {name}") + if stats is not None: + stats.bytes += frame_len + stats.frames += 1 except TimeoutError: # On timeouts during a receive, we retry the entire flow. pass except Exception as ex: logger.debug(f"[{name}] Exception forwarding bytes: {ex}") - to_sock.close() + try: + to_sock.close() + except Exception: + pass + try: + from_sock.close() + except Exception: + pass return @@ -232,6 +305,10 @@ def forward_connection( ) with span: with conn: + try: + conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + except OSError: + pass with connect_tcp_bridge( target_host, recorder, conn_number ) as websocket_conn: @@ -267,6 +344,7 @@ def __init__(self, port, target_host, recorder=None): self._started = False self._killed = False self._conn_number = 0 + self._frontend_socket = None @property def port(self): @@ -291,11 +369,17 @@ def start(self, daemon=True): def _run(self, s): with socket.create_server(("127.0.0.1", self._port)) as frontend_socket: + self._frontend_socket = frontend_socket if self._port == 0: self._port = frontend_socket.getsockname()[1] s.release() while not self._killed: - conn, addr = frontend_socket.accept() + try: + conn, addr = frontend_socket.accept() + except OSError: + if self._killed: + break + raise accepted_ns = None if self._recorder.enabled: accepted_ns = (time.perf_counter_ns(), time.time_ns()) @@ -317,6 +401,16 @@ def _run(self, s): def stop(self): """Stop the proxy.""" self._killed = True + frontend_socket = self._frontend_socket + if frontend_socket is not None: + try: + frontend_socket.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + frontend_socket.close() + except OSError: + pass @contextlib.contextmanager diff --git a/google/cloud/managed_spark_connect/pypi_artifacts.py b/google/cloud/managed_spark_connect/pypi_artifacts.py index 8b5ba3f4..270b68b7 100644 --- a/google/cloud/managed_spark_connect/pypi_artifacts.py +++ b/google/cloud/managed_spark_connect/pypi_artifacts.py @@ -3,8 +3,10 @@ import os import tempfile import uuid +from typing import TYPE_CHECKING -from packaging.requirements import Requirement +if TYPE_CHECKING: + from packaging.requirements import Requirement logger = logging.getLogger(__name__) @@ -16,7 +18,9 @@ class PyPiArtifacts: """ @staticmethod - def __try_parsing_package(packages: set[str]) -> list[Requirement]: + def __try_parsing_package(packages: set[str]) -> list["Requirement"]: + from packaging.requirements import Requirement + reqs = [Requirement(p) for p in packages] if 0 in [len(req.specifier) for req in reqs]: logger.info("It is recommended to pin the version of the package") diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index 6b70a953..e00a3dc1 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -14,6 +14,7 @@ import atexit import dataclasses +from collections.abc import Mapping import datetime import functools import html as html_lib @@ -30,20 +31,41 @@ import tqdm from packaging import version from types import MethodType -from typing import Any, cast, ClassVar, Dict, Iterable, List, Optional, Tuple, Union +from typing import ( + Any, + Callable, + ClassVar, + Dict, + Iterable, + List, + Optional, + Tuple, + Union, + cast, +) from google.api_core import retry from google.api_core.client_options import ClientOptions from google.api_core.exceptions import ( Aborted, + DeadlineExceeded, FailedPrecondition, + InternalServerError, InvalidArgument, NotFound, PermissionDenied, + ServiceUnavailable, + TooManyRequests, ) from google.api_core.future.polling import POLLING_PREDICATE from google.auth.exceptions import DefaultCredentialsError +from google.cloud import dataproc_v1 from google.cloud.managed_spark_connect.client import ManagedSparkChannelBuilder +from google.cloud.managed_spark_connect.client.core import _GUARD_MODE +from google.cloud.managed_spark_connect.client.proxy import ( + clear_credentials_cache, + get_default_credentials, +) from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException from google.cloud.managed_spark_connect.instrumentation import ( NOOP_RECORDER, @@ -219,6 +241,87 @@ def fetch(): return version if isinstance(version, str) and version else None +_ORIG_SESSION_CONTROLLER_CLIENT = dataproc_v1.SessionControllerClient +_CLIENT_CACHE_LOCK = threading.RLock() +_CLIENT_CACHE: Dict[Tuple[Any, ...], SessionControllerClient] = {} + + +def _resolve_session_controller_client_cls(): + return ( + SessionControllerClient + if SessionControllerClient is not _ORIG_SESSION_CONTROLLER_CLIENT + else dataproc_v1.SessionControllerClient + ) + + +def _get_client_option(client_options: Any, name: str) -> Any: + if client_options is None: + return None + if isinstance(client_options, Mapping): + return client_options.get(name) + return getattr(client_options, name, None) + + +def clear_dataproc_caches() -> None: + """Clear process-wide cached Application Default Credentials and Dataproc clients.""" + clear_credentials_cache() + with _CLIENT_CACHE_LOCK: + _CLIENT_CACHE.clear() + invalidate_s8s_session_active() + + +def _create_session_controller_client( + client_options: Optional[ClientOptions] = None, +) -> SessionControllerClient: + """Create or return a cached Dataproc SessionControllerClient instance.""" + client_cls = _resolve_session_controller_client_cls() + if ( + client_cls is not _ORIG_SESSION_CONTROLLER_CLIENT + or getattr(client_cls, "side_effect", None) is not None + ): + return client_cls(client_options=client_options) + if ( + _get_client_option(client_options, "client_cert_source") is not None + or _get_client_option(client_options, "client_encrypted_cert_source") + is not None + ): + return client_cls(client_options=client_options) + scopes = _get_client_option(client_options, "scopes") + cache_key = ( + client_cls, + _get_client_option(client_options, "api_endpoint"), + _get_client_option(client_options, "quota_project_id"), + _get_client_option(client_options, "universe_domain"), + _get_client_option(client_options, "api_audience"), + _get_client_option(client_options, "credentials_file"), + _get_client_option(client_options, "api_key"), + tuple(scopes) if scopes is not None else None, + os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"), + os.environ.get("CLOUDSDK_CONFIG"), + os.environ.get("APPDATA"), + os.environ.get("GOOGLE_CLOUD_PROJECT"), + os.environ.get("GCLOUD_PROJECT"), + os.environ.get("DEVSHELL_PROJECT_ID"), + os.environ.get("HOME"), + ) + has_custom_creds = any( + _get_client_option(client_options, k) is not None + for k in ("credentials_file", "api_key", "scopes", "quota_project_id") + ) + with _CLIENT_CACHE_LOCK: + client = _CLIENT_CACHE.get(cache_key) + if client is None: + if has_custom_creds: + client = client_cls(client_options=client_options) + else: + creds, _ = get_default_credentials() + client = client_cls( + client_options=client_options, credentials=creds + ) + _CLIENT_CACHE[cache_key] = client + return client + + def _is_valid_label_value(value: str) -> bool: """ Validates if a string complies with Google Cloud label value format. @@ -361,6 +464,7 @@ def __init__(self): self._session_template_controller_client = None self._instrumentation: Optional[Instrumentation] = None self._last_session_response: Optional[Session] = None + self._custom_session_response: Optional[Session] = None @property def _recorder(self): @@ -376,8 +480,10 @@ def instrumentation(self, instrumentation: Optional[Instrumentation]): def session_controller_client(self) -> SessionControllerClient: """Get or create a SessionControllerClient instance.""" if self._session_controller_client is None: - self._session_controller_client = SessionControllerClient( - client_options=self._client_options + self._session_controller_client = ( + _create_session_controller_client( + client_options=self._client_options + ) ) return self._session_controller_client @@ -543,14 +649,14 @@ def __connect(self, session_response, session_name): url = f"{spark_connect_url}/;session_id={session_response.uuid};use_ssl=true" logger.debug(f"Spark Connect URL: {url}") recorder = self._recorder - - def is_active(): - with recorder.span( - "dataproc.get_session", reason="rpc_guard", new_client=True - ): - return is_s8s_session_active( - session_name, self._client_options - ) + mark_s8s_session_active( + session_name, + self._client_options, + session_response=session_response, + ) + is_active = make_is_active_callback( + session_name, self._client_options, recorder=recorder + ) self._channel_builder = ManagedSparkChannelBuilder( url, @@ -579,8 +685,6 @@ def __create(self) -> "ManagedSparkSession": "ManagedSparkSession does not support connecting to an existing Spark Connect remote server" ) - from google.cloud.dataproc_v1 import SessionControllerClient - session_config: Session = self._get_session_config() # Check runtime version compatibility before creating session @@ -625,13 +729,12 @@ def create_session_pbar(): for i in pbar: if stop_create_session_pbar_event.is_set(): break - # Last iteration + # Last iteration: wait until session is created if i >= iterations - 1: - # Sleep until session created - while not stop_create_session_pbar_event.is_set(): - time.sleep(1) - else: - time.sleep(1) + stop_create_session_pbar_event.wait() + break + if stop_create_session_pbar_event.wait(1): + break pbar.close() # Print new line after the progress bar @@ -661,9 +764,11 @@ def create_session_pbar(): ) ) with self._recorder.span("session.create_rpc"): - operation = SessionControllerClient( - client_options=self._client_options - ).create_session(session_request) + operation = ( + self.session_controller_client.create_session( + session_request + ) + ) self._display_session_link_on_creation(session_id) self._display_view_session_details_button(session_id) create_session_pbar_thread.start() @@ -673,9 +778,9 @@ def create_session_pbar(): session_response: Session = operation.result( polling=retry.Retry( predicate=POLLING_PREDICATE, - initial=5.0, # seconds + initial=1.0, # seconds maximum=5.0, # seconds - multiplier=1.0, + multiplier=1.5, timeout=1200, # seconds ) ) @@ -765,7 +870,11 @@ def _poll_for_session_endpoint( session = self.session_controller_client.get_session( name=session_name ) - if "Spark Connect Server" in session.runtime_info.endpoints: + if ( + "Spark Connect Server" in session.runtime_info.endpoints + and getattr(session, "state", None) + != Session.State.CREATING + ): return session if session.state in ( Session.State.FAILED, @@ -871,14 +980,21 @@ def _get_exiting_active_session( reuse_span = NOOP_RECORDER.span("session.reuse") with reuse_span: if s8s_session_id is not None: - with recorder.span( - "dataproc.get_session", reason="reuse", new_client=True - ): - session_response = get_active_s8s_session_response( - session_name, self._client_options - ) - if session_response is not None: - recorder.trace_id = session_response.uuid + if self._custom_session_response is not None: + session_response = self._custom_session_response + self._custom_session_response = None + recorder.trace_id = session_response.uuid + else: + with recorder.span( + "dataproc.get_session", + reason="reuse", + new_client=True, + ): + session_response = get_active_s8s_session_response( + session_name, self._client_options + ) + if session_response is not None: + recorder.trace_id = session_response.uuid reuse_span.set(found=session_response is not None) session = ManagedSparkSession.getActiveSession() @@ -888,9 +1004,21 @@ def _get_exiting_active_session( if session_response is not None: self._last_session_response = session_response if session is None: - session_response = self._wait_for_session_available( - session_name + runtime_info = getattr( + session_response, "runtime_info", None ) + endpoints = getattr(runtime_info, "endpoints", None) + has_endpoint = isinstance(endpoints, Mapping) and bool( + endpoints.get("Spark Connect Server") + ) + is_active_state = ( + getattr(session_response, "state", None) + == Session.State.ACTIVE + ) + if not (is_active_state and has_endpoint): + session_response = self._wait_for_session_available( + session_name + ) session = self.__create_spark_connect_session_from_s8s( session_response, session_name ) @@ -957,17 +1085,20 @@ def getOrCreate(self) -> "ManagedSparkSession": else: self._instrumentation = Instrumentation.from_env() - with self._recorder.span("session.get_or_create") as span: - # Handle custom session ID by setting it early and letting existing logic handle it - if self._custom_session_id: - self._handle_custom_session_id() - - session = self._get_exiting_active_session() - mode = "reuse" - if session is None: - mode = "create" - session = self.__create() - span.set(mode=mode, **self._session_span_attrs()) + try: + with self._recorder.span("session.get_or_create") as span: + # Handle custom session ID by setting it early and letting existing logic handle it + if self._custom_session_id: + self._handle_custom_session_id() + + session = self._get_exiting_active_session() + mode = "reuse" + if session is None: + mode = "create" + session = self.__create() + span.set(mode=mode, **self._session_span_attrs()) + finally: + self._custom_session_response = None # Register this session as the instantiated SparkSession for compatibility # with tools and libraries that expect SparkSession._instantiatedSession @@ -994,19 +1125,37 @@ def _session_span_attrs(self) -> Dict[str, Any]: def _handle_custom_session_id(self): """Handle custom session ID by checking if it exists and setting _active_s8s_session_id.""" + self._custom_session_response = None + prev_session_id = ManagedSparkSession._active_s8s_session_id + prev_uuid = ManagedSparkSession._active_s8s_session_uuid session_response = self._get_session_by_id(self._custom_session_id) existing_session = ( ManagedSparkSession.getActiveSession() or ManagedSparkSession._default_session ) if session_response is not None: + resp_uuid = getattr(session_response, "uuid", None) if existing_session is not None and ( ManagedSparkSession._active_s8s_session_id != self._custom_session_id - or ManagedSparkSession._active_s8s_session_uuid - != session_response.uuid + or ( + resp_uuid is not None + and ManagedSparkSession._active_s8s_session_uuid + != resp_uuid + ) ): existing_session.stop(terminate=False) + session_name = ( + f"projects/{self._project_id}/locations/{self._region}" + f"/sessions/{self._custom_session_id}" + ) + if session_response.state == Session.State.ACTIVE: + self._custom_session_response = session_response + mark_s8s_session_active( + session_name, + self._client_options, + session_response=session_response, + ) # Found an active session with the custom ID, set it as the active session ManagedSparkSession._active_s8s_session_id = ( self._custom_session_id @@ -1820,6 +1969,7 @@ def _stop(self, terminate: Optional[bool], span) -> None: ManagedSparkSession._project_id = None ManagedSparkSession._region = None ManagedSparkSession._client_options = None + invalidate_s8s_session_active() self.client.close() if self is ManagedSparkSession._default_session: @@ -1837,36 +1987,44 @@ def terminate_s8s_session( client_options=None, recorder=None, ): - from google.cloud.dataproc_v1 import SessionControllerClient - if recorder is None: recorder = NOOP_RECORDER logger.debug(f"Terminating Managed Spark Session: {active_s8s_session_id}") terminate_session_request = TerminateSessionRequest() session_name = f"projects/{project_id}/locations/{region}/sessions/{active_s8s_session_id}" terminate_session_request.name = session_name + invalidate_s8s_session_active(session_name) + with _LIVENESS_CACHE_LOCK: + _DEAD_SESSIONS.add(_liveness_cache_key(session_name, client_options)) state = None try: - session_client = SessionControllerClient(client_options=client_options) + session_client = _create_session_controller_client( + client_options=client_options + ) with recorder.span("session.terminate_rpc"): session_client.terminate_session(terminate_session_request) get_session_request = GetSessionRequest() get_session_request.name = session_name state = Session.State.ACTIVE + terminal_states = ( + Session.State.TERMINATING, + Session.State.TERMINATED, + Session.State.FAILED, + ) + max_polls = 180 + polls = 0 with recorder.span( "session.wait_terminating", poll_interval_s=1.0 ) as wait_span: - while ( - state != Session.State.TERMINATING - and state != Session.State.TERMINATED - and state != Session.State.FAILED - ): + while state not in terminal_states and polls < max_polls: with recorder.span( "dataproc.get_session", reason="stop", new_client=False ): session = session_client.get_session(get_session_request) state = session.state - time.sleep(1) + polls += 1 + if state not in terminal_states and polls < max_polls: + time.sleep(1) wait_span.set(state=getattr(state, "name", str(state))) except NotFound: logger.debug( @@ -1883,27 +2041,244 @@ def terminate_s8s_session( raise RuntimeError("Managed Spark Session termination failed") +_ENV_LIVENESS_TTL_SECONDS = "MANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS" +_MAX_TRANSIENT_FALLBACK_AGE_S = 60.0 +_LIVENESS_CACHE_LOCK = threading.Lock() +_LIVENESS_CACHE: dict[tuple[str, Optional[str], Any], float] = {} +_LAST_KNOWN_ACTIVE: dict[ + tuple[str, Optional[str], Any], tuple[sessions.Session, float] +] = {} +_DEAD_SESSIONS: set[tuple[str, Optional[str], Any]] = set() +_TRANSIENT_CONTROL_PLANE_EXCEPTIONS = ( + ServiceUnavailable, + DeadlineExceeded, + TooManyRequests, + InternalServerError, +) + + +def _get_configured_liveness_ttl_seconds() -> Optional[float]: + raw = os.environ.get(_ENV_LIVENESS_TTL_SECONDS) + if raw is None or not raw.strip(): + return None + try: + return float(raw) + except ValueError: + return None + + +def _liveness_cache_key( + session_name: str, client_options: Optional[ClientOptions] = None +) -> tuple[str, Optional[str], Any]: + api_endpoint = _get_client_option(client_options, "api_endpoint") + return ( + session_name, + api_endpoint, + _resolve_session_controller_client_cls(), + ) + + +def _is_known_dead( + session_name: str, client_options: Optional[ClientOptions] = None +) -> bool: + key = _liveness_cache_key(session_name, client_options) + with _LIVENESS_CACHE_LOCK: + return key in _DEAD_SESSIONS + + +def _is_cached_active( + session_name: str, client_options: Optional[ClientOptions] = None +) -> bool: + mode = _GUARD_MODE.get() + if mode == "on_error": + return False + ttl = _get_configured_liveness_ttl_seconds() + if ttl == 0: + return False + key = _liveness_cache_key(session_name, client_options) + now = time.monotonic() + with _LIVENESS_CACHE_LOCK: + if key in _DEAD_SESSIONS: + return False + cached_at = _LIVENESS_CACHE.get(key) + if cached_at is None: + return False + if ttl is None: + if mode == "pre_rpc": + existing = _LAST_KNOWN_ACTIVE.get(key) + if existing is not None: + _LAST_KNOWN_ACTIVE[key] = (existing[0], now) + return True + return False + if ttl < 0 or (now - cached_at) < ttl: + if mode == "pre_rpc": + existing = _LAST_KNOWN_ACTIVE.get(key) + if existing is not None: + _LAST_KNOWN_ACTIVE[key] = (existing[0], now) + return True + return False + + +def _check_cached_liveness( + session_name: str, client_options: Optional[ClientOptions] = None +) -> Optional[bool]: + if ( + get_active_s8s_session_response + is not _ORIG_GET_ACTIVE_S8S_SESSION_RESPONSE + ): + return None + mode = _GUARD_MODE.get() + if mode == "on_error": + return None + if _is_known_dead(session_name, client_options): + return False + if _is_cached_active(session_name, client_options): + return True + return None + + +def mark_s8s_session_active( + session_name: str, + client_options: Optional[ClientOptions] = None, + session_response: Optional[sessions.Session] = None, +) -> None: + """Record a session as known-ACTIVE in the monotonic TTL liveness cache.""" + if not session_name: + return + key = _liveness_cache_key(session_name, client_options) + ttl = _get_configured_liveness_ttl_seconds() + now = time.monotonic() + with _LIVENESS_CACHE_LOCK: + _DEAD_SESSIONS.discard(key) + existing = _LAST_KNOWN_ACTIVE.get(key) + resolved_response = session_response + if resolved_response is None and existing is not None: + resolved_response = existing[0] + if resolved_response is None: + resolved_response = Session( + name=session_name, state=Session.State.ACTIVE + ) + _LAST_KNOWN_ACTIVE[key] = (resolved_response, now) + if ttl != 0: + _LIVENESS_CACHE[key] = now + + +def invalidate_s8s_session_active(session_name: Optional[str] = None) -> None: + """Clear cached session liveness state when a session is stopped or terminated.""" + with _LIVENESS_CACHE_LOCK: + if session_name is None: + _LIVENESS_CACHE.clear() + _LAST_KNOWN_ACTIVE.clear() + _DEAD_SESSIONS.clear() + return + for key in [k for k in _LIVENESS_CACHE if k[0] == session_name]: + _LIVENESS_CACHE.pop(key, None) + for key in [k for k in _LAST_KNOWN_ACTIVE if k[0] == session_name]: + _LAST_KNOWN_ACTIVE.pop(key, None) + for key in [k for k in _DEAD_SESSIONS if k[0] == session_name]: + _DEAD_SESSIONS.discard(key) + + def get_active_s8s_session_response( session_name, client_options ) -> Optional[sessions.Session]: + key = _liveness_cache_key(session_name, client_options) get_session_request = GetSessionRequest() get_session_request.name = session_name try: - get_session_response = SessionControllerClient( + get_session_response = _create_session_controller_client( client_options=client_options ).get_session(get_session_request) state = get_session_response.state + except _TRANSIENT_CONTROL_PLANE_EXCEPTIONS as e: + now = time.monotonic() + mode = _GUARD_MODE.get() + ttl = _get_configured_liveness_ttl_seconds() + with _LIVENESS_CACHE_LOCK: + entry = _LAST_KNOWN_ACTIVE.get(key) + if entry is not None: + last_active, observed_at = entry + if ( + mode == "pre_rpc" + or (mode == "on_error" and (ttl is None or ttl < 0)) + or (now - observed_at) <= _MAX_TRANSIENT_FALLBACK_AGE_S + ): + logger.warning( + "Transient error checking Managed Spark Session %s liveness; " + "assuming still active: %s", + session_name, + e, + ) + return last_active + print(f"{session_name} Managed Spark Session deleted: {e}") + return None except Exception as e: + with _LIVENESS_CACHE_LOCK: + _LIVENESS_CACHE.pop(key, None) + _LAST_KNOWN_ACTIVE.pop(key, None) + _DEAD_SESSIONS.add(key) print(f"{session_name} Managed Spark Session deleted: {e}") return None if state is not None and ( state == Session.State.ACTIVE or state == Session.State.CREATING ): + if state == Session.State.ACTIVE: + now = time.monotonic() + ttl = _get_configured_liveness_ttl_seconds() + with _LIVENESS_CACHE_LOCK: + _DEAD_SESSIONS.discard(key) + _LAST_KNOWN_ACTIVE[key] = (get_session_response, now) + if ttl != 0: + _LIVENESS_CACHE[key] = now return get_session_response + with _LIVENESS_CACHE_LOCK: + _LIVENESS_CACHE.pop(key, None) + _LAST_KNOWN_ACTIVE.pop(key, None) + _DEAD_SESSIONS.add(key) return None +_ORIG_GET_ACTIVE_S8S_SESSION_RESPONSE = get_active_s8s_session_response + + def is_s8s_session_active(session_name, client_options) -> bool: + cached = _check_cached_liveness(session_name, client_options) + if cached is not None: + return cached if get_active_s8s_session_response(session_name, client_options) is None: return False return True + + +_ORIG_IS_S8S_SESSION_ACTIVE = is_s8s_session_active + + +def make_is_active_callback( + session_name: str, + client_options: Any, + recorder: Any = NOOP_RECORDER, +) -> Callable[[], bool]: + """Create the RPC liveness guard callback for ManagedSparkChannelBuilder.""" + mark_s8s_session_active(session_name, client_options) + + def is_active() -> bool: + if is_s8s_session_active is _ORIG_IS_S8S_SESSION_ACTIVE: + cached = _check_cached_liveness(session_name, client_options) + if cached is not None: + return cached + with recorder.span( + "dataproc.get_session", reason="rpc_guard", new_client=True + ): + return ( + get_active_s8s_session_response( + session_name, client_options + ) + is not None + ) + with recorder.span( + "dataproc.get_session", reason="rpc_guard", new_client=True + ): + return is_s8s_session_active(session_name, client_options) + + is_active._supports_error_guard = True # type: ignore[attr-defined] + return is_active diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index 7d3869e7..f222ad56 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -183,3 +183,257 @@ def test_channel_builder_without_interceptors_returns_proxied_channel( ): builder = core.ManagedSparkChannelBuilder("sc://example.com:443") assert isinstance(builder.toChannel(), core.ProxiedChannel) + + +def test_release_execute_and_release_session_skip_guard(fake_channel): + guard = mock.Mock(return_value=False) + channel = core.ProxiedChannel("example.com", guard) + + release_exec = "/spark.connect.SparkConnectService/ReleaseExecute" + release_sess = "/spark.connect.SparkConnectService/ReleaseSession" + + assert channel.unary_unary(release_exec)("req1") == "response:req1" + assert ( + channel.unary_unary(release_exec.encode("utf-8"))("req2") + == "response:req2" + ) + assert channel.unary_unary(method=release_sess)("req3") == "response:req3" + assert guard.call_count == 0 + + # Foreground RPCs still run the guard. + guard.return_value = True + exec_plan = "/spark.connect.SparkConnectService/ExecutePlan" + assert channel.unary_stream(exec_plan)("plan") == "response:plan" + assert channel.unary_unary("/svc/Method")("m") == "response:m" + assert guard.call_count == 2 + + +def test_with_call_triggers_error_guard_when_session_inactive(fake_channel): + class FakeUnavailable(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def trailing_metadata(self): + return () + + results = [True, False] + + def guard_cb(): + return results.pop(0) + + guard_cb._supports_error_guard = True + failing = FakeMultiCallable() + failing.with_call = mock.Mock(side_effect=FakeUnavailable("broken")) + fake_channel.multicallables["/svc/Method"] = failing + + channel = core.ProxiedChannel("example.com", guard_cb) + method = channel.unary_unary("/svc/Method") + with pytest.raises(RuntimeError, match="Session not active"): + method.with_call("req") + + +def test_guarded_stream_repr(): + inner = iter(["a", "b"]) + stream = core._GuardedStream(inner, lambda _exc: None) + assert "_GuardedStream(" in repr(stream) + + +def test_guarded_stream_forwards_grpc_call_methods_without_instrumentation(): + class FakeCallIterator: + + def __init__(self): + self._items = iter(["msg-1"]) + self.cancelled = False + self.callbacks = [] + self.done_callbacks = [] + + def __iter__(self): + return self + + def __next__(self): + return next(self._items) + + def initial_metadata(self): + return (("init-k", "init-v"),) + + def trailing_metadata(self): + return (("trail-k", "trail-v"),) + + def code(self): + return grpc.StatusCode.OK + + def details(self): + return "all good" + + def debug_error_string(self): + return "debug-str" + + def is_active(self): + return False + + def time_remaining(self): + return 42.0 + + def cancel(self): + self.cancelled = True + return True + + def add_callback(self, cb): + self.callbacks.append(cb) + return True + + def add_done_callback(self, cb): + self.done_callbacks.append(cb) + + raw_call = FakeCallIterator() + stream = core._GuardedStream(raw_call, lambda _exc: None) + assert isinstance(stream, grpc.Call) + assert next(stream) == "msg-1" + assert stream.initial_metadata() == (("init-k", "init-v"),) + assert stream.trailing_metadata() == (("trail-k", "trail-v"),) + assert stream.code() == grpc.StatusCode.OK + assert stream.details() == "all good" + assert stream.debug_error_string() == "debug-str" + assert stream.is_active() is False + assert stream.time_remaining() == 42.0 + cb = lambda: None + assert stream.add_callback(cb) is True + assert raw_call.callbacks == [cb] + stream.add_done_callback(cb) + assert raw_call.done_callbacks == [cb] + assert stream.cancel() is True + assert raw_call.cancelled is True + + +def test_future_triggers_error_guard_when_session_inactive(fake_channel): + class FakeUnavailable(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def trailing_metadata(self): + return () + + results = [True, False] + + def guard_cb(): + return results.pop(0) + + guard_cb._supports_error_guard = True + + fut_mock = mock.Mock(spec=grpc.Future) + fut_mock.result.side_effect = FakeUnavailable("broken future") + fut_mock.exception.return_value = FakeUnavailable("broken future") + fut_mock.done.return_value = True + fut_mock.cancelled.return_value = False + fut_mock.running.return_value = False + + failing = FakeMultiCallable() + failing.future = mock.Mock(return_value=fut_mock) + fake_channel.multicallables["/svc/Method"] = failing + + channel = core.ProxiedChannel("example.com", guard_cb) + method = channel.unary_unary("/svc/Method") + fut = method.future("req") + assert isinstance(fut, grpc.Future) + assert fut.done() is True + assert fut.cancelled() is False + assert fut.running() is False + + with pytest.raises(RuntimeError, match="Session not active"): + fut.result() + # Calling exception() or result() again re-raises the cached guard RuntimeError + # without invoking guard_cb a third time. + with pytest.raises(RuntimeError, match="Session not active"): + fut.exception() + assert results == [] + + +def test_future_reraises_original_error_when_session_active(fake_channel): + class FakeUnavailable(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def trailing_metadata(self): + return () + + guard_calls = [] + + def guard_cb(): + guard_calls.append(core._GUARD_MODE.get()) + return True + + guard_cb._supports_error_guard = True + + err = FakeUnavailable("transient") + fut_mock = mock.Mock(spec=grpc.Future) + fut_mock.result.side_effect = err + fut_mock.exception.return_value = err + + failing = FakeMultiCallable() + failing.future = mock.Mock(return_value=fut_mock) + fake_channel.multicallables["/svc/Method"] = failing + + channel = core.ProxiedChannel("example.com", guard_cb) + method = channel.unary_unary("/svc/Method") + fut = method.future("req") + + assert fut.exception() is err + with pytest.raises(FakeUnavailable, match="transient"): + fut.result() + assert guard_calls == ["pre_rpc", "on_error"] + + +@pytest.mark.parametrize( + "url,expected_dest", + [ + ("sc://example.com:443/;use_ssl=true", "example.com:443"), + ("sc://example.com/;use_ssl=true", "example.com:443"), + ("sc://example.com:8443/;use_ssl=true", "example.com:8443"), + ], +) +def test_direct_channel_uses_port_options_and_composite_credentials( + url, expected_dest +): + builder = core.ManagedSparkChannelBuilder(url) + fake_creds = mock.Mock() + fake_ssl = mock.Mock() + fake_call_creds = mock.Mock() + fake_composite = mock.Mock() + fake_secure_channel = mock.Mock() + + with ( + mock.patch.object( + core, "get_default_credentials", return_value=(fake_creds, "proj") + ) as get_creds, + mock.patch( + "google.auth.transport.grpc.AuthMetadataPlugin" + ) as plugin_cls, + mock.patch("google.auth.transport.requests.Request") as request_cls, + mock.patch.object( + grpc, "ssl_channel_credentials", return_value=fake_ssl + ), + mock.patch.object( + grpc, "metadata_call_credentials", return_value=fake_call_creds + ) as meta_call_creds, + mock.patch.object( + grpc, "composite_channel_credentials", return_value=fake_composite + ) as comp_creds, + mock.patch.object( + grpc, "secure_channel", return_value=fake_secure_channel + ) as secure_ch, + ): + result = builder._direct_channel() + + assert result is fake_secure_channel + get_creds.assert_called_once_with() + plugin_cls.assert_called_once_with(fake_creds, request_cls.return_value) + meta_call_creds.assert_called_once_with(plugin_cls.return_value) + comp_creds.assert_called_once_with(fake_ssl, fake_call_creds) + secure_ch.assert_called_once_with( + expected_dest, + fake_composite, + options=getattr(builder, "_channel_options", None), + ) diff --git a/tests/unit/test_dataproc.py b/tests/unit/test_dataproc.py new file mode 100644 index 00000000..6c4717e2 --- /dev/null +++ b/tests/unit/test_dataproc.py @@ -0,0 +1,439 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime +import os +import threading +import time +from unittest import mock + +import pytest +from google.api_core.client_options import ClientOptions +from google.auth import credentials as google_auth_credentials +from google.cloud.managed_spark_connect import session as session_mod +from google.cloud.managed_spark_connect.client import proxy as proxy_mod +from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, +) + + +@pytest.fixture(autouse=True) +def _reset_caches(): + session_mod.clear_dataproc_caches() + yield + session_mod.clear_dataproc_caches() + + +class _FakeRealCredentials: + """Minimal stand-in for google.auth.credentials.Credentials with bool .valid.""" + + def __init__(self, valid: bool = False, token: str | None = None): + self.valid = valid + self.token = token + self.refresh_calls = [] + + def refresh(self, request): + self.refresh_calls.append(request) + self.valid = True + self.token = "refreshed-bearer-token" + + +class _RealSubclassedCredentials(google_auth_credentials.Credentials): + """Concrete subclass of google.auth.credentials.Credentials for expiry lifecycle testing.""" + + def __init__(self): + super().__init__() + self.refresh_calls = [] + + def refresh(self, request): + self.refresh_calls.append(request) + self.token = f"token-{len(self.refresh_calls)}" + self.expiry = datetime.datetime.now(datetime.timezone.utc).replace( + tzinfo=None + ) + datetime.timedelta(hours=1) + + +def _set_unpatched_client_cls(monkeypatch, cls): + monkeypatch.setattr(session_mod.dataproc_v1, "SessionControllerClient", cls) + monkeypatch.setattr(session_mod, "SessionControllerClient", cls) + monkeypatch.setattr(session_mod, "_ORIG_SESSION_CONTROLLER_CLIENT", cls) + + +def test_get_default_credentials_caches_by_scopes_and_env(monkeypatch): + calls = [] + fake_creds = _FakeRealCredentials(valid=True, token="tok") + + def fake_default(scopes=None): + calls.append(tuple(scopes)) + return fake_creds, "proj-1" + + monkeypatch.setattr(proxy_mod.googleauth, "default", fake_default) + + c1, p1 = proxy_mod.get_default_credentials() + c2, p2 = proxy_mod.get_default_credentials() + assert c1 is fake_creds and c2 is fake_creds + assert p1 == "proj-1" and p2 == "proj-1" + assert calls == [proxy_mod._DEFAULT_SCOPES] + + # Distinct scopes tuple populates a separate cache key. + custom_scopes = ["https://www.googleapis.com/auth/userinfo.email"] + proxy_mod.get_default_credentials(scopes=custom_scopes) + proxy_mod.get_default_credentials(scopes=custom_scopes) + assert len(calls) == 2 + + # Changing GOOGLE_CLOUD_PROJECT invalidates the cache key. + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "other-proj") + proxy_mod.get_default_credentials() + assert len(calls) == 3 + + # Changing other google.auth.default() environment variables also + # invalidates the cache key. + for idx, env_var in enumerate( + ( + "CLOUDSDK_CONFIG", + "APPDATA", + "GCLOUD_PROJECT", + "DEVSHELL_PROJECT_ID", + ), + start=4, + ): + monkeypatch.setenv(env_var, f"val-{idx}") + proxy_mod.get_default_credentials() + assert len(calls) == idx + + # Clearing caches forces a fresh lookup. + session_mod.clear_dataproc_caches() + proxy_mod.get_default_credentials() + assert len(calls) == 8 + + +def test_get_default_credentials_bypasses_cache_when_side_effect_set(): + m1 = (mock.Mock(), "p1") + m2 = (mock.Mock(), "p2") + with mock.patch("google.auth.default", side_effect=[m1, m2]) as mock_def: + assert proxy_mod.get_default_credentials() == m1 + assert proxy_mod.get_default_credentials() == m2 + assert mock_def.call_count == 2 + assert proxy_mod._CREDENTIALS_CACHE == {} + + +def test_get_refreshed_credentials_skips_refresh_when_valid_real_creds( + monkeypatch, +): + creds = _FakeRealCredentials(valid=False, token=None) + + def fake_default(scopes=None): + return creds, "proj" + + monkeypatch.setattr(proxy_mod.googleauth, "default", fake_default) + + req_mod = mock.Mock() + monkeypatch.setattr(proxy_mod, "googleauthrequests", req_mod) + inst = Instrumentation([InMemorySink()]) + + # First call: valid=False, token=None -> refreshes once and sets valid=True. + out1 = proxy_mod.get_refreshed_credentials(recorder=inst.recorder) + assert out1 is creds + assert creds.valid is True + assert creds.token == "refreshed-bearer-token" + assert len(creds.refresh_calls) == 1 + + # Second call: valid=True (bool) and token is non-empty -> skips refresh, + # but still records both auth.default and auth.refresh spans. + out2 = proxy_mod.get_refreshed_credentials(recorder=inst.recorder) + assert out2 is creds + assert len(creds.refresh_calls) == 1 + + span_names = [r["name"] for r in inst.records if r["type"] == "span"] + assert span_names == [ + "auth.default", + "auth.refresh", + "auth.default", + "auth.refresh", + ] + + +def test_real_subclassed_credentials_refresh_expiry_lifecycle(monkeypatch): + creds = _RealSubclassedCredentials() + monkeypatch.setattr( + proxy_mod.googleauth, "default", lambda scopes=None: (creds, "proj") + ) + monkeypatch.setattr(proxy_mod, "googleauthrequests", mock.Mock()) + + # 1. Initially unrefreshed -> refreshes once. + proxy_mod.get_refreshed_credentials() + assert len(creds.refresh_calls) == 1 + assert creds.token == "token-1" + + # 2. Still valid -> skips refresh. + proxy_mod.get_refreshed_credentials() + assert len(creds.refresh_calls) == 1 + + # 3. Expire the token -> refreshes again. + creds.expiry = datetime.datetime.now(datetime.timezone.utc).replace( + tzinfo=None + ) - datetime.timedelta(minutes=5) + proxy_mod.get_refreshed_credentials() + assert len(creds.refresh_calls) == 2 + assert creds.token == "token-2" + + # 4. Credentials with expiry=None and a non-empty token (e.g. token-only + # Credentials or UserAccessTokenCredentials after initial refresh) report + # valid=True and do not re-refresh on every call. + creds.expiry = None + proxy_mod.get_refreshed_credentials() + assert len(creds.refresh_calls) == 2 + + +def test_token_only_and_user_access_token_credentials_with_none_expiry( + monkeypatch, +): + class _TokenOnlyCredentials(google_auth_credentials.Credentials): + + def __init__(self, token: str): + super().__init__() + self.token = token + + def refresh(self, request): + raise RuntimeError("Token-only credentials cannot be refreshed") + + token_creds = _TokenOnlyCredentials("ya29.static-token") + monkeypatch.setattr( + proxy_mod.googleauth, + "default", + lambda scopes=None: (token_creds, "proj"), + ) + monkeypatch.setattr(proxy_mod, "googleauthrequests", mock.Mock()) + + assert proxy_mod.get_refreshed_credentials() is token_creds + assert proxy_mod.get_refreshed_credentials() is token_creds + + class _UserAccessTokenLikeCredentials(google_auth_credentials.Credentials): + + def __init__(self): + super().__init__() + self.refresh_calls = 0 + + def refresh(self, request): + self.refresh_calls += 1 + self.token = f"gcloud-token-{self.refresh_calls}" + # gcloud CLI ADC leaves expiry as None after refresh. + + session_mod.clear_dataproc_caches() + gcloud_creds = _UserAccessTokenLikeCredentials() + monkeypatch.setattr( + proxy_mod.googleauth, + "default", + lambda scopes=None: (gcloud_creds, "proj"), + ) + assert proxy_mod.get_refreshed_credentials() is gcloud_creds + assert gcloud_creds.refresh_calls == 1 + assert proxy_mod.get_refreshed_credentials() is gcloud_creds + assert gcloud_creds.refresh_calls == 1 + + +def test_get_refreshed_credentials_evicts_cache_on_refresh_failure(monkeypatch): + bad_creds = _FakeRealCredentials(valid=False, token=None) + + def failing_refresh(_request): + raise RuntimeError("refresh failed") + + bad_creds.refresh = failing_refresh + monkeypatch.setattr( + proxy_mod.googleauth, "default", lambda scopes=None: (bad_creds, "proj") + ) + monkeypatch.setattr(proxy_mod, "googleauthrequests", mock.Mock()) + + with pytest.raises(RuntimeError, match="refresh failed"): + proxy_mod.get_refreshed_credentials() + assert proxy_mod._CREDENTIALS_CACHE == {} + + +@pytest.mark.parametrize( + "valid,token", + [ + (False, "existing-token"), + (True, None), + (True, ""), + ], +) +def test_get_refreshed_credentials_refreshes_when_invalid_or_missing_token( + monkeypatch, valid, token +): + creds = _FakeRealCredentials(valid=valid, token=token) + + def fake_default(scopes=None): + return creds, "proj" + + monkeypatch.setattr(proxy_mod.googleauth, "default", fake_default) + monkeypatch.setattr(proxy_mod, "googleauthrequests", mock.Mock()) + + proxy_mod.get_refreshed_credentials() + assert len(creds.refresh_calls) == 1 + + +def test_get_refreshed_credentials_refreshes_mock_credentials(monkeypatch): + """unittest.mock.Mock(token='tok') has a Mock .valid (not bool) and must refresh.""" + mock_creds = mock.Mock(token="tok") + req_mod = mock.Mock() + monkeypatch.setattr(proxy_mod, "googleauthrequests", req_mod) + with mock.patch("google.auth.default", return_value=(mock_creds, None)): + out = proxy_mod.get_refreshed_credentials() + assert out is mock_creds + mock_creds.refresh.assert_called_once_with(req_mod.Request.return_value) + + +def test_create_session_controller_client_caches_by_options_when_unpatched( + monkeypatch, +): + fake_creds = _FakeRealCredentials(valid=True, token="tok") + default_calls = [] + + def fake_default(scopes=None): + default_calls.append(scopes) + return fake_creds, "proj" + + created_clients = [] + + class FakeClient: + + def __init__(self, client_options=None, credentials=None): + self.client_options = client_options + self.credentials = credentials + created_clients.append(self) + + monkeypatch.setattr(proxy_mod.googleauth, "default", fake_default) + _set_unpatched_client_cls(monkeypatch, FakeClient) + + opts_a1 = ClientOptions(api_endpoint="us-central1-dataproc.googleapis.com") + opts_a2 = ClientOptions(api_endpoint="us-central1-dataproc.googleapis.com") + opts_b = ClientOptions(api_endpoint="us-east4-dataproc.googleapis.com") + opts_dict = {"api_endpoint": "us-central1-dataproc.googleapis.com"} + opts_custom_creds = ClientOptions( + api_endpoint="us-central1-dataproc.googleapis.com", + credentials_file="/path/to/key.json", + ) + + c1 = session_mod._create_session_controller_client(client_options=opts_a1) + c2 = session_mod._create_session_controller_client(client_options=opts_a2) + c_dict = session_mod._create_session_controller_client( + client_options=opts_dict + ) + c3 = session_mod._create_session_controller_client(client_options=opts_b) + c4 = session_mod._create_session_controller_client(client_options=None) + c5 = session_mod._create_session_controller_client( + client_options=opts_custom_creds + ) + + assert c1 is c2 + assert c1 is c_dict + assert c1 is not c3 + assert c1 is not c4 + assert c1 is not c5 + assert c1.credentials is fake_creds + assert c3.credentials is fake_creds + # Custom credentials_file omits shared ADC credentials to avoid DuplicateCredentialArgs. + assert c5.credentials is None + assert len(created_clients) == 4 + assert len(default_calls) == 1 + + # Changing CLOUDSDK_CONFIG / GCLOUD_PROJECT / DEVSHELL_PROJECT_ID / APPDATA + # invalidates both the client cache and default credentials cache. + for env_var in ( + "CLOUDSDK_CONFIG", + "APPDATA", + "GCLOUD_PROJECT", + "DEVSHELL_PROJECT_ID", + ): + prev_clients = len(created_clients) + monkeypatch.setenv(env_var, f"custom-{env_var}") + c_env = session_mod._create_session_controller_client( + client_options=opts_a1 + ) + assert len(created_clients) == prev_clients + 1 + assert c_env is created_clients[-1] + + +@pytest.mark.parametrize( + "patch_target", + [ + "google.cloud.dataproc_v1.SessionControllerClient", + "google.cloud.managed_spark_connect.session.SessionControllerClient", + ], +) +def test_create_session_controller_client_bypasses_cache_and_omits_creds_when_patched( + patch_target, +): + opts = ClientOptions(api_endpoint="us-central1-dataproc.googleapis.com") + with mock.patch(patch_target) as mock_cls: + mock_cls.side_effect = [mock.Mock(name="c1"), mock.Mock(name="c2")] + c1 = session_mod._create_session_controller_client(client_options=opts) + c2 = session_mod._create_session_controller_client(client_options=opts) + assert c1 is not c2 + assert mock_cls.call_count == 2 + mock_cls.assert_called_with(client_options=opts) + assert session_mod._CLIENT_CACHE == {} + + +def test_thread_safety_under_concurrent_callers(monkeypatch): + creds = _FakeRealCredentials(valid=False, token=None) + default_count = 0 + client_count = 0 + counter_lock = threading.Lock() + + def slow_default(scopes=None): + nonlocal default_count + time.sleep(0.01) + with counter_lock: + default_count += 1 + return creds, "proj" + + class SlowClient: + + def __init__(self, client_options=None, credentials=None): + nonlocal client_count + time.sleep(0.01) + with counter_lock: + client_count += 1 + self.client_options = client_options + self.credentials = credentials + + monkeypatch.setattr(proxy_mod.googleauth, "default", slow_default) + monkeypatch.setattr(proxy_mod, "googleauthrequests", mock.Mock()) + _set_unpatched_client_cls(monkeypatch, SlowClient) + + barrier = threading.Barrier(16) + results_client = [None] * 16 + results_creds = [None] * 16 + opts = ClientOptions(api_endpoint="us-east4-dataproc.googleapis.com") + + def worker(idx: int): + barrier.wait() + results_creds[idx] = proxy_mod.get_refreshed_credentials() + results_client[idx] = session_mod._create_session_controller_client( + client_options=opts + ) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(16)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=5) + + assert default_count == 1 + assert len(creds.refresh_calls) == 1 + assert client_count == 1 + assert len({id(c) for c in results_client}) == 1 + assert all(c is creds for c in results_creds) diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py new file mode 100644 index 00000000..5f3dad92 --- /dev/null +++ b/tests/unit/test_guard.py @@ -0,0 +1,749 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the session liveness guard and TTL cache.""" + +import logging +from unittest import mock + +from google.api_core.client_options import ClientOptions +from google.api_core.exceptions import ( + DeadlineExceeded, + InternalServerError, + NotFound, + PermissionDenied, + ServiceUnavailable, + TooManyRequests, +) +from google.cloud.dataproc_v1 import Session +import grpc +import pytest + +from google.cloud.managed_spark_connect import session as session_mod +from google.cloud.managed_spark_connect.client import core +from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, +) + +SESSION_NAME = "projects/test-project/locations/us-central1/sessions/sess-1" +SESSION_NAME_2 = "projects/test-project/locations/us-central1/sessions/sess-2" +CLIENT_OPTIONS = ClientOptions( + api_endpoint="us-central1-dataproc.googleapis.com" +) + + +def _make_session(state=Session.State.ACTIVE, name=SESSION_NAME) -> Session: + s = Session() + s.name = name + s.state = state + return s + + +@pytest.fixture(autouse=True) +def _isolated_guard_cache(monkeypatch): + monkeypatch.delenv(session_mod._ENV_LIVENESS_TTL_SECONDS, raising=False) + session_mod.invalidate_s8s_session_active() + yield + session_mod.invalidate_s8s_session_active() + + +def test_default_lazy_mode_skips_pre_rpc_get_session_and_span(monkeypatch): + """By default (Option A), pre-RPC guard checks do 0 GetSession RPCs and 0 spans.""" + inst = Instrumentation([InMemorySink()]) + clock = [100.0] + monkeypatch.setattr(session_mod.time, "monotonic", lambda: clock[0]) + + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + cb = session_mod.make_is_active_callback( + SESSION_NAME, CLIENT_OPTIONS, recorder=inst.recorder + ) + wrapped = core._GuardedMultiCallable(lambda req: f"ok:{req}", cb) + assert wrapped("r1") == "ok:r1" + # Even after 60 seconds, default lazy mode does NOT poll pre-RPC. + clock[0] = 160.0 + assert wrapped("r2") == "ok:r2" + create_client.assert_not_called() + + guard_spans = [ + r + for r in inst.records + if r.get("type") == "span" and r.get("name") == "dataproc.get_session" + ] + assert guard_spans == [] + + +def test_lazy_mode_checks_get_session_on_rpc_error_and_raises_runtime_error(): + """On RPC error, if GetSession reports session dead, raises RuntimeError and caches dead state.""" + inst = Instrumentation([InMemorySink()]) + + class FakeUnavailable(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def trailing_metadata(self): + return () + + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client, + ): + proxy_cls.return_value.port = 12345 + fake_inner = mock.MagicMock() + failing_rpc = mock.MagicMock( + side_effect=FakeUnavailable("broken stream") + ) + fake_inner.unary_unary.return_value = failing_rpc + builder_cls.return_value.toChannel.return_value = fake_inner + + client = create_client.return_value + client.get_session.return_value = _make_session( + Session.State.TERMINATED + ) + + cb = session_mod.make_is_active_callback( + SESSION_NAME, CLIENT_OPTIONS, recorder=inst.recorder + ) + channel = core.ProxiedChannel("example.com", cb) + method = channel.unary_unary( + "/spark.connect.SparkConnectService/Config" + ) + + with pytest.raises(RuntimeError, match="Session not active"): + method("req1") + assert client.get_session.call_count == 1 + assert failing_rpc.call_count == 1 + + # Subsequent RPC fast-fails pre-RPC without calling the underlying RPC or GetSession again. + with pytest.raises(RuntimeError, match="Session not active"): + method("req2") + assert client.get_session.call_count == 1 + assert failing_rpc.call_count == 1 + + +def test_lazy_mode_reraises_rpc_error_when_session_is_still_active(): + """On RPC error, if GetSession reports session ACTIVE, the original RpcError is re-raised.""" + + class FakeInternalError(grpc.RpcError): + + def code(self): + return grpc.StatusCode.INTERNAL + + def trailing_metadata(self): + return () + + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client, + ): + proxy_cls.return_value.port = 12345 + fake_inner = mock.MagicMock() + fake_inner.unary_unary.return_value = mock.MagicMock( + side_effect=FakeInternalError("query failed") + ) + builder_cls.return_value.toChannel.return_value = fake_inner + + client = create_client.return_value + client.get_session.return_value = _make_session(Session.State.ACTIVE) + + cb = session_mod.make_is_active_callback(SESSION_NAME, CLIENT_OPTIONS) + channel = core.ProxiedChannel("example.com", cb) + method = channel.unary_unary( + "/spark.connect.SparkConnectService/Config" + ) + + with pytest.raises(FakeInternalError, match="query failed"): + method("req1") + assert client.get_session.call_count == 1 + + +def test_lazy_mode_skips_get_session_for_structured_spark_sql_error(): + """Structured Spark SQL errors carrying grpc-status-details-bin do not trigger GetSession.""" + + class StructuredSparkError(grpc.RpcError): + + def code(self): + return grpc.StatusCode.INTERNAL + + def trailing_metadata(self): + return (("grpc-status-details-bin", b"\x08\x0d\x12\x0atable_gone"),) + + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client, + ): + proxy_cls.return_value.port = 12345 + fake_inner = mock.MagicMock() + fake_inner.unary_unary.return_value = mock.MagicMock( + side_effect=StructuredSparkError("AnalysisException") + ) + builder_cls.return_value.toChannel.return_value = fake_inner + + cb = session_mod.make_is_active_callback(SESSION_NAME, CLIENT_OPTIONS) + channel = core.ProxiedChannel("example.com", cb) + method = channel.unary_unary( + "/spark.connect.SparkConnectService/Config" + ) + + with pytest.raises(StructuredSparkError): + method("req1") + create_client.assert_not_called() + + +def test_lazy_mode_checks_get_session_on_mid_stream_error(): + """Mid-stream RpcError on ExecutePlan triggers GetSession and raises RuntimeError if session died.""" + + class MidStreamRpcError(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def trailing_metadata(self): + return () + + def broken_stream(): + yield "batch-1" + raise MidStreamRpcError("stream reset") + + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client, + ): + proxy_cls.return_value.port = 12345 + fake_inner = mock.MagicMock() + fake_inner.unary_stream.return_value = lambda _req: broken_stream() + builder_cls.return_value.toChannel.return_value = fake_inner + + client = create_client.return_value + client.get_session.return_value = _make_session( + Session.State.TERMINATED + ) + + cb = session_mod.make_is_active_callback(SESSION_NAME, CLIENT_OPTIONS) + channel = core.ProxiedChannel("example.com", cb) + stream = channel.unary_stream( + "/spark.connect.SparkConnectService/ExecutePlan" + )("plan") + + assert next(stream) == "batch-1" + assert client.get_session.call_count == 0 + with pytest.raises(RuntimeError, match="Session not active"): + next(stream) + assert client.get_session.call_count == 1 + + +def test_make_is_active_callback_preseeds_ttl_cache_and_skips_span(monkeypatch): + monkeypatch.setenv(session_mod._ENV_LIVENESS_TTL_SECONDS, "5.0") + inst = Instrumentation([InMemorySink()]) + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + cb = session_mod.make_is_active_callback( + SESSION_NAME, CLIENT_OPTIONS, recorder=inst.recorder + ) + assert cb() is True + assert cb() is True + create_client.assert_not_called() + + guard_spans = [ + r + for r in inst.records + if r.get("type") == "span" and r.get("name") == "dataproc.get_session" + ] + assert guard_spans == [] + + +def test_ttl_cache_miss_after_expiry_calls_get_session_and_reseeds(monkeypatch): + monkeypatch.setenv(session_mod._ENV_LIVENESS_TTL_SECONDS, "5.0") + inst = Instrumentation([InMemorySink()]) + clock = [100.0] + monkeypatch.setattr(session_mod.time, "monotonic", lambda: clock[0]) + + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + client = create_client.return_value + client.get_session.return_value = _make_session(Session.State.ACTIVE) + + cb = session_mod.make_is_active_callback( + SESSION_NAME, CLIENT_OPTIONS, recorder=inst.recorder + ) + # Within 5.0s TTL: cache hit, no RPC, no span. + clock[0] = 104.9 + assert cb() is True + assert client.get_session.call_count == 0 + + # Past 5.0s TTL: cache miss -> 1 GetSession RPC + 1 span -> re-seeded. + clock[0] = 105.1 + assert cb() is True + assert client.get_session.call_count == 1 + + # Next call within new TTL window: cache hit again. + clock[0] = 108.0 + assert cb() is True + assert client.get_session.call_count == 1 + + guard_spans = [ + r + for r in inst.records + if r.get("type") == "span" and r.get("name") == "dataproc.get_session" + ] + assert len(guard_spans) == 1 + assert guard_spans[0]["attrs"]["reason"] == "rpc_guard" + + +def test_ttl_zero_forces_check_on_every_rpc(monkeypatch): + monkeypatch.setenv(session_mod._ENV_LIVENESS_TTL_SECONDS, "0") + inst = Instrumentation([InMemorySink()]) + + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + client = create_client.return_value + client.get_session.return_value = _make_session(Session.State.ACTIVE) + + cb = session_mod.make_is_active_callback( + SESSION_NAME, CLIENT_OPTIONS, recorder=inst.recorder + ) + assert cb() is True + assert cb() is True + assert cb() is True + assert client.get_session.call_count == 3 + + guard_spans = [ + r + for r in inst.records + if r.get("type") == "span" and r.get("name") == "dataproc.get_session" + ] + assert len(guard_spans) == 3 + + +def test_ttl_negative_disables_pre_rpc_polling(monkeypatch): + monkeypatch.setenv(session_mod._ENV_LIVENESS_TTL_SECONDS, "-1") + inst = Instrumentation([InMemorySink()]) + clock = [100.0] + monkeypatch.setattr(session_mod.time, "monotonic", lambda: clock[0]) + + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + client = create_client.return_value + client.get_session.return_value = _make_session(Session.State.ACTIVE) + + cb = session_mod.make_is_active_callback( + SESSION_NAME, CLIENT_OPTIONS, recorder=inst.recorder + ) + clock[0] = 5000.0 + assert cb() is True + assert ( + session_mod.is_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) + is True + ) + create_client.assert_not_called() + + # Explicit invalidation requires a fresh GetSession lookup. + session_mod.invalidate_s8s_session_active(SESSION_NAME) + assert ( + session_mod.is_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) + is True + ) + assert client.get_session.call_count == 1 + + guard_spans = [ + r + for r in inst.records + if r.get("type") == "span" and r.get("name") == "dataproc.get_session" + ] + assert guard_spans == [] + + +def test_invalidate_s8s_session_active_targeted_and_global(monkeypatch): + monkeypatch.setenv(session_mod._ENV_LIVENESS_TTL_SECONDS, "5.0") + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + client = create_client.return_value + client.get_session.return_value = _make_session(Session.State.ACTIVE) + + session_mod.mark_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) + session_mod.mark_s8s_session_active(SESSION_NAME_2, CLIENT_OPTIONS) + + # Targeted invalidation clears only SESSION_NAME. + session_mod.invalidate_s8s_session_active(SESSION_NAME) + assert ( + session_mod.is_s8s_session_active(SESSION_NAME_2, CLIENT_OPTIONS) + is True + ) + assert client.get_session.call_count == 0 + + assert ( + session_mod.is_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) + is True + ) + assert client.get_session.call_count == 1 + + # Global invalidation clears all entries. + session_mod.invalidate_s8s_session_active() + assert ( + session_mod.is_s8s_session_active(SESSION_NAME_2, CLIENT_OPTIONS) + is True + ) + assert client.get_session.call_count == 2 + + +@pytest.mark.parametrize( + "transient_exc", + [ + ServiceUnavailable("503 Backend unavailable"), + DeadlineExceeded("504 Deadline exceeded"), + TooManyRequests("429 Quota exceeded"), + InternalServerError("500 Internal error"), + ], +) +def test_transient_control_plane_error_falls_back_to_last_known_active( + transient_exc, caplog +): + active_session = _make_session(Session.State.ACTIVE) + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + client = create_client.return_value + client.get_session.side_effect = [active_session, transient_exc] + + first = session_mod.get_active_s8s_session_response( + SESSION_NAME, CLIENT_OPTIONS + ) + assert first is active_session + + with caplog.at_level(logging.WARNING): + second = session_mod.get_active_s8s_session_response( + SESSION_NAME, CLIENT_OPTIONS + ) + assert second is active_session + assert any("Transient error" in m for m in caplog.messages) + + +def test_transient_error_without_prior_active_returns_none(): + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + client = create_client.return_value + client.get_session.side_effect = ServiceUnavailable("503") + assert ( + session_mod.get_active_s8s_session_response( + SESSION_NAME, CLIENT_OPTIONS + ) + is None + ) + + +@pytest.mark.parametrize( + "terminal_exc", + [ + NotFound("404 Session not found"), + PermissionDenied("403 Permission denied"), + RuntimeError("unexpected error"), + ], +) +def test_terminal_exceptions_invalidate_cache_and_return_none(terminal_exc): + active_session = _make_session(Session.State.ACTIVE) + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + client = create_client.return_value + client.get_session.side_effect = [ + active_session, + terminal_exc, + ServiceUnavailable("503 after terminal error"), + ] + + assert ( + session_mod.get_active_s8s_session_response( + SESSION_NAME, CLIENT_OPTIONS + ) + is active_session + ) + assert ( + session_mod.get_active_s8s_session_response( + SESSION_NAME, CLIENT_OPTIONS + ) + is None + ) + # Because the terminal exception cleared _LAST_KNOWN_ACTIVE, a subsequent + # transient error also returns None. + assert ( + session_mod.get_active_s8s_session_response( + SESSION_NAME, CLIENT_OPTIONS + ) + is None + ) + + +def test_non_active_session_state_invalidates_cache(): + active_session = _make_session(Session.State.ACTIVE) + terminated_session = _make_session(Session.State.TERMINATED) + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + client = create_client.return_value + client.get_session.side_effect = [active_session, terminated_session] + + assert ( + session_mod.get_active_s8s_session_response( + SESSION_NAME, CLIENT_OPTIONS + ) + is active_session + ) + assert ( + session_mod.get_active_s8s_session_response( + SESSION_NAME, CLIENT_OPTIONS + ) + is None + ) + assert session_mod._LIVENESS_CACHE == {} + assert session_mod._LAST_KNOWN_ACTIVE == {} + + +def test_should_check_liveness_on_error_classifications(): + class FakeMetadatum: + + def __init__(self, key, value): + self.key = key + self.value = value + + class RpcErrWithCode(grpc.RpcError): + + def __init__(self, code, metadata=()): + self._code = code + self._metadata = metadata + + def code(self): + return self._code + + def trailing_metadata(self): + return self._metadata + + # Non-RpcError exceptions never trigger GetSession. + assert core._should_check_liveness_on_error(ValueError("bad")) is False + assert core._should_check_liveness_on_error(SystemExit(1)) is False + assert core._should_check_liveness_on_error(KeyboardInterrupt()) is False + + # Client-cancelled or argument errors do not trigger GetSession. + assert ( + core._should_check_liveness_on_error( + RpcErrWithCode(grpc.StatusCode.CANCELLED) + ) + is False + ) + assert ( + core._should_check_liveness_on_error( + RpcErrWithCode(grpc.StatusCode.INVALID_ARGUMENT) + ) + is False + ) + + # UNAVAILABLE and NOT_FOUND always check liveness, even with status details. + details = (FakeMetadatum("grpc-status-details-bin", b"\x08\x0e"),) + assert ( + core._should_check_liveness_on_error( + RpcErrWithCode(grpc.StatusCode.UNAVAILABLE, details) + ) + is True + ) + assert ( + core._should_check_liveness_on_error( + RpcErrWithCode(grpc.StatusCode.NOT_FOUND, details) + ) + is True + ) + + # INTERNAL with _Metadatum status details (live Spark SQL error) skips check. + assert ( + core._should_check_liveness_on_error( + RpcErrWithCode(grpc.StatusCode.INTERNAL, details) + ) + is False + ) + # INTERNAL without status details checks liveness. + assert ( + core._should_check_liveness_on_error( + RpcErrWithCode(grpc.StatusCode.INTERNAL) + ) + is True + ) + + +def test_mid_stream_error_reraises_original_when_session_still_active(): + class MidStreamRpcError(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def trailing_metadata(self): + return () + + def broken_stream(): + yield "batch-1" + raise MidStreamRpcError("stream reset") + + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client, + ): + proxy_cls.return_value.port = 12345 + fake_inner = mock.MagicMock() + fake_inner.unary_stream.return_value = lambda _req: broken_stream() + builder_cls.return_value.toChannel.return_value = fake_inner + + client = create_client.return_value + client.get_session.return_value = _make_session(Session.State.ACTIVE) + + cb = session_mod.make_is_active_callback(SESSION_NAME, CLIENT_OPTIONS) + channel = core.ProxiedChannel("example.com", cb) + stream = channel.unary_stream( + "/spark.connect.SparkConnectService/ExecutePlan" + )("plan") + + assert next(stream) == "batch-1" + with pytest.raises(MidStreamRpcError, match="stream reset"): + next(stream) + assert client.get_session.call_count == 1 + + +def test_transient_control_plane_error_during_on_error_guard_keeps_session_alive( + monkeypatch, +): + """In default lazy mode, a newly created session survives transient 503 on error guard even after >60s.""" + clock = [100.0] + monkeypatch.setattr(session_mod.time, "monotonic", lambda: clock[0]) + + class TransientStreamError(grpc.RpcError): + + def code(self): + return grpc.StatusCode.UNAVAILABLE + + def trailing_metadata(self): + return () + + with ( + mock.patch.object(core.proxy, "ManagedSparkSessionProxy") as proxy_cls, + mock.patch.object(core, "DefaultChannelBuilder") as builder_cls, + mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client, + ): + proxy_cls.return_value.port = 12345 + fake_inner = mock.MagicMock() + fake_inner.unary_unary.return_value = mock.MagicMock( + side_effect=TransientStreamError("socket closed") + ) + builder_cls.return_value.toChannel.return_value = fake_inner + + client = create_client.return_value + client.get_session.side_effect = ServiceUnavailable( + "503 control plane blip" + ) + + # make_is_active_callback seeds _LAST_KNOWN_ACTIVE without needing a prior GetSession call. + cb = session_mod.make_is_active_callback(SESSION_NAME, CLIENT_OPTIONS) + channel = core.ProxiedChannel("example.com", cb) + method = channel.unary_unary( + "/spark.connect.SparkConnectService/Config" + ) + + # Advance past _MAX_TRANSIENT_FALLBACK_AGE_S (60s). + clock[0] = 250.0 + with pytest.raises(TransientStreamError, match="socket closed"): + method("req1") + assert client.get_session.call_count == 1 + + +def test_invalidate_and_terminate_affect_pre_rpc_and_negative_ttl(monkeypatch): + """Invalidation and termination are respected by pre_rpc and negative-TTL checks.""" + with mock.patch.object( + session_mod, "_create_session_controller_client" + ) as create_client: + client = create_client.return_value + client.get_session.return_value = _make_session( + Session.State.TERMINATED + ) + + cb = session_mod.make_is_active_callback(SESSION_NAME, CLIENT_OPTIONS) + + def check_active(): + if not cb(): + raise RuntimeError("Session not active") + + wrapped = core._GuardedMultiCallable( + lambda req: f"ok:{req}", check_active + ) + assert wrapped("r1") == "ok:r1" + assert client.get_session.call_count == 0 + + # Invalidating the active session forces pre_rpc to re-check GetSession. + session_mod.invalidate_s8s_session_active(SESSION_NAME) + with pytest.raises(RuntimeError, match="Session not active"): + wrapped("r2") + assert client.get_session.call_count == 1 + + # Even if invalidate_s8s_session_active is called again while dead, + # pre_rpc must not treat the session as cached active without checking. + session_mod.invalidate_s8s_session_active(SESSION_NAME) + with pytest.raises(RuntimeError, match="Session not active"): + wrapped("r3") + assert client.get_session.call_count == 2 + + # With TTL=-1, a dead session still returns False on direct is_s8s_session_active(). + monkeypatch.setenv(session_mod._ENV_LIVENESS_TTL_SECONDS, "-1") + assert ( + session_mod.is_s8s_session_active(SESSION_NAME, CLIENT_OPTIONS) + is False + ) + assert client.get_session.call_count == 2 + + # terminate_s8s_session marks the session dead immediately for both pre_rpc and direct checks. + session_mod.mark_s8s_session_active(SESSION_NAME_2, CLIENT_OPTIONS) + assert ( + session_mod.is_s8s_session_active(SESSION_NAME_2, CLIENT_OPTIONS) + is True + ) + client.get_session.return_value = _make_session( + Session.State.TERMINATING, name=SESSION_NAME_2 + ) + session_mod.terminate_s8s_session( + "test-project", "us-central1", "sess-2", CLIENT_OPTIONS + ) + prev_calls = client.get_session.call_count + assert ( + session_mod.is_s8s_session_active(SESSION_NAME_2, CLIENT_OPTIONS) + is False + ) + assert client.get_session.call_count == prev_calls diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index ab57ec85..939b1922 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -647,12 +647,24 @@ def test_module_import_calls_init_extras(self): # Undoing the patch does not undo that copy, so reload again outside it # or every later test in this process sees the mock. self.addCleanup(importlib.reload, msc) - with mock.patch( - "google.cloud.managed_spark_connect._ipython._init_extras" - ) as mock_init: + with ( + mock.patch("IPython.get_ipython", return_value=mock.Mock()), + mock.patch( + "google.cloud.managed_spark_connect._ipython._init_extras" + ) as mock_init, + ): importlib.reload(msc) mock_init.assert_called_once() + with ( + mock.patch("IPython.get_ipython", return_value=None), + mock.patch( + "google.cloud.managed_spark_connect._ipython._init_extras" + ) as mock_init, + ): + importlib.reload(msc) + mock_init.assert_not_called() + _OUR_MAGICS = (("line", "dpip"), ("line", "sparksql"), ("cell", "sparksql")) @@ -1097,6 +1109,33 @@ def test_import_without_interactive_deps(self): ["True", "False", "False", "True", "True", "True"], ) + def test_non_ipython_import_skips_heavy_optional_modules(self): + import os + import subprocess + import sys + + probe = """ +import sys +import google.cloud.managed_spark_connect as msc + +print(msc.ManagedSparkSession is not None) +print("google.cloud.managed_spark_connect._ipython" in sys.modules) +print("packaging.requirements" in sys.modules) +""" + env = dict(os.environ) + env.pop("MANAGED_SPARK_CONNECT_ENABLE_EXTRAS", None) + result = subprocess.run( + [sys.executable, "-c", probe], + capture_output=True, + text=True, + env=env, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual( + result.stdout.split(), + ["True", "False", "False"], + ) + _NO_PYSPARK_PROBE = """ import importlib.machinery, importlib.util, json, os, sys, types diff --git a/tests/unit/test_polling.py b/tests/unit/test_polling.py new file mode 100644 index 00000000..3719cba6 --- /dev/null +++ b/tests/unit/test_polling.py @@ -0,0 +1,161 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import time +import unittest +from unittest import mock + +from google.cloud.dataproc_v1 import Session +from google.cloud.managed_spark_connect import ManagedSparkSession +from google.cloud.managed_spark_connect import session as session_mod +from pyspark.sql.connect.client.core import ConfigResult +from pyspark.sql.connect.proto import ConfigResponse + + +class PollingTests(unittest.TestCase): + + def setUp(self): + self.original_environment = dict(os.environ) + os.environ.clear() + os.environ["GOOGLE_CLOUD_PROJECT"] = "test-project" + os.environ["GOOGLE_CLOUD_REGION"] = "test-region" + session_mod.clear_dataproc_caches() + + def tearDown(self): + session_mod.clear_dataproc_caches() + os.environ.clear() + os.environ.update(self.original_environment) + + @mock.patch("time.sleep") + @mock.patch( + "google.cloud.managed_spark_connect.session.SessionControllerClient" + ) + def test_terminate_s8s_session_stops_polling_after_max_polls( + self, mock_client_cls, mock_sleep + ): + mock_client = mock_client_cls.return_value + active = Session() + active.state = Session.State.ACTIVE + mock_client.get_session.return_value = active + + session_mod.terminate_s8s_session( + "test-project", "test-region", "test-session" + ) + + self.assertEqual(mock_client.get_session.call_count, 180) + self.assertEqual(mock_sleep.call_count, 179) + + @mock.patch("time.sleep") + @mock.patch( + "google.cloud.managed_spark_connect.session.SessionControllerClient" + ) + def test_terminate_s8s_session_skips_sleep_when_already_terminating( + self, mock_client_cls, mock_sleep + ): + mock_client = mock_client_cls.return_value + terminating = Session() + terminating.state = Session.State.TERMINATING + mock_client.get_session.return_value = terminating + + session_mod.terminate_s8s_session( + "test-project", "test-region", "test-session" + ) + + mock_client.terminate_session.assert_called_once() + mock_client.get_session.assert_called_once() + mock_sleep.assert_not_called() + + @mock.patch("time.sleep") + @mock.patch( + "google.cloud.managed_spark_connect.session.SessionControllerClient" + ) + def test_terminate_s8s_session_sleeps_only_between_active_polls( + self, mock_client_cls, mock_sleep + ): + mock_client = mock_client_cls.return_value + active = Session() + active.state = Session.State.ACTIVE + terminated = Session() + terminated.state = Session.State.TERMINATED + mock_client.get_session.side_effect = [active, active, terminated] + + session_mod.terminate_s8s_session( + "test-project", "test-region", "test-session" + ) + + self.assertEqual(mock_client.get_session.call_count, 3) + self.assertEqual(mock_sleep.call_count, 2) + mock_sleep.assert_called_with(1) + + @mock.patch( + "google.cloud.managed_spark_connect.session._server_spark_version", + return_value=None, + ) + @mock.patch("google.auth.default") + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") + @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") + @mock.patch( + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" + ) + @mock.patch( + "google.cloud.managed_spark_connect.session.is_s8s_session_active" + ) + def test_create_uses_adaptive_lro_polling_and_wakes_pbar_immediately( + self, + mock_is_active, + mock_gen_id, + mock_client_config, + mock_client_cls, + mock_creds, + _mock_spark_ver, + ): + mock_is_active.return_value = True + mock_gen_id.return_value = "sc-20240702-103952-abcdef" + mock_client_config.return_value = ConfigResult.fromProto( + ConfigResponse() + ) + cred = mock.MagicMock(token="token") + mock_creds.return_value = (cred, "") + + mock_client = mock_client_cls.return_value + mock_op = mock.Mock() + ready = Session() + ready.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect-server.example.com:443" + } + ready.uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + mock_op.result.return_value = ready + mock_client.create_session.return_value = mock_op + + terminating = Session() + terminating.state = Session.State.TERMINATING + mock_client.get_session.return_value = terminating + + t0 = time.monotonic() + session = ManagedSparkSession.builder.getOrCreate() + try: + elapsed = time.monotonic() - t0 + self.assertLess(elapsed, 0.8) + polling = mock_op.result.call_args.kwargs["polling"] + self.assertEqual(polling._initial, 1.0) + self.assertEqual(polling._maximum, 5.0) + self.assertEqual(polling._multiplier, 1.5) + self.assertEqual(polling._timeout, 1200) + finally: + session.stop() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_proxy.py b/tests/unit/test_proxy.py index 8aa8172b..ac778ea5 100644 --- a/tests/unit/test_proxy.py +++ b/tests/unit/test_proxy.py @@ -89,7 +89,8 @@ def echo_server(server_socket): target=echo_server, args=[server_socket], daemon=True ) t.start() - yield socket.create_connection(server_socket.getsockname()) + with socket.create_connection(server_socket.getsockname()) as conn: + yield conn def test_echo_server(echo_server_conn, test_message): @@ -124,9 +125,9 @@ def proxy_server(server_socket): target=proxy_server, args=[server_socket], daemon=True ) t.start() - conn = socket.create_connection(server_socket.getsockname()) - conn.settimeout(incoming_socket_timeout) - yield conn + with socket.create_connection(server_socket.getsockname()) as conn: + conn.settimeout(incoming_socket_timeout) + yield conn def retry_on_timeouts(method, arg): @@ -337,3 +338,343 @@ def test_proxy_start_span(): p.stop() (rec,) = _spans(inst, "proxy.start") assert rec["attrs"]["port"] == p.port > 0 + + +def test_recv_buffer_size_is_64_kib(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + assert proxy.RECV_BUFFER_SIZE == 65536 + + payload = b"x" * 32768 + recv_sizes = [] + chunks = [payload, b""] + + class FakeFromSock: + + def recv(self, buf_size): + recv_sizes.append(buf_size) + return chunks.pop(0) + + to_sock = mock.Mock() + stats = proxy._DirectionStats() + proxy.forward_bytes("test-64k", FakeFromSock(), to_sock, stats) + + assert recv_sizes == [65536, 65536] + to_sock.send.assert_called_once_with(payload) + to_sock.close.assert_called_once_with() + assert stats.bytes == 32768 + assert stats.frames == 1 + + +def test_forward_bytes_partial_send_loop_continuation(): + from google.cloud.managed_spark_connect.client import proxy + + payload = b"0123456789" * 1500 + assert len(payload) == 15000 + chunks = [payload, b""] + + class FakeFromSock: + + def recv(self, buf_size): + return chunks.pop(0) + + class PartialToSock: + + def __init__(self, max_chunk=1000): + self.max_chunk = max_chunk + self.received = bytearray() + self.send_calls = 0 + self.closed = False + + def send(self, bs): + self.send_calls += 1 + n = min(len(bs), self.max_chunk) + self.received.extend(bs[:n]) + return n + + def close(self): + self.closed = True + + to_sock = PartialToSock(max_chunk=1000) + stats = proxy._DirectionStats() + proxy.forward_bytes("test-partial", FakeFromSock(), to_sock, stats) + + assert to_sock.closed + assert to_sock.send_calls == 15 + assert bytes(to_sock.received) == payload + assert stats.bytes == len(payload) + assert stats.frames == 1 + + +def test_forward_bytes_partial_send_resets_timeout_attempts(): + from google.cloud.managed_spark_connect.client import proxy + + payload = b"abcdefghij" * 100 + chunks = [payload, b""] + + class FakeFromSock: + + def recv(self, buf_size): + return chunks.pop(0) + + class FlakyPartialToSock: + + def __init__(self): + self.received = bytearray() + self.step = 0 + self.closed = False + + def send(self, bs): + self.step += 1 + if self.step % 6 != 0: + raise TimeoutError("simulated send timeout") + n = min(len(bs), 250) + self.received.extend(bs[:n]) + return n + + def close(self): + self.closed = True + + to_sock = FlakyPartialToSock() + stats = proxy._DirectionStats() + proxy.forward_bytes("test-flaky-partial", FakeFromSock(), to_sock, stats) + + assert to_sock.closed + assert bytes(to_sock.received) == payload + assert stats.bytes == len(payload) + assert stats.frames == 1 + + +def test_forward_bytes_bridged_socket_send_returns_byte_count(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + ws_conn = mock.Mock() + b_sock = proxy.bridged_socket(ws_conn) + assert b_sock.send(b"\xde\xad\xbe\xef") == 4 + ws_conn.send.assert_called_once_with("deadbeef") + ws_conn.reset_mock() + + chunks = [b"\xde\xad\xbe\xef", b""] + + class FakeFromSock: + + def recv(self, buf_size): + return chunks.pop(0) + + proxy.forward_bytes("test-bridged", FakeFromSock(), b_sock) + ws_conn.send.assert_called_once_with("deadbeef") + ws_conn.close.assert_called_once_with() + + +def test_forward_bytes_zero_send_closes_both_sockets_and_skips_stats(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + from_sock = mock.Mock() + from_sock.recv.return_value = b"payload" + to_sock = mock.Mock() + to_sock.send.return_value = 0 + stats = proxy._DirectionStats() + + proxy.forward_bytes("test-zero-send", from_sock, to_sock, stats) + + to_sock.close.assert_called_once_with() + from_sock.close.assert_called_once_with() + assert stats.bytes == 0 + assert stats.frames == 0 + + +def test_forward_bytes_mid_write_error_closes_both_sockets_and_skips_stats(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + from_sock = mock.Mock() + from_sock.recv.return_value = b"0123456789" + to_sock = mock.Mock() + to_sock.send.side_effect = [5, BrokenPipeError("pipe closed")] + stats = proxy._DirectionStats() + + proxy.forward_bytes("test-broken-pipe", from_sock, to_sock, stats) + + assert to_sock.send.call_count == 2 + to_sock.close.assert_called_once_with() + from_sock.close.assert_called_once_with() + assert stats.bytes == 0 + assert stats.frames == 0 + + +def test_forward_bytes_ten_timeout_exhaustion_closes_both_sockets(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + from_sock = mock.Mock() + from_sock.recv.return_value = b"payload" + to_sock = mock.Mock() + to_sock.send.side_effect = TimeoutError("timed out") + stats = proxy._DirectionStats() + + proxy.forward_bytes("test-10-timeouts", from_sock, to_sock, stats) + + assert to_sock.send.call_count == 10 + to_sock.close.assert_called_once_with() + from_sock.close.assert_called_once_with() + assert stats.bytes == 0 + assert stats.frames == 0 + + +def test_connect_tcp_bridge_disables_compression_and_sets_max_size(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + creds = mock.Mock(token="tok-123") + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ), + mock.patch.object(proxy.websocketclient, "connect") as mock_ws_connect, + ): + ws = proxy.connect_tcp_bridge("backend.example.com") + + assert ws is mock_ws_connect.return_value + mock_ws_connect.assert_called_once_with( + "wss://backend.example.com/tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96", + additional_headers={"Authorization": "Bearer tok-123"}, + open_timeout=30, + compression=None, + max_size=proxy.MAX_WS_MESSAGE_SIZE, + ) + + +def test_managed_spark_session_proxy_context_manager_exits_cleanly(): + from google.cloud.managed_spark_connect.client import proxy + + with proxy.managed_spark_session_proxy(0, "h") as p: + assert p.port > 0 + + +def test_connect_tcp_bridge_loopback_no_deflate_extension(): + from unittest import mock + + import websockets.sync.client as websocketclient + import websockets.sync.server as websocketserver + + from google.cloud.managed_spark_connect.client import proxy + + creds = mock.Mock(token="tok-loopback") + server_extensions = [] + real_connect = websocketclient.connect + + def handle_ws(ws_conn): + server_extensions.append( + [type(ext).__name__ for ext in (ws_conn.protocol.extensions or [])] + ) + msg = ws_conn.recv(timeout=10) + ws_conn.send(msg) + + ws_server = websocketserver.serve(handle_ws, "127.0.0.1", 0) + ws_port = ws_server.socket.getsockname()[1] + t = threading.Thread(target=ws_server.serve_forever, daemon=True) + t.start() + + def local_connect(uri, **kwargs): + return real_connect( + uri.replace("wss://", "ws://", 1), + proxy=None, + **kwargs, + ) + + try: + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, None) + ), + mock.patch.object( + proxy.websocketclient, "connect", new=local_connect + ), + ): + with proxy.connect_tcp_bridge(f"127.0.0.1:{ws_port}") as ws_conn: + assert (ws_conn.protocol.extensions or []) == [] + b_sock = proxy.bridged_socket(ws_conn) + b_sock.send(b"hello-no-deflate") + assert b_sock.recv(65536) == b"hello-no-deflate" + assert server_extensions == [[]] + finally: + ws_server.shutdown() + t.join(timeout=2) + + +def test_forward_connection_sets_tcp_nodelay(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + conn = mock.MagicMock() + ws_cm = mock.MagicMock() + with ( + mock.patch.object(proxy, "connect_tcp_bridge", return_value=ws_cm), + mock.patch.object(proxy, "connect_sockets"), + ): + proxy.forward_connection(1, conn, ("127.0.0.1", 1), "h") + conn.setsockopt.assert_called_once_with( + socket.IPPROTO_TCP, socket.TCP_NODELAY, 1 + ) + + +def test_forward_connection_ignores_tcp_nodelay_oserror(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + conn = mock.MagicMock() + conn.setsockopt.side_effect = OSError("not a TCP socket") + ws_cm = mock.MagicMock() + with ( + mock.patch.object(proxy, "connect_tcp_bridge", return_value=ws_cm), + mock.patch.object(proxy, "connect_sockets"), + ): + proxy.forward_connection(1, conn, ("127.0.0.1", 1), "h") + conn.setsockopt.assert_called_once_with( + socket.IPPROTO_TCP, socket.TCP_NODELAY, 1 + ) + + +def test_proxy_run_sets_tcp_nodelay_on_accepted_socket(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + accepted_nodelay = [] + done = threading.Event() + + def fake_connect_sockets(conn_number, conn, backend_socket, span=None): + accepted_nodelay.append( + conn.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY) + ) + done.set() + + with ( + mock.patch.object( + proxy, "connect_tcp_bridge", return_value=mock.MagicMock() + ), + mock.patch.object( + proxy, "connect_sockets", side_effect=fake_connect_sockets + ), + ): + p = proxy.ManagedSparkSessionProxy(0, "h") + p.start() + try: + with socket.create_connection(("127.0.0.1", p.port)): + assert done.wait(timeout=5) + finally: + p.stop() + + assert len(accepted_nodelay) == 1 + assert accepted_nodelay[0] != 0 diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index ac3d87d9..7d64d779 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -3087,6 +3087,236 @@ def test_get_or_create_dataproc_batch_returns_local_session( self.assertIs(session, mock_local_session) mock_pyspark_sql_session.builder.getOrCreate.assert_called_once_with() + @mock.patch("google.auth.default") + @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") + @mock.patch( + "google.cloud.managed_spark_connect.session.SessionControllerClient" + ) + def test_custom_session_id_active_with_endpoint_single_rpc( + self, + mock_session_controller_client, + mock_client_config, + mock_credentials, + ): + """Attaching to an ACTIVE session with Spark Connect Server endpoint makes 1 GetSession RPC.""" + mock_client = mock_session_controller_client.return_value + mock_client_config.return_value = ConfigResult.fromProto( + ConfigResponse() + ) + cred = mock.MagicMock() + cred.token = "token" + mock_credentials.return_value = (cred, "") + + active_session = Session() + active_session.name = ( + "projects/test-project/locations/test-region/sessions/my-session" + ) + active_session.state = Session.State.ACTIVE + active_session.uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + active_session.runtime_info.endpoints = { + "Spark Connect Server": "sc://example.com:443" + } + mock_client.get_session.return_value = active_session + + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .sessionId("my-session") + .getOrCreate() + ) + try: + self.assertEqual(mock_client.get_session.call_count, 1) + mock_client.create_session.assert_not_called() + finally: + self.stopSession(mock_client, session) + + @mock.patch("time.sleep", return_value=None) + @mock.patch("google.auth.default") + @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") + @mock.patch( + "google.cloud.managed_spark_connect.session.SessionControllerClient" + ) + def test_custom_session_id_creating_waits_for_endpoint( + self, + mock_session_controller_client, + mock_client_config, + mock_credentials, + mock_sleep, + ): + """Attaching to a CREATING session without endpoint still polls until available.""" + mock_client = mock_session_controller_client.return_value + mock_client_config.return_value = ConfigResult.fromProto( + ConfigResponse() + ) + cred = mock.MagicMock() + cred.token = "token" + mock_credentials.return_value = (cred, "") + + creating_session = Session() + creating_session.name = ( + "projects/test-project/locations/test-region/sessions/my-session" + ) + creating_session.state = Session.State.CREATING + creating_session.uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + + ready_session = Session() + ready_session.name = creating_session.name + ready_session.state = Session.State.ACTIVE + ready_session.uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + ready_session.runtime_info.endpoints = { + "Spark Connect Server": "sc://example.com:443" + } + + mock_client.get_session.side_effect = [ + creating_session, + creating_session, + ready_session, + ] + + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .sessionId("my-session") + .getOrCreate() + ) + try: + self.assertEqual(mock_client.get_session.call_count, 3) + mock_client.create_session.assert_not_called() + finally: + self.stopSession(mock_client, session) + + @mock.patch("time.sleep", return_value=None) + @mock.patch("google.auth.default") + @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") + @mock.patch( + "google.cloud.managed_spark_connect.session.SessionControllerClient" + ) + def test_custom_session_id_creating_with_prepopulated_endpoint_waits_until_active( + self, + mock_session_controller_client, + mock_client_config, + mock_credentials, + mock_sleep, + ): + """A CREATING session whose endpoint URI is already populated still polls until ACTIVE.""" + mock_client = mock_session_controller_client.return_value + mock_client_config.return_value = ConfigResult.fromProto( + ConfigResponse() + ) + cred = mock.MagicMock() + cred.token = "token" + mock_credentials.return_value = (cred, "") + + creating_with_endpoint = Session() + creating_with_endpoint.name = ( + "projects/test-project/locations/test-region/sessions/my-session" + ) + creating_with_endpoint.state = Session.State.CREATING + creating_with_endpoint.uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + creating_with_endpoint.runtime_info.endpoints = { + "Spark Connect Server": "sc://example.com:443" + } + + ready_session = Session() + ready_session.name = creating_with_endpoint.name + ready_session.state = Session.State.ACTIVE + ready_session.uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + ready_session.runtime_info.endpoints = { + "Spark Connect Server": "sc://example.com:443" + } + + mock_client.get_session.side_effect = [ + creating_with_endpoint, + creating_with_endpoint, + creating_with_endpoint, + ready_session, + ] + + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .sessionId("my-session") + .getOrCreate() + ) + try: + self.assertEqual(mock_client.get_session.call_count, 4) + mock_client.create_session.assert_not_called() + finally: + self.stopSession(mock_client, session) + + @mock.patch("google.auth.default") + @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") + @mock.patch( + "google.cloud.managed_spark_connect.session.SessionControllerClient" + ) + def test_custom_session_id_switches_between_distinct_active_sessions( + self, + mock_session_controller_client, + mock_client_config, + mock_credentials, + ): + """Switching from active session-a to active session-b connects to session-b.""" + mock_client = mock_session_controller_client.return_value + mock_client_config.return_value = ConfigResult.fromProto( + ConfigResponse() + ) + cred = mock.MagicMock(token="token") + mock_credentials.return_value = (cred, "") + + sess_a = Session() + sess_a.name = ( + "projects/test-project/locations/test-region/sessions/session-a" + ) + sess_a.state = Session.State.ACTIVE + sess_a.uuid = "11111111-1111-1111-1111-111111111111" + sess_a.runtime_info.endpoints = { + "Spark Connect Server": "sc://host-a.example.com:443" + } + + sess_b = Session() + sess_b.name = ( + "projects/test-project/locations/test-region/sessions/session-b" + ) + sess_b.state = Session.State.ACTIVE + sess_b.uuid = "22222222-2222-2222-2222-222222222222" + sess_b.runtime_info.endpoints = { + "Spark Connect Server": "sc://host-b.example.com:443" + } + + mock_client.get_session.side_effect = [sess_a, sess_b] + + s1 = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .sessionId("session-a") + .getOrCreate() + ) + self.assertEqual( + ManagedSparkSession._active_s8s_session_id, "session-a" + ) + self.assertEqual( + ManagedSparkSession._active_s8s_session_uuid, sess_a.uuid + ) + + s2 = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .sessionId("session-b") + .getOrCreate() + ) + try: + self.assertIsNot(s1, s2) + self.assertEqual( + ManagedSparkSession._active_s8s_session_id, "session-b" + ) + self.assertEqual( + ManagedSparkSession._active_s8s_session_uuid, sess_b.uuid + ) + self.assertEqual(mock_client.get_session.call_count, 2) + mock_client.terminate_session.assert_not_called() + finally: + self.stopSession(mock_client, s2) + class ServerSparkVersionTest(unittest.TestCase):