diff --git a/README.md b/README.md index 54c4e055..44e40c36 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,11 @@ The `ManagedSparkSession.builder` provides a fluent API to configure the session | `subnetwork(subnet)` | Sets the subnetwork URI for the session. | | `ttl(duration)` | Sets the time-to-live (TTL) for the session using a `datetime.timedelta` object. | | `instrumentation(inst)` | Records what the library does on your behalf; see [Instrumentation (opt-in)](#instrumentation-opt-in). | +| `showOperationLinks(enabled)` | Show each query's Spark UI link (default `True`). Links only appear in notebooks and interactive terminals; `False` hides them there too. | +| `showProgressBar(enabled)` | Show session-creation and query progress bars (default `True`). Like links, they only appear in notebooks and interactive terminals. | +| `progressBarLeave(leave)` | Keep a query's progress bar on screen once it finishes (default `False`). | + +The three display settings apply when `getOrCreate()` creates or attaches to a session, not to a session it reuses from this process. ### Configuration (environment variables) @@ -134,6 +139,12 @@ These environment variables change how the client behaves. | `MANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST` | `1` | When a read query (`collect`, `toPandas`, `count`, ...) fails with `INVALID_CURSOR.POSITION_NOT_AVAILABLE` because the server no longer buffers the result the client has to resume from, re-run it once under a new operation id and log a `WARNING`. Commands (writes, DDL) are never re-run. Set to `0` to disable, e.g. for queries that are expensive or not deterministic. Read when the session is created. | | `MANAGED_SPARK_CONNECT_WS_SILENCE_LIMIT_S` | `300` | Cut the proxy WebSocket connection if an RPC is in flight and the server stays silent (no frames received) for this many seconds, so a wedged stream fails instead of hanging. The check runs on each 60 s WebSocket receive poll, so the cut happens between this limit and the limit plus 60 s. Set to `0` to disable. | | `MANAGED_SPARK_CONNECT_LIVENESS_TTL_SECONDS` | unset | Before an RPC, the client checks that the Dataproc session is still alive. Once a check finds it `ACTIVE`, that answer is reused for this many seconds: unset or negative reuses it until an RPC fails, `0` checks before every RPC. A failed RPC always triggers a fresh check. The session only counts as gone when Dataproc reports it not found or in a terminal state; when the check itself fails (for example during a control-plane outage), the RPC proceeds. | +| `GRPC_VERBOSITY` | `ERROR` | gRPC's own variable. If it is unset when this package is imported, and gRPC hasn't been imported yet, the package sets it to `ERROR` for the whole process to silence gRPC's informational fork messages. Set it yourself to override. | +| `MANAGED_SPARK_CONNECT_LOG_LEVEL` | unset | A level name (`DEBUG`, `INFO`, ...) or a positive number for the `google.cloud.managed_spark_connect` loggers. If the root logger has no handlers when the package is imported, the library logs to stderr itself, and its records then don't reach root handlers configured later: configure logging before importing the package to send them elsewhere. The library never configures the root logger. Setting it also lets pyspark print its query-context JSON for failed queries (see `SPARK_CONNECT_LOG_LEVEL`). | +| `SPARK_CONNECT_LOG_LEVEL` | unset | pyspark's own variable. Unless it or `MANAGED_SPARK_CONNECT_LOG_LEVEL` is set, creating a session silences pyspark's `SQLQueryContextLogger` and `DataFrameQueryContextLogger`, which otherwise print a JSON dump to stderr for every failed query. This applies to the whole process. | +| `MANAGED_SPARK_CONNECT_SHOW_OPERATION_LINKS` | `true` | Default for `showOperationLinks()`. | +| `MANAGED_SPARK_CONNECT_PROGRESS_BAR` | `true` | Default for `showProgressBar()`. | +| `MANAGED_SPARK_CONNECT_PROGRESS_BAR_LEAVE` | `false` | Default for `progressBarLeave()`. | ### Reusing Named Sessions Across Notebooks @@ -177,6 +188,12 @@ result conversions — as timed *spans*. It is off by default; opt in via `MANAGED_SPARK_CONNECT_INSTRUMENTATION_FILE` and `MANAGED_SPARK_CONNECT_INSTRUMENTATION_LOG` environment variables. +Without instrumentation the session still hooks a few pyspark client methods +on its own client (operation ids and links, progress bars, the +`INVALID_CURSOR.POSITION_NOT_AVAILABLE` replay; see +[Configuration](#configuration-environment-variables)) and coalesces the +`ReleaseExecute` RPCs of its result streams. + See [docs/instrumentation.md](https://github.com/GoogleCloudDataproc/managed-spark-connect-python/blob/main/docs/instrumentation.md) for usage, the API, the record schema, and the span and event catalogues. diff --git a/docs/instrumentation.md b/docs/instrumentation.md index 16ad6027..61fcd3ea 100644 --- a/docs/instrumentation.md +++ b/docs/instrumentation.md @@ -76,6 +76,10 @@ JsonlFileSink(path, *, append=True) InMemorySink() LoggingSink(logger=None, they set operation ids, display operation links, close a query's progress bar when it ends and re-run a read query once on `INVALID_CURSOR.POSITION_NOT_AVAILABLE` (see `reattach.position_lost` below). +Every session also coalesces the `ReleaseExecute(until)` RPCs of its result +streams, through a process-wide patch of pyspark's +`ExecutePlanResponseReattachableIterator._release_until` that leaves the +iterators of other clients alone. ## Record schema diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index dcee82e2..791be7a7 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -13,9 +13,20 @@ # limitations under the License. import importlib.metadata import importlib.util +import os import sys import warnings +# gRPC C-core logs fork-handler noise at INFO; default it to ERROR. gRPC reads +# the variable once, when it is first imported, so this only applies when this +# package is imported before grpc; an explicit GRPC_VERBOSITY always wins. +if "grpc" not in sys.modules: + os.environ.setdefault("GRPC_VERBOSITY", "ERROR") + +from . import _logging + +_logging.configure_package_logger() + def _in_ipython() -> bool: # Not environment.is_interactive(), which is also true in a plain @@ -92,7 +103,8 @@ def _check_pyspark_version(pyspark_version: str) -> None: warnings.warn( f"Package '{old_package_name}' is already installed in your environment. " f"This might cause conflicts with '{current_package_name}'. " - f"Consider uninstalling '{old_package_name}' and only install '{current_package_name}'." + f"Consider uninstalling '{old_package_name}' and only install '{current_package_name}'.", + stacklevel=2, ) except Exception: pass diff --git a/google/cloud/managed_spark_connect/_logging.py b/google/cloud/managed_spark_connect/_logging.py new file mode 100644 index 00000000..e6f998c7 --- /dev/null +++ b/google/cloud/managed_spark_connect/_logging.py @@ -0,0 +1,78 @@ +# Copyright 2026 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. +"""Logging setup. Never configures the root logger.""" + +import logging +import os +from typing import Optional + +PACKAGE_LOGGER_NAME = "google.cloud.managed_spark_connect" +ENV_LOG_LEVEL = "MANAGED_SPARK_CONNECT_LOG_LEVEL" +_PYSPARK_QUERY_CONTEXT_LOGGERS = ( + "SQLQueryContextLogger", + "DataFrameQueryContextLogger", +) + + +def _parse_level(raw: str) -> Optional[int]: + """A positive logging level from a name (``DEBUG``) or a number.""" + raw = raw.strip() + level = int(raw) if raw.isdigit() else getattr(logging, raw.upper(), None) + return level if isinstance(level, int) and level > 0 else None + + +def configure_package_logger() -> logging.Logger: + """Give the package logger a NullHandler and apply ``LOG_LEVEL``. + + With ``MANAGED_SPARK_CONNECT_LOG_LEVEL`` set and no handler on the root + logger, the package logger gets its own ``StreamHandler`` and stops + propagating, so a later ``logging.basicConfig()`` doesn't print every + record twice. + """ + pkg_logger = logging.getLogger(PACKAGE_LOGGER_NAME) + if not any(isinstance(h, logging.NullHandler) for h in pkg_logger.handlers): + pkg_logger.addHandler(logging.NullHandler()) + level = _parse_level(os.environ.get(ENV_LOG_LEVEL, "")) + if level is None: + return pkg_logger + pkg_logger.setLevel(level) + has_stream_handler = any( + not isinstance(h, logging.NullHandler) for h in pkg_logger.handlers + ) + if not has_stream_handler and not logging.getLogger().handlers: + pkg_logger.addHandler(logging.StreamHandler()) + pkg_logger.propagate = False + return pkg_logger + + +def quiet_pyspark_query_context_loggers() -> None: + """Stop pyspark from dumping every failed query's context to stderr. + + pyspark's ``SQLQueryContextLogger`` and ``DataFrameQueryContextLogger`` + write a JSON record with the server stack trace whenever an + ``AnalysisException`` or similar is built, on top of the exception + itself. Unless ``SPARK_CONNECT_LOG_LEVEL`` or + ``MANAGED_SPARK_CONNECT_LOG_LEVEL`` is set, raise both loggers above + ``CRITICAL``. They are process-wide, so this applies to every pyspark + session in the process. + """ + if ( + os.environ.get("SPARK_CONNECT_LOG_LEVEL", "").strip() + or os.environ.get(ENV_LOG_LEVEL, "").strip() + ): + return + from pyspark.logger.logger import PySparkLogger + + for name in _PYSPARK_QUERY_CONTEXT_LOGGERS: + PySparkLogger.getLogger(name).setLevel(logging.CRITICAL + 1) diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index b06a9413..73a48106 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -401,45 +401,54 @@ 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 ) self._is_active_callback = is_active_callback self._guard_errors = _OutageLog(logger) + # Discover and refresh credentials on the calling thread, before the + # proxy and gRPC threads start, so a credential helper that forks + # (gcloud) never does so while a gRPC thread holds an active ExecCtx. + # The result is cached for the proxy's own connections. + try: + proxy.get_refreshed_credentials(recorder=recorder) + except Exception: + logger.debug("Credential pre-warm failed", exc_info=True) self._proxy = proxy.ManagedSparkSessionProxy( 0, target_host, recorder=recorder ) self._proxy.start() - self._proxied_connect_url = f"sc://localhost:{self._proxy.port}" - self._wrapped = instrumentation_lib.wrap_channel( - DefaultChannelBuilder(self._proxied_connect_url).toChannel(), - recorder, - instrumentation_lib.INNER, - ) + try: + self._proxied_connect_url = f"sc://localhost:{self._proxy.port}" + self._wrapped = instrumentation_lib.wrap_channel( + DefaultChannelBuilder(self._proxied_connect_url).toChannel(), + recorder, + instrumentation_lib.INNER, + ) + except BaseException: + self._proxy.stop() + raise def __enter__(self): return self def __exit__(self, *args): - ret = self._wrapped.__exit__(*args) - self._proxy.stop() - return ret + try: + return self._wrapped.__exit__(*args) + finally: + self._proxy.stop() def close(self): - ret = self._wrapped.close() - self._proxy.stop() - return ret + try: + return self._wrapped.close() + finally: + self._proxy.stop() def _check_active(self): if self._is_active_callback is None: return mode = _GUARD_MODE.get() - if mode == "on_error" and not self._supports_error_guard: - return try: active = self._is_active_callback() except Exception as exc: @@ -457,7 +466,7 @@ def _check_active(self): return self._guard_errors.success() if not active: - logger.warning(f"Session is no longer active") + logger.debug("Session is no longer active") raise RuntimeError( "Session not active. Please create a new session" ) diff --git a/google/cloud/managed_spark_connect/client/proxy.py b/google/cloud/managed_spark_connect/client/proxy.py index 35361ae4..b7331528 100755 --- a/google/cloud/managed_spark_connect/client/proxy.py +++ b/google/cloud/managed_spark_connect/client/proxy.py @@ -70,8 +70,6 @@ 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, @@ -117,13 +115,8 @@ def get_refreshed_credentials(recorder=NOOP_RECORDER, scopes=None): with recorder.span("auth.refresh"): refreshed_attrs = None 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: - had_token = bool(getattr(creds, "token", None)) + if not (creds.valid and creds.token): + had_token = bool(creds.token) try: creds.refresh(googleauthrequests.Request()) except Exception: @@ -180,6 +173,17 @@ def send(self, msg_bytes): def close(self): return self._conn.close() + def close_abrupt(self): + """Close the WebSocket without waiting for a Close handshake.""" + try: + self._conn.close_socket() + except Exception: + pass + try: + return self._conn.close() + except Exception: + return None + class SendTimeout(Exception): """Sending one frame timed out on every retry.""" @@ -517,6 +521,8 @@ def forward_connection( recorder=NOOP_RECORDER, accepted_ns=None, proxy_instance=None, + on_backend_open=None, + on_backend_close=None, ): """Create a connection to the target and forward `conn` to it. @@ -543,6 +549,11 @@ def forward_connection( thread got scheduled. proxy_instance: Optional :class:`ManagedSparkSessionProxy` whose in-flight RPC count arms the backend socket's silence cut. + on_backend_open: Optional callback invoked with the backend + :class:`bridged_socket` once the WebSocket bridge is connected. + If it raises, the connection is abandoned. + on_backend_close: Optional callback invoked with the backend + :class:`bridged_socket` when forwarding finishes. """ inflight_rpcs = ( proxy_instance.get_inflight_rpcs if proxy_instance is not None else None @@ -565,18 +576,28 @@ def forward_connection( backend_socket = bridged_socket( websocket_conn, inflight_rpcs=inflight_rpcs ) - # Set a timeout on how long we will allow send/recv calls to block - # - # The code that reads and writes to this connection will retry - # on timeouts, so this is a safe change. - conn.settimeout(SOCKET_TIMEOUT_S) - connect_sockets( - conn_number, - conn, - backend_socket, - span if recorder.enabled else None, - recorder, - ) + if on_backend_open is not None: + on_backend_open(backend_socket) + try: + # Set a timeout on how long we will allow send/recv calls to block + # + # The code that reads and writes to this connection will retry + # on timeouts, so this is a safe change. + conn.settimeout(SOCKET_TIMEOUT_S) + connect_sockets( + conn_number, + conn, + backend_socket, + span if recorder.enabled else None, + recorder, + ) + finally: + if on_backend_close is not None: + on_backend_close(backend_socket) + + +# Total time ManagedSparkSessionProxy.stop() waits for its threads to exit. +_STOP_JOIN_TIMEOUT_S = 5.0 class ManagedSparkSessionProxy(object): @@ -603,6 +624,17 @@ def __init__(self, port, target_host, recorder=None): self._inflight_rpcs = 0 self._inflight_lock = threading.Lock() self._frontend_socket = None + self._accept_thread = None + # Guards _killed transitions and the three sets below, so stop() + # sees every socket and thread registered before it ran, and + # anything registering afterwards sees _killed. + self._conn_lock = threading.Lock() + self._active_client_sockets = set() + self._active_backend_sockets = set() + self._conn_threads = set() + # Whether a failed connection was logged at WARNING since the last + # successful one. + self._conn_failure_logged = False @property def port(self): @@ -635,10 +667,62 @@ def start(self, daemon=True): with self._recorder.span("proxy.start") as span: s = threading.Semaphore(value=0) t = threading.Thread(target=self._run, args=[s], daemon=daemon) + self._accept_thread = t t.start() s.acquire() span.set(port=self._port) + def _register_backend_socket(self, backend_socket): + with self._conn_lock: + if not self._killed: + self._active_backend_sockets.add(backend_socket) + self._conn_failure_logged = False + return + # stop() ran while this thread was still connecting. + backend_socket.close_abrupt() + raise ConnectionAbortedError("proxy stopped") + + def _unregister_backend_socket(self, backend_socket): + with self._conn_lock: + self._active_backend_sockets.discard(backend_socket) + + def _handle_connection(self, conn_number, conn, addr, accepted_ns): + try: + forward_connection( + conn_number, + conn, + addr, + self._target_host, + self._recorder, + accepted_ns, + proxy_instance=self, + on_backend_open=self._register_backend_socket, + on_backend_close=self._unregister_backend_socket, + ) + except Exception: + if not self._killed: + # gRPC reconnects quickly, so an outage (e.g. a backend 503) + # fails one connection after another: warn once, until a + # connection succeeds again. + if self._conn_failure_logged: + level = logging.DEBUG + else: + level = logging.WARNING + self._conn_failure_logged = True + logger.log( + level, + "Proxy connection %s failed", + conn_number, + exc_info=True, + ) + finally: + # forward_connection closes conn itself once it enters + # `with conn`; this covers a failure before that point. + _shutdown_and_close(conn) + with self._conn_lock: + self._active_client_sockets.discard(conn) + self._conn_threads.discard(threading.current_thread()) + def _run(self, s): with socket.create_server(("127.0.0.1", self._port)) as frontend_socket: self._frontend_socket = frontend_socket @@ -657,33 +741,45 @@ def _run(self, s): accepted_ns = (time.perf_counter_ns(), time.time_ns()) logger.debug(f"Accepted a connection from {addr}...") self._conn_number += 1 - threading.Thread( - target=forward_connection, - args=[ - self._conn_number, - conn, - addr, - self._target_host, - self._recorder, - accepted_ns, - self, - ], + t = threading.Thread( + target=self._handle_connection, + args=[self._conn_number, conn, addr, accepted_ns], daemon=True, - ).start() + ) + with self._conn_lock: + if self._killed: + _shutdown_and_close(conn) + break + self._active_client_sockets.add(conn) + self._conn_threads.add(t) + t.start() 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 + """Stop the proxy, close its sockets and wait for its threads. + + The threads get one shared deadline of ``_STOP_JOIN_TIMEOUT_S``. A + connection still opening its WebSocket can't be interrupted; it is + abandoned as soon as the WebSocket connects. + """ + with self._conn_lock: + self._killed = True + client_sockets = list(self._active_client_sockets) + backend_sockets = list(self._active_backend_sockets) + threads = [self._accept_thread, *self._conn_threads] + self._active_client_sockets.clear() + self._active_backend_sockets.clear() + + if self._frontend_socket is not None: + _shutdown_and_close(self._frontend_socket) + _shutdown_and_close(*client_sockets) + for backend_socket in backend_sockets: + backend_socket.close_abrupt() + + current = threading.current_thread() + deadline = time.monotonic() + _STOP_JOIN_TIMEOUT_S + for t in threads: + if t is not None and t is not current: + t.join(timeout=max(0.0, deadline - time.monotonic())) @contextlib.contextmanager diff --git a/google/cloud/managed_spark_connect/client/reattach_patch.py b/google/cloud/managed_spark_connect/client/reattach_patch.py index 15eeb816..0602bfc7 100644 --- a/google/cloud/managed_spark_connect/client/reattach_patch.py +++ b/google/cloud/managed_spark_connect/client/reattach_patch.py @@ -25,11 +25,21 @@ ``operation_id``. Command plans are never re-run. Every such failure emits ``reattach.position_lost``, and a replay also logs a ``WARNING``. ``MANAGED_SPARK_CONNECT_REPLAY_ON_POSITION_LOST=0`` turns the replay off. + +:func:`install_release_coalescing` makes the reattachable iterators of one +client send at most one ``ReleaseExecute(until)`` at a time. pyspark sends +one per response, each on its release thread pool, so a fast stream behind a +slow round trip queues up hundreds of RPCs that are stale by the time they +run. While one is in flight, later responses only move the release point, +and the next RPC sends the latest one. """ +import dataclasses import logging +import threading from typing import Any, Callable, Optional import uuid +import weakref from google.cloud.managed_spark_connect import environment from google.cloud.managed_spark_connect.instrumentation import NOOP_RECORDER @@ -137,3 +147,96 @@ def _execute_and_fetch(req, *args, **kwargs): return _execute_and_fetch install_client_hooks(client, {"_execute_and_fetch": replay_hook}) + + +# Stubs of the clients install_release_coalescing() was called for. An +# iterator coalesces only if it talks through one of them. +_COALESCED_STUBS: "weakref.WeakSet[Any]" = weakref.WeakSet() +_INSTALL_LOCK = threading.Lock() +_pyspark_release_until: Optional[Callable[..., None]] = None + + +@dataclasses.dataclass +class _ReleaseState: + """Per-iterator coalescing state, guarded by ``lock``.""" + + lock: threading.Lock = dataclasses.field(default_factory=threading.Lock) + next_id: Optional[str] = None # The latest response id not yet sent. + busy: bool = False # A drain is scheduled or running. + + +def install_release_coalescing(client: Any) -> None: + """Coalesce the ``ReleaseExecute(until)`` RPCs of ``client``'s iterators. + + pyspark's ``ExecutePlanResponseReattachableIterator._release_until`` is + replaced once per process, but iterators of other clients still run + pyspark's own version. + """ + global _pyspark_release_until + from pyspark.sql.connect.client.reattach import ( + ExecutePlanResponseReattachableIterator as Iterator, + ) + + with _INSTALL_LOCK: + _COALESCED_STUBS.add(client._stub) + if _pyspark_release_until is None: + _pyspark_release_until = Iterator._release_until + Iterator._release_until = _release_until + + +def _release_until(self, until_response_id: str) -> None: + if self._stub not in _COALESCED_STUBS: + return _pyspark_release_until(self, until_response_id) + if self._result_complete: + return + state = self.__dict__.get("_msc_release") + if state is None: + # Only send() calls this, from the consuming thread. + state = self._msc_release = _ReleaseState() + with state.lock: + state.next_id = until_response_id + if state.busy: + return + state.busy = True + with self._lock: + # pyspark <= 4.1 creates the shared pool with the first iterator and + # skips releases once shutdown() has cleared it. 4.2+ creates it on + # first use and drops it when no iterator is left, so a missing pool + # only means none was needed yet. + eager_pool = hasattr(type(self), "_get_or_create_release_thread_pool") + if not eager_pool or self._release_thread_pool_instance is not None: + try: + future = self._release_thread_pool.submit( + _drain_releases, self, state + ) + except RuntimeError: + # The pool is shut down. Skip the release, as pyspark does + # once the pool is gone; the server releases on completion. + pass + else: + # 4.2+ collects these so that client.close() can wait for + # releases in flight. + futures = getattr(self, "_release_futures", None) + if futures is not None: + futures.add(future) + return + with state.lock: + state.busy = False + + +def _drain_releases(it: Any, state: _ReleaseState) -> None: + """Send the latest pending release until none is left (pool thread).""" + while True: + with state.lock: + until_id, state.next_id = state.next_id, None + if until_id is None or it._result_complete: + state.busy = False + return + # What pyspark's _release_until submits, run synchronously here. + request = it._create_release_execute_request(until_id) + try: + for attempt in it._retrying(): + with attempt: + it._stub.ReleaseExecute(request, metadata=it._metadata) + except Exception as e: + logger.warning("ReleaseExecute failed with exception: %s.", e) diff --git a/google/cloud/managed_spark_connect/pypi_artifacts.py b/google/cloud/managed_spark_connect/pypi_artifacts.py index 270b68b7..e8ed4974 100644 --- a/google/cloud/managed_spark_connect/pypi_artifacts.py +++ b/google/cloud/managed_spark_connect/pypi_artifacts.py @@ -23,7 +23,7 @@ def __try_parsing_package(packages: set[str]) -> list["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") + logger.debug("It is recommended to pin the version of the package") return reqs def __init__(self, packages: set[str]): diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index 2adfa812..4702ef28 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -67,6 +67,7 @@ get_default_credentials, ) from google.cloud.managed_spark_connect.client.reattach_patch import ( + install_release_coalescing, install_replay_on_position_lost, ) from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException @@ -92,14 +93,13 @@ TerminateSessionRequest, ) from google.cloud.dataproc_v1.types import sessions +from google.cloud.managed_spark_connect import _logging from google.cloud.managed_spark_connect import _runtimes from google.cloud.managed_spark_connect import environment from google.cloud.managed_spark_connect import ide from pyspark.sql.connect.session import SparkSession from pyspark.sql.utils import to_str -# Set up logging -logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # System labels that should not be overridden by user @@ -249,27 +249,46 @@ 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] = {} #: Set to a true value to run the execution progress handler outside a #: notebook / interactive shell. A test hook, deliberately undocumented. _FORCE_PROGRESS_BAR_ENV = "MANAGED_SPARK_CONNECT_FORCE_PROGRESS_BAR" +SHOW_OPERATION_LINKS_ENV = "MANAGED_SPARK_CONNECT_SHOW_OPERATION_LINKS" +PROGRESS_BAR_ENV = "MANAGED_SPARK_CONNECT_PROGRESS_BAR" +PROGRESS_BAR_LEAVE_ENV = "MANAGED_SPARK_CONNECT_PROGRESS_BAR_LEAVE" def _progress_bar_forced() -> bool: return environment.env_bool(_FORCE_PROGRESS_BAR_ENV, default=False) -def _resolve_session_controller_client_cls(): - return ( - SessionControllerClient - if SessionControllerClient is not _ORIG_SESSION_CONTROLLER_CLIENT - else dataproc_v1.SessionControllerClient +def _display_flag( + explicit: Optional[bool], env_var: str, default: bool +) -> bool: + """The builder or constructor value, else ``env_var``, else ``default``.""" + if explicit is not None: + return explicit + from_env = environment.env_optional_bool(env_var) + return default if from_env is None else from_env + + +def _operation_links_enabled(explicit: Optional[bool]) -> bool: + # Links are only ever shown in an interactive environment; the settings + # can only turn them off there. + return environment.is_interactive() and _display_flag( + explicit, SHOW_OPERATION_LINKS_ENV, True ) +def _progress_bar_enabled(explicit: Optional[bool]) -> bool: + # Like _operation_links_enabled, unless a test forces the bar on. + return ( + _progress_bar_forced() or environment.is_interactive() + ) and _display_flag(explicit, PROGRESS_BAR_ENV, True) + + def _get_client_option(client_options: Any, name: str) -> Any: if client_options is None: return None @@ -290,12 +309,7 @@ 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) + client_cls = dataproc_v1.SessionControllerClient if ( _get_client_option(client_options, "client_cert_source") is not None or _get_client_option(client_options, "client_encrypted_cert_source") @@ -454,6 +468,10 @@ class ManagedSparkSession(SparkSession): _client_options = None _active_s8s_session_id: ClassVar[Optional[str]] = None _active_session_uses_custom_id: ClassVar[bool] = False + #: The session getOrCreate() made, whose server session this class + #: manages. Shadows pyspark's SparkSession._default_session, which a + #: plain pyspark session may hold. + _default_session: ClassVar[Optional["ManagedSparkSession"]] = None _execution_progress_bar = dict() #: Test hook: a ``tqdm``-like class used for execution progress bars #: instead of the environment's default (see :meth:`_set_tqdm_factory`). @@ -484,6 +502,9 @@ def __init__(self): self._instrumentation: Optional[Instrumentation] = None self._last_session_response: Optional[Session] = None self._custom_session_response: Optional[Session] = None + self._show_operation_links: Optional[bool] = None + self._show_progress_bar: Optional[bool] = None + self._progress_bar_leave: Optional[bool] = None @property def _recorder(self): @@ -495,6 +516,35 @@ def instrumentation(self, instrumentation: Optional[Instrumentation]): self._instrumentation = instrumentation return self + def showOperationLinks(self, enabled: bool = True): + """Show each query's Spark UI link (default: True). + + Links only ever appear in an interactive environment (a notebook + or terminal); ``False`` hides them there too. Overrides + ``MANAGED_SPARK_CONNECT_SHOW_OPERATION_LINKS``. Applies to a new + session, not to one ``getOrCreate()`` reuses. + """ + self._show_operation_links = bool(enabled) + return self + + def showProgressBar(self, enabled: bool = True): + """Show progress bars (default: True). + + Progress bars only ever appear in an interactive environment; + ``False`` hides them there too. Overrides + ``MANAGED_SPARK_CONNECT_PROGRESS_BAR``. Applies to a new session, + not to one ``getOrCreate()`` reuses. + """ + self._show_progress_bar = bool(enabled) + return self + + def progressBarLeave(self, leave: bool = True): + """Keep a query's progress bar on screen after it finishes + (default: False). Overrides + ``MANAGED_SPARK_CONNECT_PROGRESS_BAR_LEAVE``.""" + self._progress_bar_leave = bool(leave) + return self + @property def session_controller_client(self) -> SessionControllerClient: """Get or create a SessionControllerClient instance.""" @@ -648,7 +698,7 @@ def create(self) -> "ManagedSparkSession": "ManagedSparkSession allows session creation only through getOrCreate" ) - def __create_spark_connect_session_from_s8s( + def _create_spark_connect_session_from_s8s( self, session_response, session_name ) -> "ManagedSparkSession": recorder = self._recorder @@ -658,6 +708,7 @@ def __create_spark_connect_session_from_s8s( return self.__connect(session_response, session_name) def __connect(self, session_response, session_name): + os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" ManagedSparkSession._active_s8s_session_uuid = session_response.uuid ManagedSparkSession._project_id = self._project_id ManagedSparkSession._region = self._region @@ -683,12 +734,24 @@ def __connect(self, session_response, session_name): session = ManagedSparkSession( connection=self._channel_builder, instrumentation=self._instrumentation, + show_operation_links=self._show_operation_links, + show_progress_bar=self._show_progress_bar, + progress_bar_leave=self._progress_bar_leave, ) # Register handler for Cell Execution Progress bar session._register_progress_execution_handler() - ManagedSparkSession._set_default_and_active_session(session) + # This class owns its default session: the one whose server + # session stop() terminates. pyspark's _set_default_and_active_session + # would keep a default that another (plain pyspark) session set. + with ManagedSparkSession._lock: + ManagedSparkSession._default_session = session + if ( + getattr(ManagedSparkSession._active_session, "session", None) + is None + ): + ManagedSparkSession._active_session.session = session return session @@ -733,31 +796,34 @@ def __create(self) -> "ManagedSparkSession": s8s_creation_start_time = time.time() stop_create_session_pbar_event = threading.Event() + create_session_pbar_thread: Optional[threading.Thread] = None - def create_session_pbar(): - iterations = 150 - pbar = tqdm.trange( - iterations, - bar_format="{bar}", - ncols=80, - ) - for i in pbar: - if stop_create_session_pbar_event.is_set(): - break - # Last iteration: wait until session is created - if i >= iterations - 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 - print() + if _progress_bar_enabled(self._show_progress_bar): - create_session_pbar_thread = threading.Thread( - target=create_session_pbar - ) + def create_session_pbar(): + iterations = 150 + pbar = tqdm.trange( + iterations, + bar_format="{bar}", + ncols=80, + ) + for i in pbar: + if stop_create_session_pbar_event.is_set(): + break + # Last iteration: wait until session is created + if i >= iterations - 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 + print() + + create_session_pbar_thread = threading.Thread( + target=create_session_pbar + ) # Activate Spark Connect mode for Spark client os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" @@ -786,7 +852,8 @@ def create_session_pbar(): ) self._display_session_link_on_creation(session_id) self._display_view_session_details_button(session_id) - create_session_pbar_thread.start() + if create_session_pbar_thread is not None: + create_session_pbar_thread.start() with self._recorder.span( "session.wait_active", poll_interval_s=5.0 ) as wait_span: @@ -800,8 +867,6 @@ def create_session_pbar(): ) ) wait_span.set(**_session_state_attrs(session_response)) - stop_create_session_pbar_event.set() - create_session_pbar_thread.join() file_path = ( ManagedSparkSession._get_active_session_file_path() ) @@ -821,27 +886,18 @@ def create_session_pbar(): f"Exception while writing active session to file {file_path}, {e}" ) except (InvalidArgument, PermissionDenied) as e: - stop_create_session_pbar_event.set() - if create_session_pbar_thread.is_alive(): - create_session_pbar_thread.join() ManagedSparkSession._active_s8s_session_id = None ManagedSparkSession._active_session_uses_custom_id = False raise ManagedSparkConnectException( f"Error while creating Managed Spark Session: {e.message}" ) except DefaultCredentialsError as e: - stop_create_session_pbar_event.set() - if create_session_pbar_thread.is_alive(): - create_session_pbar_thread.join() ManagedSparkSession._active_s8s_session_id = None ManagedSparkSession._active_session_uses_custom_id = False raise ManagedSparkConnectException( "Credentials error while creating Managed Spark Session (see https://docs.cloud.google.com/docs/authentication/provide-credentials-adc for more info)" ) from e except Exception as e: - stop_create_session_pbar_event.set() - if create_session_pbar_thread.is_alive(): - create_session_pbar_thread.join() ManagedSparkSession._active_s8s_session_id = None ManagedSparkSession._active_session_uses_custom_id = False raise RuntimeError( @@ -849,11 +905,16 @@ def create_session_pbar(): ) from e finally: stop_create_session_pbar_event.set() + if ( + create_session_pbar_thread is not None + and create_session_pbar_thread.is_alive() + ): + create_session_pbar_thread.join() logger.debug( f"Managed Spark Session created: {session_id} in {int(time.time() - s8s_creation_start_time)} seconds" ) - session = self.__create_spark_connect_session_from_s8s( + session = self._create_spark_connect_session_from_s8s( session_response, session_config.name ) # After connecting, so the message can carry the server's @@ -1058,30 +1119,51 @@ def _get_exiting_active_session( session_response = self._wait_for_session_available( session_name ) - session = self.__create_spark_connect_session_from_s8s( + session = self._create_spark_connect_session_from_s8s( session_response, session_name ) - links = _session_links( - self._region, - self._project_id, - s8s_session_id, - html=False, - ) - versions = self._session_versions(session) - subject = ( - f"Using existing Managed Spark Session {s8s_session_id} " - f"{versions}" - ).rstrip() - print( - f"{subject}; configuration changes may not be applied: " - f"{links.url}" - ) - self._display_view_session_details_button(s8s_session_id) + links = _session_links( + self._region, + self._project_id, + s8s_session_id, + html=False, + ) + versions = self._session_versions(session) + subject = ( + f"Using existing Managed Spark Session {s8s_session_id} " + f"{versions}" + ).rstrip() + status = f"{subject}; configuration changes may not be applied" + self._output_element_or_message( + f"{status}: {links.url}", + _status_line_html( + html_lib.escape(status), + links=_session_links( + self._region, + self._project_id, + s8s_session_id, + ), + ), + ) + self._display_view_session_details_button( + s8s_session_id + ) + else: + logger.debug( + "Reusing active in-process Managed Spark Session: %s", + s8s_session_id, + ) return session else: if session is not None: - print( - f"{s8s_session_id} Managed Spark Session is not active, stopping and creating a new one" + plain_message = ( + f"{s8s_session_id} Managed Spark Session is not active, " + "stopping and creating a new one" + ) + logger.debug(plain_message) + self._output_element_or_message( + plain_message, + _status_line_html(html_lib.escape(plain_message)), ) session.stop() @@ -1139,6 +1221,11 @@ def getOrCreate(self) -> "ManagedSparkSession": finally: self._custom_session_response = None + # Whichever path got it (new, reattached or reused in-process), + # this session now owns the server session and the class-level + # state, which only the default session's stop() handles. + ManagedSparkSession._default_session = session + # Register this session as the instantiated SparkSession for compatibility # with tools and libraries that expect SparkSession._instantiatedSession from pyspark.sql import SparkSession as PySparkSQLSession @@ -1403,6 +1490,9 @@ def _check_python_version_compatibility(self, runtime_version): import sys import warnings + if not runtime_version: + return + runtime_info = _runtimes.get_runtime_info(runtime_version) if runtime_info is None: return @@ -1543,7 +1633,7 @@ def _get_session_by_id(self, session_id: str) -> Optional[Session]: Session.State.CREATING, ]: # Reuse the active session - logger.info(f"Reusing existing session: {session_id}") + logger.debug(f"Reusing existing session: {session_id}") return session else: # Session exists but is not usable (terminated/failed/terminating) @@ -1631,6 +1721,9 @@ def __init__( connection: Union[str, ManagedSparkChannelBuilder], user_id: Optional[str] = None, instrumentation: Optional[Instrumentation] = None, + show_operation_links: Optional[bool] = None, + show_progress_bar: Optional[bool] = None, + progress_bar_leave: Optional[bool] = None, ): """ Creates a new ManagedSparkSession for the Spark Connect interface. @@ -1647,11 +1740,19 @@ def __init__( instrumentation : :class:`Instrumentation`, optional Records spans for the client-side work of this session (see the ``instrumentation`` package). Exposed as ``self.instrumentation``. + show_operation_links, show_progress_bar, progress_bar_leave : bool, optional + As the builder's ``showOperationLinks``, ``showProgressBar`` and + ``progressBarLeave``; ``None`` falls back to the environment. """ + os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" super().__init__(connection, user_id) + _logging.quiet_pyspark_query_context_loggers() self.instrumentation = instrumentation + self._show_operation_links: Optional[bool] = show_operation_links + self._show_progress_bar: Optional[bool] = show_progress_bar + self._progress_bar_leave: Optional[bool] = progress_bar_leave # The _execute_and_fetch hooks stack in install order, innermost # first: the position-lost replay (so the hooks around it see one # call, with the caller's request, however many attempts it takes), @@ -1663,47 +1764,35 @@ def __init__( ) if instrumentation is not None and instrumentation.wrap_client: install_client_wrappers(self.client, instrumentation.recorder) + install_release_coalescing(self.client) - execute_plan_request_base_method = ( - self.client._execute_plan_request_with_metadata - ) - execute_base_method = self.client._execute - execute_and_fetch_as_iterator_base_method = ( - self.client._execute_and_fetch_as_iterator - ) - - def execute_plan_request_wrapped_method(*args, **kwargs): - req = execute_plan_request_base_method(*args, **kwargs) - if not req.operation_id: - req.operation_id = str(uuid.uuid4()) - logger.debug( - f"No operation_id found. Setting operation_id: {req.operation_id}" - ) - return req + def execute_plan_request_hook(inner): + def _execute_plan_request_with_metadata(*args, **kwargs): + req = inner(*args, **kwargs) + if not req.operation_id: + req.operation_id = str(uuid.uuid4()) + logger.debug( + f"No operation_id found. Setting operation_id: {req.operation_id}" + ) + return req - self.client._execute_plan_request_with_metadata = ( - execute_plan_request_wrapped_method - ) + return _execute_plan_request_with_metadata - def execute_wrapped_method(client_self, req, *args, **kwargs): - if not self._sql_lazy_transformation(req): - self._display_operation_link(req.operation_id) - execute_base_method(req, *args, **kwargs) + def execute_hook(inner): + def _execute(req, *args, **kwargs): + if not self._sql_lazy_transformation(req): + self._display_operation_link(req.operation_id) + return inner(req, *args, **kwargs) - self.client._execute = MethodType(execute_wrapped_method, self.client) + return _execute - def execute_and_fetch_as_iterator_wrapped_method( - client_self, req, *args, **kwargs - ): - if not self._sql_lazy_transformation(req): - self._display_operation_link(req.operation_id) - return execute_and_fetch_as_iterator_base_method( - req, *args, **kwargs - ) + def execute_and_fetch_as_iterator_hook(inner): + def _execute_and_fetch_as_iterator(req, *args, **kwargs): + if not self._sql_lazy_transformation(req): + self._display_operation_link(req.operation_id) + return inner(req, *args, **kwargs) - self.client._execute_and_fetch_as_iterator = MethodType( - execute_and_fetch_as_iterator_wrapped_method, self.client - ) + return _execute_and_fetch_as_iterator def execute_and_fetch_hook(inner): def _execute_and_fetch(req, *args, **kwargs): @@ -1717,9 +1806,18 @@ def _execute_and_fetch(req, *args, **kwargs): return _execute_and_fetch - # Always installed (instrumented or not), outermost; stop() removes it. + # Always installed (instrumented or not), outermost; _close_client() + # removes them along with every other hook on the client. install_client_hooks( - self.client, {"_execute_and_fetch": execute_and_fetch_hook} + self.client, + { + "_execute_plan_request_with_metadata": execute_plan_request_hook, + "_execute": execute_hook, + "_execute_and_fetch_as_iterator": ( + execute_and_fetch_as_iterator_hook + ), + "_execute_and_fetch": execute_and_fetch_hook, + }, ) # Patching clearProgressHandlers method to not remove Managed Spark Progress Handler @@ -1756,6 +1854,12 @@ def get_tqdm_bar(): except ImportError: return cli_tqdm + def _should_show_operation_links(self) -> bool: + return _operation_links_enabled(self._show_operation_links) + + def _should_show_progress_bar(self) -> bool: + return _progress_bar_enabled(self._show_progress_bar) + def _register_progress_execution_handler(self): from pyspark.sql.connect.shell.progress import StageInfo @@ -1769,8 +1873,8 @@ def handler( return # Don't build / render progress bar for non-interactive (despite - # Ipython or non-IPython), unless tests forced the bar on. - if not (_progress_bar_forced() or environment.is_interactive()): + # Ipython or non-IPython) or when disabled, unless tests forced the bar on. + if not self._should_show_progress_bar(): return total_tasks = 0 @@ -1801,7 +1905,11 @@ def handler( tqdm_pbar = type(self)._tqdm_factory or self.get_tqdm_bar() pbar = tqdm_pbar( total=total_tasks, - leave=True, + leave=_display_flag( + self._progress_bar_leave, + PROGRESS_BAR_LEAVE_ENV, + False, + ), dynamic_ncols=True, bar_format="{l_bar}{bar} {n_fmt}/{total_fmt} Tasks", ) @@ -1876,6 +1984,16 @@ def _close_progress_bar( reason=reason, ) + @classmethod + def getActiveSession(cls) -> Optional["ManagedSparkSession"]: + """The thread's active session, if it is a ``ManagedSparkSession``. + + pyspark keeps one active session per thread for every + ``SparkSession`` class, so it can be a plain pyspark session. + """ + session = super().getActiveSession() + return session if isinstance(session, ManagedSparkSession) else None + @classmethod def progress_bars(cls) -> Dict[str, Any]: """A snapshot of the open execution progress bars. @@ -1934,18 +2052,19 @@ def _repr_html_(self) -> str: def _display_operation_link(self, operation_id: str): # Don't print per-operation Spark UI link for non-interactive (despite - # Ipython or non-IPython) - if not environment.is_interactive(): + # Ipython or non-IPython) or when disabled via Builder/env. + if not self._should_show_operation_links(): return - assert all( + if not all( [ operation_id is not None, self._region is not None, self._active_s8s_session_id is not None, self._project_id is not None, ] - ) + ): + return url = ( f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/" @@ -2089,76 +2208,134 @@ def stop(self, terminate: Optional[bool] = None) -> None: with recorder.span("session.stop") as span: self._stop(terminate, span) finally: - uninstall_client_hooks(self.client) if self.instrumentation is not None: - uninstall_client_wrappers(self.client) recorder.trace_id = "" self.instrumentation.close() + def __del__(self) -> None: + # pyspark's SparkSession.__del__ closes the client again. After + # stop() the GC may run that on any thread, including a + # ReleaseExecute pool worker, where shutting the shared release pool + # down fails half-way and leaves it unusable for the whole process. + if self.__dict__.get("_client_closed"): + return + super().__del__() + + def _close_client(self) -> None: + """Close the client and release everything that pins this session. + + Unregistering pyspark's atexit hook lets a stopped session be garbage + collected. Removing the client hooks, the progress handler and the + clearProgressHandlers wrapper breaks the session <-> client reference + cycles, so that mostly happens by reference counting rather than in + the cyclic GC. + """ + client = self.client + # pyspark's __del__ closes the StreamingQueryManager; ours doesn't. + sqm = self.__dict__.get("_sqm") + if sqm is not None: + try: + sqm.close() + except Exception: + logger.debug( + "Failed to close the StreamingQueryManager", exc_info=True + ) + try: + client.close() + except Exception: + logger.warning( + "Failed to close the Spark Connect client", exc_info=True + ) + finally: + # Even when close() failed: __del__ must not retry it. + self._client_closed = True + # pyspark registers one exit hook per client, under a name that + # changes: _cleanup_ml (4.0), _cleanup_ml_cache (4.1), _on_exit (4.2+). + # Unregistering a function that is not registered is a no-op. + for hook_name in ("_on_exit", "_cleanup_ml_cache", "_cleanup_ml"): + hook = getattr(client, hook_name, None) + if hook is not None: + atexit.unregister(hook) + if self.instrumentation is not None: + uninstall_client_wrappers(client) + uninstall_client_hooks(client) + client.clear_progress_handlers() + self.__dict__.pop("clearProgressHandlers", None) + def _stop(self, terminate: Optional[bool], span) -> None: with ManagedSparkSession._lock: - if ManagedSparkSession._active_s8s_session_id is not None: - # Determine if we should terminate the server-side session - if terminate is None: - # Auto-detect: managed sessions terminate, named sessions don't - should_terminate = ( - not ManagedSparkSession._active_session_uses_custom_id - ) - else: - should_terminate = terminate + try: + # The server session and the class-level state belong to the + # default session. Stopping any other one (attached directly, + # or replaced by a later getOrCreate()) only closes its client. + if self is ManagedSparkSession._default_session: + self._stop_server_session(terminate, span) + finally: + self._close_client() + if self is ManagedSparkSession._default_session: + ManagedSparkSession._default_session = None + if self is getattr( + ManagedSparkSession._active_session, "session", None + ): + ManagedSparkSession._active_session.session = None - span.set( - terminate=should_terminate, - session_id=ManagedSparkSession._active_s8s_session_id, - ) - if should_terminate: - # Terminate the server-side session - logger.debug( - f"Terminating session {ManagedSparkSession._active_s8s_session_id}" - ) - terminate_s8s_session( - ManagedSparkSession._project_id, - ManagedSparkSession._region, - ManagedSparkSession._active_s8s_session_id, - self._client_options, - recorder=get_recorder(self.instrumentation), - ) - else: - # Client-side cleanup only - logger.debug( - f"Stopping session {ManagedSparkSession._active_s8s_session_id} without termination" - ) + def _stop_server_session(self, terminate: Optional[bool], span) -> None: + if ManagedSparkSession._active_s8s_session_id is None: + return + # Determine if we should terminate the server-side session + if terminate is None: + # Auto-detect: managed sessions terminate, named sessions don't + should_terminate = ( + not ManagedSparkSession._active_session_uses_custom_id + ) + else: + should_terminate = terminate - self._remove_stopped_session_from_file() + span.set( + terminate=should_terminate, + session_id=ManagedSparkSession._active_s8s_session_id, + ) + try: + if should_terminate: + # Terminate the server-side session + logger.debug( + f"Terminating session {ManagedSparkSession._active_s8s_session_id}" + ) + terminate_s8s_session( + ManagedSparkSession._project_id, + ManagedSparkSession._region, + ManagedSparkSession._active_s8s_session_id, + self._client_options, + recorder=get_recorder(self.instrumentation), + ) + else: + # Client-side cleanup only + logger.debug( + f"Stopping session {ManagedSparkSession._active_s8s_session_id} without termination" + ) + finally: + self._remove_stopped_session_from_file() - # Clean up SparkSession._instantiatedSession if it points to this session - try: - from pyspark.sql import SparkSession as PySparkSQLSession + # Clean up SparkSession._instantiatedSession if it points to this session + try: + from pyspark.sql import SparkSession as PySparkSQLSession - if PySparkSQLSession._instantiatedSession is self: - PySparkSQLSession._instantiatedSession = None - logger.debug( - "Cleared SparkSession._instantiatedSession reference" - ) - except (ImportError, AttributeError): - # PySpark not available or _instantiatedSession doesn't exist - pass + if PySparkSQLSession._instantiatedSession is self: + PySparkSQLSession._instantiatedSession = None + logger.debug( + "Cleared SparkSession._instantiatedSession reference" + ) + except (ImportError, AttributeError): + # PySpark not available or _instantiatedSession doesn't exist + pass - ManagedSparkSession._active_s8s_session_uuid = None - ManagedSparkSession._active_s8s_session_id = None - ManagedSparkSession._active_session_uses_custom_id = False - ManagedSparkSession._project_id = None - ManagedSparkSession._region = None - ManagedSparkSession._client_options = None - invalidate_s8s_session_active() - - self.client.close() - if self is ManagedSparkSession._default_session: - ManagedSparkSession._default_session = None - if self is getattr( - ManagedSparkSession._active_session, "session", None - ): - ManagedSparkSession._active_session.session = None + ManagedSparkSession._active_s8s_session_uuid = None + ManagedSparkSession._active_s8s_session_id = None + ManagedSparkSession._active_session_uses_custom_id = False + ManagedSparkSession._project_id = None + ManagedSparkSession._region = None + ManagedSparkSession._client_options = None + invalidate_s8s_session_active() def terminate_s8s_session( @@ -2234,7 +2411,7 @@ class _LivenessEntry(NamedTuple): _LIVENESS_CACHE_LOCK = threading.Lock() -_LIVENESS: dict[tuple[str, Optional[str], Any], _LivenessEntry] = {} +_LIVENESS: dict[tuple[str, Optional[str]], _LivenessEntry] = {} # getOrCreate's GetSession check of the session it is about to reuse. _reuse_check_errors = _OutageLog(logger) @@ -2245,16 +2422,11 @@ def _get_configured_liveness_ttl_seconds() -> Optional[float]: 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(), - ) +) -> tuple[str, Optional[str]]: + return (session_name, _get_client_option(client_options, "api_endpoint")) -def _record_liveness(key: tuple[str, Optional[str], Any], active: bool) -> None: +def _record_liveness(key: tuple[str, Optional[str]], active: bool) -> None: with _LIVENESS_CACHE_LOCK: _LIVENESS[key] = _LivenessEntry(active, _now()) @@ -2270,11 +2442,6 @@ def _check_cached_liveness( negative, never when 0). The guard that runs after a failed RPC always asks. """ - if ( - get_active_s8s_session_response - is not _ORIG_GET_ACTIVE_S8S_SESSION_RESPONSE - ): - return None if _GUARD_MODE.get() == "on_error": return None key = _liveness_cache_key(session_name, client_options) @@ -2328,7 +2495,7 @@ def get_active_s8s_session_response( ).get_session(get_session_request) except NotFound as e: _record_liveness(key, False) - print(f"{session_name} Managed Spark Session deleted: {e}") + logger.debug("%s Managed Spark Session deleted: %s", session_name, e) return None state = get_session_response.state if state == Session.State.ACTIVE: @@ -2340,9 +2507,6 @@ def get_active_s8s_session_response( 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: @@ -2352,9 +2516,6 @@ def is_s8s_session_active(session_name, client_options) -> bool: return True -_ORIG_IS_S8S_SESSION_ACTIVE = is_s8s_session_active - - def make_is_active_callback( session_name: str, client_options: Any, @@ -2374,19 +2535,9 @@ def make_is_active_callback( def check() -> Tuple[bool, bool]: """``(active, cached)`` for one guard invocation.""" - 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, True - 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 - ), False + cached = _check_cached_liveness(session_name, client_options) + if cached is not None: + return cached, True with recorder.span( "dataproc.get_session", reason="rpc_guard", new_client=True ): @@ -2414,5 +2565,4 @@ def is_active() -> bool: ) return active - is_active._supports_error_guard = True # type: ignore[attr-defined] return is_active diff --git a/tests/local/test_sparksql_local.py b/tests/local/test_sparksql_local.py index 90f820f6..03709a1b 100644 --- a/tests/local/test_sparksql_local.py +++ b/tests/local/test_sparksql_local.py @@ -320,8 +320,11 @@ def test_create_local_spark_session(monkeypatch): monkeypatch.setenv("DATAPROC_WORKLOAD_TYPE", "batch") monkeypatch.setenv("SPARK_LOCAL_IP", "127.0.0.1") - if ManagedSparkSession.getActiveSession(): - ManagedSparkSession.getActiveSession().stop() + # The module-scoped `spark` fixture may have left its session active. + # ManagedSparkSession.getActiveSession() only reports managed sessions, + # so ask pyspark itself. + if SparkSession.getActiveSession(): + SparkSession.getActiveSession().stop() local_spark_session = ( SparkSession.builder.master("local") diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 00000000..1a675c07 --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,133 @@ +# Copyright 2026 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 +from unittest import mock + +import pytest + +from google.cloud.managed_spark_connect.client import proxy + + +class FakeCredentials: + """Valid credentials that never need a refresh.""" + + valid = True + token = "unit-test-token" + expiry = None + + def refresh(self, request): + raise AssertionError("FakeCredentials are always valid") + + +@pytest.fixture(autouse=True) +def _fake_application_default_credentials(): + """Unit tests never look up real credentials. + + Opening a channel or a Dataproc client looks up Application Default + Credentials, which would otherwise read the developer's gcloud config + or probe the metadata server. Tests that need other credentials patch + ``google.auth.default`` themselves. + """ + with mock.patch( + "google.auth.default", return_value=(FakeCredentials(), "test-project") + ): + yield + proxy.clear_credentials_cache() + + +def session_response(session_id, uuid): + from google.cloud.dataproc_v1 import Session + + response = Session() + response.name = ( + f"projects/test-project/locations/test-region/sessions/{session_id}" + ) + response.uuid = uuid + response.state = Session.State.ACTIVE + response.runtime_config.version = "3.0" + response.session_template = "projects/p/locations/r/sessionTemplates/t" + response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect-server.example.com:443" + } + response.create_time = "2024-01-01T00:00:00Z" + response.state_time = "2024-01-01T00:01:00Z" + response.state_history = [ + { + "state": Session.State.CREATING, + "state_start_time": "2024-01-01T00:00:00Z", + } + ] + return response + + +@pytest.fixture +def dataproc_mocks(monkeypatch): + """Everything getOrCreate()/stop() touch, without a network.""" + from google.cloud.dataproc_v1 import Session + from pyspark.sql.connect.client.core import ConfigResult + from pyspark.sql.connect.proto import ConfigResponse + + from google.cloud.managed_spark_connect import ManagedSparkSession + + for key in list(os.environ): + if key.startswith("MANAGED_SPARK_CONNECT_"): + monkeypatch.delenv(key) + session_id = "sc-20240702-103952-abcdef" + uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + response = session_response(session_id, uuid) + + with ( + mock.patch("google.auth.default") as auth_default, + mock.patch( + "google.cloud.dataproc_v1.SessionControllerClient" + ) as client_cls, + mock.patch( + "pyspark.sql.connect.client.SparkConnectClient.config", + return_value=ConfigResult.fromProto(ConfigResponse()), + ), + mock.patch.object( + ManagedSparkSession.Builder, + "generate_session_id", + return_value=session_id, + ), + mock.patch( + "google.cloud.managed_spark_connect.session.is_s8s_session_active", + return_value=True, + ), + mock.patch( + "google.cloud.managed_spark_connect.session._server_spark_version", + return_value=None, + ), + mock.patch( + "google.cloud.managed_spark_connect.environment.get_client_environment_label", + return_value="unknown", + ), + mock.patch("time.sleep"), + ): + cred = mock.MagicMock() + cred.token = "token" + auth_default.return_value = (cred, "") + client = client_cls.return_value + operation = mock.Mock() + operation.result.return_value = response + client.create_session.return_value = operation + terminating = Session() + terminating.state = Session.State.TERMINATING + client.get_session.return_value = terminating + yield { + "client": client, + "session_id": session_id, + "uuid": uuid, + "response": response, + } diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index f3c8425e..08006996 100644 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -225,7 +225,6 @@ def trailing_metadata(self): 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 @@ -361,8 +360,6 @@ def trailing_metadata(self): 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") @@ -406,8 +403,6 @@ 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 diff --git a/tests/unit/test_dataproc.py b/tests/unit/test_dataproc.py index 6c4717e2..edf6958b 100644 --- a/tests/unit/test_dataproc.py +++ b/tests/unit/test_dataproc.py @@ -65,10 +65,8 @@ def refresh(self, request): ) + datetime.timedelta(hours=1) -def _set_unpatched_client_cls(monkeypatch, cls): +def _set_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): @@ -119,16 +117,6 @@ def fake_default(scopes=None): 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, ): @@ -284,18 +272,7 @@ def fake_default(scopes=None): 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( +def test_create_session_controller_client_caches_by_options( monkeypatch, ): fake_creds = _FakeRealCredentials(valid=True, token="tok") @@ -315,7 +292,7 @@ def __init__(self, client_options=None, credentials=None): created_clients.append(self) monkeypatch.setattr(proxy_mod.googleauth, "default", fake_default) - _set_unpatched_client_cls(monkeypatch, FakeClient) + _set_client_cls(monkeypatch, FakeClient) opts_a1 = ClientOptions(api_endpoint="us-central1-dataproc.googleapis.com") opts_a2 = ClientOptions(api_endpoint="us-central1-dataproc.googleapis.com") @@ -366,27 +343,6 @@ def __init__(self, client_options=None, credentials=None): 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 @@ -412,7 +368,7 @@ def __init__(self, client_options=None, credentials=None): monkeypatch.setattr(proxy_mod.googleauth, "default", slow_default) monkeypatch.setattr(proxy_mod, "googleauthrequests", mock.Mock()) - _set_unpatched_client_cls(monkeypatch, SlowClient) + _set_client_cls(monkeypatch, SlowClient) barrier = threading.Barrier(16) results_client = [None] * 16 diff --git a/tests/unit/test_grpc_fork_noise.py b/tests/unit/test_grpc_fork_noise.py new file mode 100644 index 00000000..64e43b4e --- /dev/null +++ b/tests/unit/test_grpc_fork_noise.py @@ -0,0 +1,731 @@ +# Copyright 2026 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. +"""Regression tests for gRPC fork-noise suppression, auth pre-warming, and proxy teardown.""" + +from concurrent import futures +import datetime +import gc +import json +import os +import socket +import subprocess +import sys +import textwrap +import threading +from unittest import mock +import warnings + +import grpc +from pyspark.sql.connect import proto +from pyspark.sql.connect.proto import base_pb2_grpc +import pytest +import websockets.exceptions +import websockets.sync.client as websocketclient +import websockets.sync.server as websocketserver + +from google.cloud.managed_spark_connect.client import core +from google.cloud.managed_spark_connect.client import proxy + + +class _PrewarmTrackingCredentials: + """Real-shaped credentials that fork ``true`` on refresh and record thread info.""" + + def __init__(self): + self.valid = False + self.token = None + self.expiry = None + self.refresh_threads = [] + + def refresh(self, _request): + self.refresh_threads.append(threading.current_thread().name) + subprocess.run([sys.executable, "-c", "pass"], check=True) + self.token = "prewarmed-token" + self.valid = True + self.expiry = datetime.datetime.now(datetime.timezone.utc).replace( + tzinfo=None + ) + datetime.timedelta(hours=1) + + +def _start_ws_to_grpc_bridge(grpc_port: int): + """Start a local WebSocket server that bridges hex frames to 127.0.0.1:grpc_port.""" + + def handle_ws(ws_conn): + with socket.create_connection(("127.0.0.1", grpc_port)) as tcp_sock: + tcp_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + tcp_sock.settimeout(0.2) + closed = threading.Event() + + def ws_to_tcp(): + try: + while not closed.is_set(): + msg = ws_conn.recv(timeout=10) + if not msg: + break + tcp_sock.sendall(bytes.fromhex(msg)) + except Exception: + pass + finally: + closed.set() + proxy._shutdown_and_close(tcp_sock) + try: + ws_conn.close_socket() + except Exception: + pass + + def tcp_to_ws(): + try: + while not closed.is_set(): + try: + data = tcp_sock.recv(65536) + except TimeoutError: + continue + if not data: + break + ws_conn.send(data.hex()) + except Exception: + pass + finally: + closed.set() + try: + ws_conn.close_socket() + except Exception: + pass + try: + ws_conn.close() + except Exception: + pass + + t1 = threading.Thread(target=ws_to_tcp, daemon=True) + t2 = threading.Thread(target=tcp_to_ws, daemon=True) + t1.start() + t2.start() + t1.join() + t2.join() + + ws_server = websocketserver.serve(handle_ws, "127.0.0.1", 0) + ws_port = ws_server.socket.getsockname()[1] + server_thread = threading.Thread( + target=ws_server.serve_forever, daemon=True + ) + server_thread.start() + return ws_server, ws_port, server_thread + + +def test_import_sets_default_grpc_verbosity_before_grpc_loads(): + """Importing managed_spark_connect sets GRPC_VERBOSITY=ERROR before grpc loads.""" + env = dict(os.environ) + env.pop("GRPC_VERBOSITY", None) + env.pop("GRPC_ENABLE_FORK_SUPPORT", None) + + code = textwrap.dedent(""" + import builtins + import json + import os + import sys + + assert "grpc" not in sys.modules + grpc_import_env = {} + orig_import = builtins.__import__ + + def tracking_import(name, globals=None, locals=None, fromlist=(), level=0): + if (name == "grpc" or name.startswith("grpc.")) and "verbosity_at_grpc_import" not in grpc_import_env: + grpc_import_env["verbosity_at_grpc_import"] = os.environ.get("GRPC_VERBOSITY") + grpc_import_env["fork_support_at_grpc_import"] = os.environ.get("GRPC_ENABLE_FORK_SUPPORT") + return orig_import(name, globals, locals, fromlist, level) + + builtins.__import__ = tracking_import + import google.cloud.managed_spark_connect # noqa: F401 + + print(json.dumps({ + "verbosity_after": os.environ.get("GRPC_VERBOSITY"), + "fork_support_after": os.environ.get("GRPC_ENABLE_FORK_SUPPORT"), + **grpc_import_env, + })) + """) + proc = subprocess.run( + # Code on stdin, not -c: see _run_python in test_init.py. + [sys.executable, "-"], + input=code, + env=env, + capture_output=True, + text=True, + check=True, + ) + data = json.loads(proc.stdout.strip().splitlines()[-1]) + assert data["verbosity_at_grpc_import"] == "ERROR" + assert data["verbosity_after"] == "ERROR" + assert data["fork_support_at_grpc_import"] is None + assert data["fork_support_after"] is None + + +def test_import_respects_user_provided_grpc_verbosity(): + """User-supplied GRPC_VERBOSITY is preserved via os.environ.setdefault.""" + env = dict(os.environ) + env["GRPC_VERBOSITY"] = "DEBUG" + env["GRPC_ENABLE_FORK_SUPPORT"] = "1" + + code = textwrap.dedent(""" + import json + import os + import google.cloud.managed_spark_connect # noqa: F401 + + print(json.dumps({ + "verbosity": os.environ.get("GRPC_VERBOSITY"), + "fork_support": os.environ.get("GRPC_ENABLE_FORK_SUPPORT"), + })) + """) + proc = subprocess.run( + # Code on stdin, not -c: see _run_python in test_init.py. + [sys.executable, "-"], + input=code, + env=env, + capture_output=True, + text=True, + check=True, + ) + data = json.loads(proc.stdout.strip().splitlines()[-1]) + assert data["verbosity"] == "DEBUG" + assert data["fork_support"] == "1" + + +def test_proxied_channel_prewarms_credentials_on_init(monkeypatch): + """ProxiedChannel.__init__ warms and refreshes credentials before any RPC.""" + from google.cloud.managed_spark_connect.instrumentation import ( + InMemorySink, + Instrumentation, + ) + + proxy.clear_credentials_cache() + creds = _PrewarmTrackingCredentials() + default_calls = [] + + def fake_default(scopes=None): + default_calls.append((threading.current_thread().name, scopes)) + return creds, "test-project" + + monkeypatch.setattr(proxy.googleauth, "default", fake_default) + + inst = Instrumentation([InMemorySink()]) + ch = core.ProxiedChannel("localhost", lambda: True, instrumentation=inst) + try: + # Credentials were discovered and refreshed synchronously on the calling + # thread during ProxiedChannel.__init__, before any RPC connected. + assert len(default_calls) == 1 + assert default_calls[0][0] == threading.current_thread().name + assert creds.refresh_threads == [threading.current_thread().name] + assert creds.valid is True + assert creds.token == "prewarmed-token" + # The initial refresh is traced with the channel's recorder. + assert [ + r["attrs"]["had_token"] + for r in inst.records + if r.get("name") == "auth.token_refreshed" + ] == [False] + + # Subsequent get_refreshed_credentials() in the proxy thread hits the + # warmed cache without re-running google.auth.default or creds.refresh. + reused = proxy.get_refreshed_credentials() + assert reused is creds + assert len(default_calls) == 1 + assert len(creds.refresh_threads) == 1 + finally: + ch.close() + proxy.clear_credentials_cache() + + +def test_proxied_channel_opens_when_credential_prewarm_fails(): + """No credentials yet only skips the pre-warm.""" + from google.auth.exceptions import DefaultCredentialsError + + def no_credentials(*args, **kwargs): + raise DefaultCredentialsError("no credentials") + + with mock.patch("google.auth.default", side_effect=no_credentials): + ch = core.ProxiedChannel("localhost", lambda: True) + try: + assert ch._proxy.port > 0 + finally: + ch.close() + + +def test_session_get_or_create_prewarms_and_stop_cleans_up_proxy(monkeypatch): + """Builder.getOrCreate() pre-warms credentials and session.stop() stops the proxy cleanly.""" + from google.cloud.dataproc_v1 import Session + from pyspark.sql.connect.client.core import ConfigResult + from pyspark.sql.connect.proto import ConfigResponse + + from google.cloud.managed_spark_connect import ManagedSparkSession + from google.cloud.managed_spark_connect import session as session_mod + + for proxy_var in ( + "http_proxy", + "https_proxy", + "HTTP_PROXY", + "HTTPS_PROXY", + "all_proxy", + "ALL_PROXY", + ): + monkeypatch.delenv(proxy_var, raising=False) + + session_mod.clear_dataproc_caches() + creds = _PrewarmTrackingCredentials() + + session_id = "sc-20261019-120000-prewrm" + session_uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + response = Session() + response.name = ( + f"projects/test-project/locations/test-region/sessions/{session_id}" + ) + response.uuid = session_uuid + response.state = Session.State.ACTIVE + response.runtime_config.version = "3.0" + response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect-server.example.com:443" + } + + with warnings.catch_warnings(record=True) as caught_warnings: + warnings.simplefilter("always", ResourceWarning) + with ( + mock.patch( + "google.auth.default", return_value=(creds, "test-project") + ), + mock.patch( + "google.cloud.dataproc_v1.SessionControllerClient" + ) as client_cls, + mock.patch( + "pyspark.sql.connect.client.SparkConnectClient.config", + return_value=ConfigResult.fromProto(ConfigResponse()), + ), + mock.patch.object( + ManagedSparkSession.Builder, + "generate_session_id", + return_value=session_id, + ), + mock.patch( + "google.cloud.managed_spark_connect.session.is_s8s_session_active", + return_value=True, + ), + mock.patch( + "google.cloud.managed_spark_connect.session._server_spark_version", + return_value=None, + ), + mock.patch("time.sleep"), + ): + client_instance = client_cls.return_value + operation = mock.Mock() + operation.result.return_value = response + client_instance.create_session.return_value = operation + terminating = Session() + terminating.state = Session.State.TERMINATING + client_instance.get_session.return_value = terminating + + builder = ManagedSparkSession.builder.projectId( + "test-project" + ).location("test-region") + session = builder.getOrCreate() + try: + assert creds.valid is True + assert creds.token == "prewarmed-token" + assert len(creds.refresh_threads) == 1 + proxy_obj = session.client._channel._proxy + assert proxy_obj._accept_thread is not None + assert proxy_obj._accept_thread.is_alive() + finally: + session.stop() + + assert not proxy_obj._accept_thread.is_alive() + gc.collect() + + resource_warnings = [ + str(w.message) + for w in caught_warnings + if issubclass(w.category, ResourceWarning) + ] + assert resource_warnings == [] + session_mod.clear_dataproc_caches() + + +def test_proxy_stop_and_channel_close_cleanly_teardown_active_connections( + monkeypatch, +): + """ManagedSparkSessionProxy.stop() and ProxiedChannel.close() close all active sockets and threads without ResourceWarnings.""" + for proxy_var in ( + "http_proxy", + "https_proxy", + "HTTP_PROXY", + "HTTPS_PROXY", + "all_proxy", + "ALL_PROXY", + ): + monkeypatch.delenv(proxy_var, raising=False) + + class Service(base_pb2_grpc.SparkConnectServiceServicer): + + def ExecutePlan(self, request, context): + for idx in range(3): + yield proto.ExecutePlanResponse( + session_id=request.session_id, + operation_id=f"{request.operation_id}-{idx}", + ) + + grpc_server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) + base_pb2_grpc.add_SparkConnectServiceServicer_to_server( + Service(), grpc_server + ) + grpc_port = grpc_server.add_insecure_port("127.0.0.1:0") + grpc_server.start() + + ws_server, ws_port, ws_thread = _start_ws_to_grpc_bridge(grpc_port) + real_ws_connect = websocketclient.connect + + def local_ws_connect(uri, **kwargs): + return real_ws_connect( + uri.replace("wss://", "ws://", 1), + proxy=None, + **kwargs, + ) + + creds = _PrewarmTrackingCredentials() + proxy.clear_credentials_cache() + + try: + with warnings.catch_warnings(record=True) as caught_warnings: + warnings.simplefilter("always", ResourceWarning) + with ( + mock.patch.object( + proxy.googleauth, "default", return_value=(creds, "proj") + ), + mock.patch.object( + proxy.websocketclient, + "connect", + side_effect=local_ws_connect, + ), + ): + proxied_channel = core.ProxiedChannel( + f"127.0.0.1:{ws_port}", is_active_callback=lambda: True + ) + stub = base_pb2_grpc.SparkConnectServiceStub(proxied_channel) + call = stub.ExecutePlan( + proto.ExecutePlanRequest( + session_id="s1", operation_id="op-close" + ) + ) + # Read one response so the proxy connection and backend websocket are actively open. + first = next(call) + assert first.session_id == "s1" + + proxy_obj = proxied_channel._proxy + assert len(proxy_obj._active_client_sockets) >= 1 + assert len(proxy_obj._active_backend_sockets) >= 1 + + # Closing the channel while stream is open must cleanly close + # frontend, client, and backend sockets and join proxy threads. + proxied_channel.close() + + assert len(proxy_obj._active_client_sockets) == 0 + assert len(proxy_obj._active_backend_sockets) == 0 + assert ( + proxy_obj._accept_thread is None + or not proxy_obj._accept_thread.is_alive() + ) + for t in list(proxy_obj._conn_threads): + t.join(timeout=5.0) + assert not t.is_alive() + + del call + gc.collect() + + resource_warnings = [ + str(w.message) + for w in caught_warnings + if issubclass(w.category, ResourceWarning) + ] + assert ( + resource_warnings == [] + ), f"Unexpected ResourceWarnings: {resource_warnings}" + finally: + ws_server.shutdown() + ws_thread.join(timeout=2) + grpc_server.stop(grace=None) + proxy.clear_credentials_cache() + + +def test_session_stop_clears_global_refs_even_when_client_close_raises(caplog): + """If client.close() raises during stop(), global/thread-local session refs are still cleared.""" + from google.cloud.managed_spark_connect import ManagedSparkSession + + fake_session = object.__new__(ManagedSparkSession) + fake_session.instrumentation = None + fake_session._client = mock.MagicMock() + fake_session._client.close.side_effect = RuntimeError("close boom") + fake_session._sqm = mock.Mock() + + ManagedSparkSession._default_session = fake_session + ManagedSparkSession._active_session.session = fake_session + ManagedSparkSession._active_s8s_session_id = None + try: + fake_session.stop(terminate=False) + assert ManagedSparkSession._default_session is None + assert ( + getattr(ManagedSparkSession._active_session, "session", None) + is None + ) + assert "Failed to close the Spark Connect client" in [ + r.getMessage() for r in caplog.records if r.levelname == "WARNING" + ] + fake_session._sqm.close.assert_called_once_with() + # The failed close still counts: pyspark's __del__ doesn't retry it. + fake_session.__del__() + fake_session._client.close.assert_called_once_with() + finally: + ManagedSparkSession._default_session = None + ManagedSparkSession._active_session.session = None + + +def _attached_session(): + from google.cloud.managed_spark_connect import ManagedSparkSession + from google.cloud.managed_spark_connect.client import ManagedSparkChannelBuilder + + cb = ManagedSparkChannelBuilder( + "sc://localhost/;session_id=00000000-0000-0000-0000-000000000000;use_ssl=true" + ) + return ManagedSparkSession(connection=cb) + + +def test_session_close_unregisters_atexit_and_allows_gc(monkeypatch): + """ManagedSparkSession.__init__ sets SPARK_CONNECT_MODE_ENABLED=1 and _close_client/stop unregisters the client's pyspark atexit hook so sessions and channels are GC'd.""" + import weakref + + monkeypatch.delenv("SPARK_CONNECT_MODE_ENABLED", raising=False) + + def make_and_stop(): + s = _attached_session() + assert os.environ.get("SPARK_CONNECT_MODE_ENABLED") == "1" + sref = weakref.ref(s) + cref = weakref.ref(s.client) + chref = weakref.ref(s.client._channel) + s._close_client() + return sref, cref, chref + + sref, cref, chref = make_and_stop() + gc.collect() + assert sref() is None + assert cref() is None + assert chref() is None + + +def test_closed_session_is_freed_without_the_cyclic_gc(monkeypatch): + """_close_client breaks the session <-> client cycles (hooks, wrappers, + progress handler), so a stopped session is freed by reference counting + on the thread that drops it, not later by the GC on an arbitrary one.""" + import weakref + + from google.cloud.managed_spark_connect import ManagedSparkSession + + monkeypatch.setattr(ManagedSparkSession, "_active_s8s_session_id", None) + s = _attached_session() + s._register_progress_execution_handler() + s.stop() + sref = weakref.ref(s) + gc.disable() + try: + del s + assert sref() is None + finally: + gc.enable() + + +def test_gc_of_stopped_session_on_release_pool_thread_keeps_pool_usable( + monkeypatch, +): + """Regression: pyspark's SparkSession.__del__ closed the client again. + + When the GC collected a stopped session on a ReleaseExecute pool worker, + that close shut the shared release pool down from one of its own + threads, which fails after marking it shut down but before clearing it, + so every later ReleaseExecute in the process raised "cannot schedule new + futures after shutdown". + """ + import weakref + + from pyspark.sql.connect.client.reattach import ( + ExecutePlanResponseReattachableIterator as ReattachableIterator, + ) + + from google.cloud.managed_spark_connect import ManagedSparkSession + + monkeypatch.setattr(ManagedSparkSession, "_active_s8s_session_id", None) + s = _attached_session() + s.stop() + s._keep_alive_until_gc = s # Only the cyclic GC can free it now. + sref = weakref.ref(s) + del s + + def release_pool(): + """The shared pool, created if needed: by a classmethod up to pyspark + 4.1, by an instance property (of class-level state) from 4.2.""" + get = getattr( + ReattachableIterator, "_get_or_create_release_thread_pool", None + ) + if get is not None: + return get() + return ReattachableIterator._release_thread_pool.fget( + ReattachableIterator + ) + + pool = release_pool() + pool.submit(gc.collect).result(timeout=10) + assert sref() is None + + assert release_pool() is pool + pool.submit(lambda: None).result(timeout=10) + + +def test_stopping_a_non_default_session_leaves_the_default_alone(monkeypatch): + """Stopping a directly attached session must not terminate the default + session's server session or clear its class-level state.""" + from google.cloud.managed_spark_connect import ManagedSparkSession + from google.cloud.managed_spark_connect import session as session_mod + + default = object() + for name, value in ( + ("_default_session", default), + ("_active_s8s_session_id", "sc-default"), + ("_project_id", "p"), + ("_region", "r"), + ): + monkeypatch.setattr(ManagedSparkSession, name, value) + terminate = mock.Mock() + monkeypatch.setattr(session_mod, "terminate_s8s_session", terminate) + + attached = _attached_session() + attached.stop(terminate=True) + + terminate.assert_not_called() + assert attached.client.is_closed + assert ManagedSparkSession._default_session is default + assert ManagedSparkSession._active_s8s_session_id == "sc-default" + assert ManagedSparkSession._project_id == "p" + + +def test_close_client_restores_every_hooked_client_method(monkeypatch): + from google.cloud.managed_spark_connect import ManagedSparkSession + + monkeypatch.setattr(ManagedSparkSession, "_active_s8s_session_id", None) + s = _attached_session() + hooked = ( + "_execute_plan_request_with_metadata", + "_execute", + "_execute_and_fetch_as_iterator", + "_execute_and_fetch", + ) + assert all(name in s.client.__dict__ for name in hooked) + + s.stop() + + assert not [name for name in hooked if name in s.client.__dict__] + + +def test_stop_terminates_even_if_pyspark_default_is_foreign( + dataproc_mocks, monkeypatch +): + """A plain pyspark Connect session created earlier holds pyspark's + default session; the managed session must still own (and terminate) + its own server session.""" + from pyspark.sql.connect.session import SparkSession as ConnectSession + + from google.cloud.managed_spark_connect import ManagedSparkSession + + # Without ManagedSparkSession's own _default_session, this one would be + # inherited, and the managed session would never become the default. + monkeypatch.setattr(ConnectSession, "_default_session", object()) + + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + session.stop() + + dataproc_mocks["client"].terminate_session.assert_called_once() + assert ManagedSparkSession._active_s8s_session_id is None + + +def test_get_or_create_ignores_a_foreign_active_pyspark_session( + dataproc_mocks, monkeypatch +): + """pyspark's per-thread active session is shared with plain pyspark + sessions; the managed builder must not mistake one for its own.""" + from pyspark.sql.connect.session import SparkSession as ConnectSession + + from google.cloud.managed_spark_connect import ManagedSparkSession + + foreign = mock.Mock(spec=["is_stopped"], is_stopped=False) + monkeypatch.setattr( + ConnectSession._active_session, "session", foreign, raising=False + ) + assert ConnectSession.getActiveSession() is foreign + assert ManagedSparkSession.getActiveSession() is None + + client = dataproc_mocks["client"] + stop_response = client.get_session.return_value + session = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + try: + assert isinstance(session, ManagedSparkSession) + # A second getOrCreate() reuses the managed session (still ACTIVE), + # not the foreign one. + client.get_session.return_value = dataproc_mocks["response"] + again = ( + ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + assert again is session + finally: + client.get_session.return_value = stop_response + session.stop() + client.terminate_session.assert_called_once() + + +def test_get_or_create_reusing_the_active_session_makes_it_the_default( + dataproc_mocks, +): + """getOrCreate() reusing the in-process session makes it the default + again, even if the default was cleared, so that its stop() still + terminates the server session and clears the class-level state.""" + from google.cloud.managed_spark_connect import ManagedSparkSession + + def builder(): + return ManagedSparkSession.builder.projectId("test-project").location( + "test-region" + ) + + client = dataproc_mocks["client"] + stop_response = client.get_session.return_value + session = builder().getOrCreate() + try: + ManagedSparkSession._default_session = None + client.get_session.return_value = dataproc_mocks["response"] + again = builder().getOrCreate() + assert again is session + assert ManagedSparkSession._default_session is session + finally: + client.get_session.return_value = stop_response + session.stop() + client.terminate_session.assert_called_once() + assert ManagedSparkSession._active_s8s_session_id is None diff --git a/tests/unit/test_instrumentation.py b/tests/unit/test_instrumentation.py index 3ddd5ccb..c0a445ad 100644 --- a/tests/unit/test_instrumentation.py +++ b/tests/unit/test_instrumentation.py @@ -15,7 +15,6 @@ import json import logging -import os import threading from unittest import mock @@ -735,7 +734,7 @@ def test_guard_failure_is_an_error_on_the_outer_span(proxied): assert rec["error"]["type"] == "RuntimeError" assert rec["attrs"]["guard_ns"] >= 0 assert spans(inst, "grpc.Config") == [rec] # inner never ran - (emitted,) = [r for r in snapshots if r["type"] == "span"] + (emitted,) = [r for r in snapshots if r.get("name") == "grpc.Config"] assert emitted["attrs"]["guard_ns"] >= 0 @@ -922,93 +921,6 @@ def wrapper(*args, **kwargs): # -------------------------------------------------------------------------- -def _session_response(session_id, uuid): - from google.cloud.dataproc_v1 import Session - - response = Session() - response.name = ( - f"projects/test-project/locations/test-region/sessions/{session_id}" - ) - response.uuid = uuid - response.state = Session.State.ACTIVE - response.runtime_config.version = "3.0" - response.session_template = "projects/p/locations/r/sessionTemplates/t" - response.runtime_info.endpoints = { - "Spark Connect Server": "sc://spark-connect-server.example.com:443" - } - response.create_time = "2024-01-01T00:00:00Z" - response.state_time = "2024-01-01T00:01:00Z" - response.state_history = [ - { - "state": Session.State.CREATING, - "state_start_time": "2024-01-01T00:00:00Z", - } - ] - return response - - -@pytest.fixture -def dataproc_mocks(monkeypatch): - """Everything getOrCreate()/stop() touch, without a network.""" - from google.cloud.dataproc_v1 import Session - from pyspark.sql.connect.client.core import ConfigResult - from pyspark.sql.connect.proto import ConfigResponse - - from google.cloud.managed_spark_connect import ManagedSparkSession - - for key in list(os.environ): - if key.startswith("MANAGED_SPARK_CONNECT_"): - monkeypatch.delenv(key) - session_id = "sc-20240702-103952-abcdef" - uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" - response = _session_response(session_id, uuid) - - with ( - mock.patch("google.auth.default") as auth_default, - mock.patch( - "google.cloud.dataproc_v1.SessionControllerClient" - ) as client_cls, - mock.patch( - "pyspark.sql.connect.client.SparkConnectClient.config", - return_value=ConfigResult.fromProto(ConfigResponse()), - ), - mock.patch.object( - ManagedSparkSession.Builder, - "generate_session_id", - return_value=session_id, - ), - mock.patch( - "google.cloud.managed_spark_connect.session.is_s8s_session_active", - return_value=True, - ), - mock.patch( - "google.cloud.managed_spark_connect.session._server_spark_version", - return_value=None, - ), - mock.patch( - "google.cloud.managed_spark_connect.environment.get_client_environment_label", - return_value="unknown", - ), - mock.patch("time.sleep"), - ): - cred = mock.MagicMock() - cred.token = "token" - auth_default.return_value = (cred, "") - client = client_cls.return_value - operation = mock.Mock() - operation.result.return_value = response - client.create_session.return_value = operation - terminating = Session() - terminating.state = Session.State.TERMINATING - client.get_session.return_value = terminating - yield { - "client": client, - "session_id": session_id, - "uuid": uuid, - "response": response, - } - - def _by_name(inst): out = {} for rec in spans(inst): @@ -1365,7 +1277,8 @@ def test_real_grpc_round_trip_through_proxied_channel(inst, grpc_server): finally: channel.close() - assert guard.call_count == 3 + # Three pre-RPC checks, and one after the INTERNAL error. + assert guard.call_count == 4 by_name = _by_name(inst) config = by_name["grpc.Config"] assert [s["attrs"]["layer"] for s in config] == ["inner", "outer"] diff --git a/tests/unit/test_instrumentation_events.py b/tests/unit/test_instrumentation_events.py index 9d49e9c8..21e767c5 100644 --- a/tests/unit/test_instrumentation_events.py +++ b/tests/unit/test_instrumentation_events.py @@ -20,7 +20,6 @@ import contextlib import datetime -import os import socket import threading import time @@ -53,6 +52,12 @@ def one_event(inst, name): return found[0] +def progress_events(inst): + # Sessions also emit auth.* events when their channel pre-warms + # credentials. + return [r for r in events(inst) if r["name"].startswith("progress.")] + + @pytest.fixture def inst(): return Instrumentation([InMemorySink()], {"run_id": "r1"}) @@ -598,79 +603,6 @@ def test_token_refresh_disabled_recorder_still_refreshes(): # -------------------------------------------------------------------------- -def _session_response(session_id, uuid): - from google.cloud.dataproc_v1 import Session - - response = Session() - response.name = ( - f"projects/test-project/locations/test-region/sessions/{session_id}" - ) - response.uuid = uuid - response.state = Session.State.ACTIVE - response.runtime_config.version = "3.0" - response.runtime_info.endpoints = { - "Spark Connect Server": "sc://spark-connect-server.example.com:443" - } - return response - - -@pytest.fixture -def dataproc_mocks(monkeypatch): - """Everything getOrCreate()/stop() touch, without a network.""" - from google.cloud.dataproc_v1 import Session - from pyspark.sql.connect.client.core import ConfigResult - from pyspark.sql.connect.proto import ConfigResponse - - from google.cloud.managed_spark_connect import ManagedSparkSession - - for key in list(os.environ): - if key.startswith("MANAGED_SPARK_CONNECT_"): - monkeypatch.delenv(key) - session_id = "sc-20240702-103952-abcdef" - uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" - response = _session_response(session_id, uuid) - - with ( - mock.patch("google.auth.default") as auth_default, - mock.patch( - "google.cloud.dataproc_v1.SessionControllerClient" - ) as client_cls, - mock.patch( - "pyspark.sql.connect.client.SparkConnectClient.config", - return_value=ConfigResult.fromProto(ConfigResponse()), - ), - mock.patch.object( - ManagedSparkSession.Builder, - "generate_session_id", - return_value=session_id, - ), - mock.patch( - "google.cloud.managed_spark_connect.session.is_s8s_session_active", - return_value=True, - ), - mock.patch( - "google.cloud.managed_spark_connect.session._server_spark_version", - return_value=None, - ), - mock.patch( - "google.cloud.managed_spark_connect.environment.get_client_environment_label", - return_value="unknown", - ), - mock.patch("time.sleep"), - ): - cred = mock.MagicMock() - cred.token = "token" - auth_default.return_value = (cred, "") - client = client_cls.return_value - operation = mock.Mock() - operation.result.return_value = response - client.create_session.return_value = operation - terminating = Session() - terminating.state = Session.State.TERMINATING - client.get_session.return_value = terminating - yield {"client": client, "session_id": session_id, "uuid": uuid} - - class FakeTqdm: """A headless stand-in for ``tqdm`` that records what happened.""" @@ -837,7 +769,7 @@ def test_progress_handler_skips_a_bar_for_a_query_already_done( ) assert forced_progress == [] assert session.progress_bars() == {} - assert not events(inst) + assert not progress_events(inst) def test_progress_handler_ignores_empty_stages_and_no_operation( @@ -850,7 +782,7 @@ def test_progress_handler_ignores_empty_stages_and_no_operation( handler( stages=[_stage(3, 1)], inflight_tasks=0, operation_id=None, done=False ) - assert not events(inst) + assert not progress_events(inst) assert session.progress_bars() == {} @@ -875,7 +807,7 @@ def test_progress_handler_is_off_when_not_interactive_and_not_forced( done=False, ) assert session.progress_bars() == {} - assert not events(inst) + assert not progress_events(inst) finally: ManagedSparkSession._set_tqdm_factory(None) @@ -1070,7 +1002,7 @@ def test_execute_and_fetch_without_a_bar_emits_nothing( session = session_factory(inst) fake_execute_and_fetch["fn"] = lambda client, req: "r" session.client._execute_and_fetch(_execute_request("op-none"), {}) - assert not events(inst) + assert not progress_events(inst) @pytest.mark.parametrize("exc", [ValueError("boom"), KeyboardInterrupt()]) @@ -1265,13 +1197,14 @@ def test_guard_result_records_refresh_error_from_get_session( assert session_mod._check_cached_liveness(SESSION_NAME, None) is True -def test_guard_result_with_patched_is_s8s_session_active(clean_liveness, inst): +def test_guard_result_pre_rpc_asks_when_nothing_is_cached(clean_liveness, inst): session_mod = clean_liveness is_active = session_mod.make_is_active_callback( SESSION_NAME, None, recorder=inst.recorder ) + session_mod.invalidate_s8s_session_active() with mock.patch.object( - session_mod, "is_s8s_session_active", return_value=False + session_mod, "get_active_s8s_session_response", return_value=None ): with _guard_mode("pre_rpc"): assert is_active() is False @@ -1287,6 +1220,7 @@ def test_guard_result_mode_is_direct_outside_the_channel(clean_liveness, inst): is_active = session_mod.make_is_active_callback( SESSION_NAME, None, recorder=inst.recorder ) + session_mod.invalidate_s8s_session_active() with mock.patch.object( session_mod, "get_active_s8s_session_response", return_value=None ): @@ -1314,8 +1248,15 @@ def method(request): return "response" # A guard that has to ask (the RPC path) reports through the channel. + from google.cloud.dataproc_v1 import Session + + active_session = Session() + active_session.state = Session.State.ACTIVE + session_mod.invalidate_s8s_session_active() with mock.patch.object( - session_mod, "is_s8s_session_active", return_value=True + session_mod, + "get_active_s8s_session_response", + return_value=active_session, ): wrapped = core._GuardedMultiCallable(method, guard, mock.Mock()) assert wrapped("req") == "response" @@ -1331,7 +1272,6 @@ def test_guard_disabled_recorder_emits_nothing(clean_liveness): session_mod = clean_liveness is_active = session_mod.make_is_active_callback(SESSION_NAME, None) - assert getattr(is_active, "_supports_error_guard") is True with _guard_mode("pre_rpc"): assert is_active() is True with mock.patch.object( diff --git a/tests/unit/test_logging_hygiene.py b/tests/unit/test_logging_hygiene.py new file mode 100644 index 00000000..cf98f4ca --- /dev/null +++ b/tests/unit/test_logging_hygiene.py @@ -0,0 +1,774 @@ +# Copyright 2026 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 json +import logging +import os +import subprocess +import sys +from unittest import mock +import warnings + +from google.api_core.exceptions import NotFound, ServiceUnavailable +from google.cloud.dataproc_v1 import Session +from google.cloud.managed_spark_connect import ManagedSparkSession +from google.cloud.managed_spark_connect import _logging +from google.cloud.managed_spark_connect import session as session_mod +from google.cloud.managed_spark_connect.session import ( + clear_dataproc_caches, + get_active_s8s_session_response, + invalidate_s8s_session_active, +) +import pytest + + +@pytest.fixture(autouse=True) +def _clean_session_state(monkeypatch): + """Reset global session state and package logger around each test.""" + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "test-project") + monkeypatch.setenv("GOOGLE_CLOUD_REGION", "test-region") + monkeypatch.delenv("MANAGED_SPARK_CONNECT_LOG_LEVEL", raising=False) + monkeypatch.delenv( + "MANAGED_SPARK_CONNECT_SHOW_OPERATION_LINKS", raising=False + ) + monkeypatch.delenv("MANAGED_SPARK_CONNECT_PROGRESS_BAR", raising=False) + monkeypatch.delenv( + "MANAGED_SPARK_CONNECT_PROGRESS_BAR_LEAVE", raising=False + ) + monkeypatch.delenv( + "MANAGED_SPARK_CONNECT_FORCE_PROGRESS_BAR", raising=False + ) + pkg_logger = logging.getLogger("google.cloud.managed_spark_connect") + orig_level = pkg_logger.level + orig_handlers = list(pkg_logger.handlers) + orig_propagate = pkg_logger.propagate + + def reset(): + clear_dataproc_caches() + invalidate_s8s_session_active() + ManagedSparkSession._active_s8s_session_id = None + ManagedSparkSession._active_s8s_session_uuid = None + ManagedSparkSession._active_session_uses_custom_id = False + ManagedSparkSession._project_id = None + ManagedSparkSession._region = None + ManagedSparkSession._default_session = None + ManagedSparkSession._set_tqdm_factory(None) + ManagedSparkSession._execution_progress_bar.clear() + + reset() + yield + reset() + pkg_logger.setLevel(orig_level) + pkg_logger.handlers[:] = orig_handlers + pkg_logger.propagate = orig_propagate + + +def _bare_session( + show_operation_links=None, show_progress_bar=None, progress_bar_leave=None +): + """A ManagedSparkSession with only the display settings initialized.""" + session = object.__new__(ManagedSparkSession) + session.instrumentation = None + session._show_operation_links = show_operation_links + session._show_progress_bar = show_progress_bar + session._progress_bar_leave = progress_bar_leave + return session + + +def test_import_does_not_mutate_root_logger_and_attaches_null_handler(): + """Importing the package in a fresh subprocess must not touch root logger.""" + code = """ +import json +import logging +import warnings + +# Other packages' import-time warnings (google-api-core on old Pythons, +# pyspark on pandas 3) are not ours to silence; this package's still show. +warnings.filterwarnings("ignore", module="(?!google[.]cloud[.]managed_spark_connect)") + +root = logging.getLogger() +before_handlers = len(root.handlers) +before_level = root.level + +import google.cloud.managed_spark_connect # noqa: F401 + +pkg = logging.getLogger("google.cloud.managed_spark_connect") +print(json.dumps({ + "before_handlers": before_handlers, + "after_handlers": len(root.handlers), + "before_level": before_level, + "after_level": root.level, + "pkg_level": pkg.level, + "pkg_handler_types": [type(h).__name__ for h in pkg.handlers], +})) +""" + env = dict(os.environ) + env.pop("MANAGED_SPARK_CONNECT_LOG_LEVEL", None) + proc = subprocess.run( + # Code on stdin, not -c: see _run_python in test_init.py. + [sys.executable, "-"], + input=code, + env=env, + capture_output=True, + text=True, + check=True, + ) + assert proc.stderr == "" + data = json.loads(proc.stdout.strip()) + assert data["before_handlers"] == 0 + assert data["after_handlers"] == 0 + assert data["before_level"] == logging.WARNING + assert data["after_level"] == logging.WARNING + assert data["pkg_level"] == logging.NOTSET + assert "NullHandler" in data["pkg_handler_types"] + assert all(h == "NullHandler" for h in data["pkg_handler_types"]) + + +def test_log_level_env_configures_only_package_logger_in_subprocess(): + """MANAGED_SPARK_CONNECT_LOG_LEVEL configures only package logger, not root.""" + code = """ +import json +import logging + +import google.cloud.managed_spark_connect # noqa: F401 + +root = logging.getLogger() +pkg = logging.getLogger("google.cloud.managed_spark_connect") +print(json.dumps({ + "root_handlers": len(root.handlers), + "root_level": root.level, + "pkg_level": pkg.level, + "pkg_handler_types": [type(h).__name__ for h in pkg.handlers], +})) +""" + env = dict(os.environ) + env["MANAGED_SPARK_CONNECT_LOG_LEVEL"] = "DEBUG" + proc = subprocess.run( + # Code on stdin, not -c: see _run_python in test_init.py. + [sys.executable, "-"], + input=code, + env=env, + capture_output=True, + text=True, + check=True, + ) + data = json.loads(proc.stdout.strip()) + assert data["root_handlers"] == 0 + assert data["root_level"] == logging.WARNING + assert data["pkg_level"] == logging.DEBUG + assert "NullHandler" in data["pkg_handler_types"] + assert "StreamHandler" in data["pkg_handler_types"] + + +def test_configure_package_logger_in_process(monkeypatch): + """configure_package_logger respects MANAGED_SPARK_CONNECT_LOG_LEVEL without touching root.""" + root_logger = logging.getLogger() + root_handlers_before = list(root_logger.handlers) + root_level_before = root_logger.level + + pkg_logger = logging.getLogger("google.cloud.managed_spark_connect") + pkg_logger.handlers[:] = [] + pkg_logger.setLevel(logging.NOTSET) + + monkeypatch.setenv("MANAGED_SPARK_CONNECT_LOG_LEVEL", "DEBUG") + _logging.configure_package_logger() + + assert pkg_logger.level == logging.DEBUG + assert any(isinstance(h, logging.NullHandler) for h in pkg_logger.handlers) + assert root_logger.level == root_level_before + assert list(root_logger.handlers) == root_handlers_before + + +def test_log_level_stream_handler_does_not_propagate(monkeypatch): + """With its own StreamHandler, the package logger doesn't also propagate + to a root handler configured later (no duplicate lines).""" + pkg_logger = logging.getLogger("google.cloud.managed_spark_connect") + pkg_logger.handlers[:] = [] + monkeypatch.setattr(logging.getLogger(), "handlers", []) + monkeypatch.setenv("MANAGED_SPARK_CONNECT_LOG_LEVEL", "INFO") + + _logging.configure_package_logger() + + assert any(type(h) is logging.StreamHandler for h in pkg_logger.handlers) + assert pkg_logger.propagate is False + + +@pytest.mark.parametrize("raw", ["0", "-10", "NOTSET", "loud", ""]) +def test_log_level_ignores_non_positive_and_unknown_levels(monkeypatch, raw): + pkg_logger = logging.getLogger("google.cloud.managed_spark_connect") + pkg_logger.handlers[:] = [] + pkg_logger.setLevel(logging.NOTSET) + monkeypatch.setenv("MANAGED_SPARK_CONNECT_LOG_LEVEL", raw) + + _logging.configure_package_logger() + + assert pkg_logger.level == logging.NOTSET + assert all(isinstance(h, logging.NullHandler) for h in pkg_logger.handlers) + + +def test_non_interactive_lifecycle_zero_stdout_stderr_and_warnings( + capsys, monkeypatch +): + """Non-interactive getOrCreate(), 404/503 liveness checks, actions, and stop() emit 0 bytes and 0 warnings.""" + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive", + lambda: False, + ) + + session_response = Session() + session_response.name = ( + "projects/test-project/locations/test-region/sessions/sc-test-123456" + ) + session_response.uuid = "11111111-2222-3333-4444-555555555555" + session_response.state = Session.State.ACTIVE + session_response.runtime_config.version = "3.0.2" + session_response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect.example.com:443" + } + + mock_operation = mock.Mock() + mock_operation.result.return_value = session_response + + mock_controller = mock.Mock() + mock_controller.create_session.return_value = mock_operation + mock_controller.get_session.return_value = session_response + mock_controller.terminate_session.return_value = mock.Mock() + + cred = mock.MagicMock() + cred.token = "fake-token" + + with ( + mock.patch("google.auth.default", return_value=(cred, "test-project")), + mock.patch( + "google.cloud.dataproc_v1.SessionControllerClient", + return_value=mock_controller, + ), + mock.patch.object( + ManagedSparkSession.Builder, + "generate_session_id", + return_value="sc-test-123456", + ), + # These warn by design when the client's Python or PySpark does not + # match the runtime's, as on most CI lanes (see test_init.py). + mock.patch.object( + ManagedSparkSession.Builder, "_check_python_version_compatibility" + ), + mock.patch.object( + ManagedSparkSession.Builder, "_check_pyspark_version_compatibility" + ), + mock.patch("tqdm.trange") as mock_trange, + warnings.catch_warnings(record=True) as caught_warnings, + ): + warnings.simplefilter("always") + + # 1. Initial getOrCreate() in non-interactive mode + session = ManagedSparkSession.builder.getOrCreate() + mock_trange.assert_not_called() + + # 2. Repeated getOrCreate() returning the active singleton + session2 = ManagedSparkSession.builder.getOrCreate() + assert session2 is session + + # 3. _get_session_by_id on ACTIVE and TERMINATED sessions + assert ( + ManagedSparkSession.builder._get_session_by_id("sc-test-123456") + is session_response + ) + terminated_session = Session() + terminated_session.state = Session.State.TERMINATED + mock_controller.get_session.return_value = terminated_session + assert ( + ManagedSparkSession.builder._get_session_by_id("sc-test-123456") + is None + ) + + # 4. get_active_s8s_session_response on 404 NotFound and 503 ServiceUnavailable + mock_controller.get_session.side_effect = NotFound("Session not found") + assert ( + get_active_s8s_session_response( + "projects/test-project/locations/test-region/sessions/missing-404", + None, + ) + is None + ) + + # A 503 is not a verdict on the session: it propagates (silently). + mock_controller.get_session.side_effect = ServiceUnavailable( + "Transient 503" + ) + with pytest.raises(ServiceUnavailable): + get_active_s8s_session_response( + "projects/test-project/locations/test-region/sessions/missing-503", + None, + ) + + # 5. Simulated DataFrame action (_display_operation_link + progress handler) + session._display_operation_link("op-123") + for handler in list(session.client._progress_handlers): + stage = mock.Mock(num_tasks=4, num_completed_tasks=2) + handler([stage], 2, "op-123", False) + stage_done = mock.Mock(num_tasks=4, num_completed_tasks=4) + handler([stage_done], 0, "op-123", True) + + # 6. Stop session + mock_controller.get_session.side_effect = None + terminating_session = Session() + terminating_session.state = Session.State.TERMINATED + mock_controller.get_session.return_value = terminating_session + session.stop() + + captured = capsys.readouterr() + assert captured.out == "" + assert captured.err == "" + unexpected_warnings = [ + w + for w in caught_warnings + if issubclass(w.category, (UserWarning, ResourceWarning)) + ] + assert unexpected_warnings == [] + + +def test_interactive_attach_prints_once_and_singleton_reuse_is_silent( + capsys, monkeypatch +): + """First-time interactive attach emits 'Using existing Managed Spark Session...', while singleton reuse is silent.""" + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive", + lambda: True, + ) + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive_terminal", + lambda: True, + ) + + session_response = Session() + session_response.name = ( + "projects/test-project/locations/test-region/sessions/sc-existing" + ) + session_response.uuid = "aaaa1111-bbbb-2222-cccc-333333333333" + session_response.state = Session.State.ACTIVE + session_response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect.example.com:443" + } + + fake_session = _bare_session() + ManagedSparkSession._active_s8s_session_id = "sc-existing" + + with ( + mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=session_response, + ), + mock.patch.object( + ManagedSparkSession.Builder, + "_create_spark_connect_session_from_s8s", + return_value=fake_session, + ), + ): + # First-time attach (no in-process active session yet) + ManagedSparkSession._default_session = None + with mock.patch.object( + ManagedSparkSession, "getActiveSession", return_value=None + ): + attached = ManagedSparkSession.builder._get_exiting_active_session() + assert attached is fake_session + first_out = capsys.readouterr() + assert "Using existing Managed Spark Session" in first_out.out + assert "configuration changes may not be applied" in first_out.out + + # Subsequent call when in-process singleton already exists -> must be + # silent, and must not change the live session's display settings. + with mock.patch.object( + ManagedSparkSession, "getActiveSession", return_value=fake_session + ): + reused = ManagedSparkSession.builder.showProgressBar( + False + )._get_exiting_active_session() + assert reused is fake_session + assert fake_session._show_progress_bar is None + second_out = capsys.readouterr() + assert second_out.out == "" + assert second_out.err == "" + + +def test_non_interactive_first_time_attach_is_silent(capsys, monkeypatch): + """Attaching to an existing remote session prints nothing in a script.""" + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive", + lambda: False, + ) + session_response = Session() + session_response.name = ( + "projects/test-project/locations/test-region/sessions/sc-existing" + ) + session_response.state = Session.State.ACTIVE + session_response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect.example.com:443" + } + fake_session = _bare_session() + ManagedSparkSession._active_s8s_session_id = "sc-existing" + + with ( + mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=session_response, + ), + mock.patch.object( + ManagedSparkSession.Builder, + "_create_spark_connect_session_from_s8s", + return_value=fake_session, + ), + mock.patch.object( + ManagedSparkSession, "getActiveSession", return_value=None + ), + ): + attached = ManagedSparkSession.builder._get_exiting_active_session() + + assert attached is fake_session + out = capsys.readouterr() + assert out.out == "" + assert out.err == "" + + +def test_inactive_session_message_escapes_the_session_id(monkeypatch): + """The session ID can come from .sessionId(...), so the HTML escapes it.""" + shown = [] + monkeypatch.setattr( + ManagedSparkSession.Builder, + "_output_element_or_message", + lambda self, plain, html: shown.append(html), + ) + fake_session = mock.Mock(spec=ManagedSparkSession) + ManagedSparkSession._active_s8s_session_id = "x" + + with ( + mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=None, + ), + mock.patch.object( + ManagedSparkSession, "getActiveSession", return_value=fake_session + ), + ): + ManagedSparkSession.builder._get_exiting_active_session() + + (html,) = shown + assert "" not in html + assert "<b>x</b>" in html + + +def test_inactive_session_message_gated_on_interactive(capsys, monkeypatch): + """When active session becomes inactive on server, non-interactive is silent and interactive prints via _output_element_or_message.""" + fake_session = mock.Mock(spec=ManagedSparkSession) + ManagedSparkSession._active_s8s_session_id = "sc-expired" + + with ( + mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=None, + ), + mock.patch.object( + ManagedSparkSession, "getActiveSession", return_value=fake_session + ), + ): + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive", + lambda: False, + ) + res = ManagedSparkSession.builder._get_exiting_active_session() + assert res is None + fake_session.stop.assert_called_once() + out_non_interactive = capsys.readouterr() + assert out_non_interactive.out == "" + assert out_non_interactive.err == "" + + fake_session.stop.reset_mock() + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive", + lambda: True, + ) + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive_terminal", + lambda: True, + ) + res = ManagedSparkSession.builder._get_exiting_active_session() + assert res is None + fake_session.stop.assert_called_once() + out_interactive = capsys.readouterr() + assert ( + "sc-expired Managed Spark Session is not active, stopping and creating a new one" + in out_interactive.out + ) + + +def test_show_operation_links_env_and_builder(capsys, monkeypatch): + """MANAGED_SPARK_CONNECT_SHOW_OPERATION_LINKS and Builder.showOperationLinks control _display_operation_link.""" + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive", + lambda: True, + ) + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive_terminal", + lambda: True, + ) + + session = _bare_session() + ManagedSparkSession._project_id = "test-project" + ManagedSparkSession._region = "test-region" + ManagedSparkSession._active_s8s_session_id = "sc-op-test" + + # Default in interactive mode: prints Spark Query link + session._display_operation_link("op-default") + assert "Spark Query:" in capsys.readouterr().out + + # Disabled via environment variable + monkeypatch.setenv("MANAGED_SPARK_CONNECT_SHOW_OPERATION_LINKS", "0") + session._display_operation_link("op-env-off") + assert capsys.readouterr().out == "" + + # An explicit True (Builder.showOperationLinks(True)) overrides env=0 + session._show_operation_links = True + session._display_operation_link("op-builder-on") + assert "Spark Query:" in capsys.readouterr().out + + # An explicit False disables even when env=true + monkeypatch.setenv("MANAGED_SPARK_CONNECT_SHOW_OPERATION_LINKS", "true") + session._show_operation_links = False + session._display_operation_link("op-builder-off") + assert capsys.readouterr().out == "" + + # Settings only suppress: nothing is shown outside interactive use. + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive", + lambda: False, + ) + session._show_operation_links = True + session._display_operation_link("op-non-interactive") + assert capsys.readouterr().out == "" + + +def test_builder_passes_display_settings_to_the_new_session(): + builder = ( + ManagedSparkSession.Builder() + .showOperationLinks(False) + .showProgressBar(False) + .progressBarLeave(True) + ) + session_response = Session() + session_response.uuid = "aaaa1111-bbbb-2222-cccc-333333333333" + session_response.runtime_info.endpoints = { + "Spark Connect Server": "sc://spark-connect.example.com:443" + } + with ( + mock.patch.object(session_mod, "ManagedSparkSession") as session_cls, + mock.patch.object(session_mod, "ManagedSparkChannelBuilder"), + mock.patch.object(session_mod, "mark_s8s_session_active"), + ): + builder._Builder__connect(session_response, "projects/p/sessions/s") + + kwargs = session_cls.call_args.kwargs + assert kwargs["show_operation_links"] is False + assert kwargs["show_progress_bar"] is False + assert kwargs["progress_bar_leave"] is True + + +def test_progress_bar_and_leave_env_and_builder(monkeypatch): + """MANAGED_SPARK_CONNECT_PROGRESS_BAR, Builder.showProgressBar, and MANAGED_SPARK_CONNECT_PROGRESS_BAR_LEAVE work as expected.""" + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive", + lambda: True, + ) + + created_bars = [] + + class RecordingTqdm: + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.total = kwargs.get("total", 0) + self.n = 0 + self.closed = False + created_bars.append(self) + + def reset(self, total=None): + if total is not None: + self.total = total + + def refresh(self): + pass + + def close(self): + self.closed = True + + ManagedSparkSession._set_tqdm_factory(RecordingTqdm) + + registered_handlers = [] + session = _bare_session() + session.registerProgressHandler = lambda h: registered_handlers.append(h) + session._register_progress_execution_handler() + assert len(registered_handlers) == 1 + handler = registered_handlers[0] + stage = mock.Mock(num_tasks=5, num_completed_tasks=2) + + # 1. Default in interactive mode: bar created with leave=False + handler([stage], 3, "op-1", False) + assert len(created_bars) == 1 + assert created_bars[-1].kwargs["leave"] is False + + # 2. MANAGED_SPARK_CONNECT_PROGRESS_BAR_LEAVE=1 -> leave=True + monkeypatch.setenv("MANAGED_SPARK_CONNECT_PROGRESS_BAR_LEAVE", "1") + handler([stage], 3, "op-2", False) + assert len(created_bars) == 2 + assert created_bars[-1].kwargs["leave"] is True + + # 3. MANAGED_SPARK_CONNECT_PROGRESS_BAR=0 -> no bar created + monkeypatch.setenv("MANAGED_SPARK_CONNECT_PROGRESS_BAR", "0") + handler([stage], 3, "op-3", False) + assert len(created_bars) == 2 + + # 4. Explicit settings (showProgressBar(True), progressBarLeave(False)) + # override the env vars + session._show_progress_bar = True + session._progress_bar_leave = False + handler([stage], 3, "op-4", False) + assert len(created_bars) == 3 + assert created_bars[-1].kwargs["leave"] is False + + # 5. showProgressBar(False) disables the progress bar + session._show_progress_bar = False + handler([stage], 3, "op-5", False) + assert len(created_bars) == 3 + assert session_mod._progress_bar_enabled(False) is False + assert session_mod._progress_bar_enabled(True) is True + + +def test_check_python_version_compatibility_normalizes_patch_versions(): + """_check_python_version_compatibility('3.0.2') normalizes '3.0.2' to '3.0'.""" + builder = ManagedSparkSession.Builder() + + # Matching Python 3.12 -> no warning for "3.0" or "3.0.2" + with mock.patch.object(sys, "version_info", (3, 12, 4, "final", 0)): + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + builder._check_python_version_compatibility("3.0.2") + assert len(w) == 0 + + # Mismatched Python 3.11 -> emits warning for "3.0.2" by normalizing to "3.0" + with mock.patch.object(sys, "version_info", (3, 11, 9, "final", 0)): + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + builder._check_python_version_compatibility("3.0.2") + assert len(w) == 1 + assert "Python version mismatch detected" in str(w[0].message) + assert "Client is using Python 3.11" in str(w[0].message) + assert "runtime 3.0.2 uses Python 3.12" in str(w[0].message) + + +def test_pypi_artifacts_unpinned_package_logs_at_debug_not_info(caplog): + """PyPiArtifacts logs unpinned package recommendation at DEBUG, not INFO.""" + from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts + + with caplog.at_level( + logging.DEBUG, logger="google.cloud.managed_spark_connect" + ): + PyPiArtifacts({"spacy"}) + + matching = [ + r + for r in caplog.records + if "It is recommended to pin the version of the package" in r.message + ] + assert len(matching) == 1 + assert matching[0].levelno == logging.DEBUG + + +def test_proxied_channel_check_active_logs_at_debug_not_warning(caplog): + """ProxiedChannel._check_active logs at DEBUG (not WARNING) before raising RuntimeError.""" + from google.cloud.managed_spark_connect.client.core import ProxiedChannel + + from google.cloud.managed_spark_connect.client import core + + ch = object.__new__(ProxiedChannel) + ch._is_active_callback = lambda: False + ch._guard_errors = core._OutageLog(core.logger) + + with caplog.at_level( + logging.DEBUG, logger="google.cloud.managed_spark_connect" + ): + with pytest.raises(RuntimeError, match="Session not active"): + ch._check_active() + + matching = [ + r for r in caplog.records if "Session is no longer active" in r.message + ] + assert len(matching) == 1 + assert matching[0].levelno == logging.DEBUG + + +def test_display_operation_link_returns_cleanly_when_session_metadata_missing( + capsys, monkeypatch +): + """_display_operation_link returns cleanly instead of raising AssertionError when metadata is None.""" + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive", + lambda: True, + ) + monkeypatch.setattr( + "google.cloud.managed_spark_connect.environment.is_interactive_terminal", + lambda: True, + ) + + session = _bare_session() + ManagedSparkSession._project_id = None + ManagedSparkSession._region = None + ManagedSparkSession._active_s8s_session_id = None + + # Must not raise AssertionError + session._display_operation_link("op-missing-meta") + assert capsys.readouterr().out == "" + + +def test_pyspark_query_context_loggers_do_not_dump_json_to_stderr( + capsys, monkeypatch +): + """SQLQueryContextLogger and DataFrameQueryContextLogger do not dump + duplicate JSON/JVM stacktraces to stderr when AnalysisException is constructed. + """ + monkeypatch.delenv("SPARK_CONNECT_LOG_LEVEL", raising=False) + monkeypatch.delenv("MANAGED_SPARK_CONNECT_LOG_LEVEL", raising=False) + _logging.quiet_pyspark_query_context_loggers() + + from pyspark.logger.logger import PySparkLogger + + for logger_name in ("SQLQueryContextLogger", "DataFrameQueryContextLogger"): + q_logger = PySparkLogger.getLogger(logger_name) + q_logger.exception( + "[TABLE_OR_VIEW_NOT_FOUND] The table `non_existent` cannot be found.", + errorClass="TABLE_OR_VIEW_NOT_FOUND", + ) + + captured = capsys.readouterr() + assert captured.out == "" + assert captured.err == "" + + +def test_session_construction_quiets_pyspark_query_context_loggers(): + with mock.patch.object( + _logging, "quiet_pyspark_query_context_loggers" + ) as quiet: + session = ManagedSparkSession(connection="sc://localhost:15002") + try: + quiet.assert_called_once_with() + finally: + session.stop() diff --git a/tests/unit/test_polling.py b/tests/unit/test_polling.py index 3719cba6..5e640b4b 100644 --- a/tests/unit/test_polling.py +++ b/tests/unit/test_polling.py @@ -39,9 +39,7 @@ def tearDown(self): os.environ.update(self.original_environment) @mock.patch("time.sleep") - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_terminate_s8s_session_stops_polling_after_max_polls( self, mock_client_cls, mock_sleep ): @@ -58,9 +56,7 @@ def test_terminate_s8s_session_stops_polling_after_max_polls( self.assertEqual(mock_sleep.call_count, 179) @mock.patch("time.sleep") - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_terminate_s8s_session_skips_sleep_when_already_terminating( self, mock_client_cls, mock_sleep ): @@ -78,9 +74,7 @@ def test_terminate_s8s_session_skips_sleep_when_already_terminating( mock_sleep.assert_not_called() @mock.patch("time.sleep") - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_terminate_s8s_session_sleeps_only_between_active_polls( self, mock_client_cls, mock_sleep ): diff --git a/tests/unit/test_proxy.py b/tests/unit/test_proxy.py index 6db96ec4..c03bc231 100644 --- a/tests/unit/test_proxy.py +++ b/tests/unit/test_proxy.py @@ -219,7 +219,7 @@ def test_connect_tcp_bridge_spans(): ) inst = Instrumentation([InMemorySink()]) - creds = mock.Mock(token="tok") + creds = mock.Mock(token="tok", valid=False) with ( mock.patch.object( proxy.googleauth, "default", return_value=(creds, None) @@ -264,6 +264,32 @@ def test_forward_connection_ends_span_on_failure(): assert rec["duration_ns"] >= 5_000_000 +def test_bridged_socket_close_abrupt_closes_underlying_ws_socket(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + ws_conn = mock.Mock() + b_sock = proxy.bridged_socket(ws_conn) + + b_sock.close_abrupt() + + assert ws_conn.method_calls == [mock.call.close_socket(), mock.call.close()] + + +def test_bridged_socket_close_abrupt_ignores_errors(): + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + ws_conn = mock.Mock() + ws_conn.close_socket.side_effect = OSError("already closed") + ws_conn.close.side_effect = RuntimeError("closing") + + assert proxy.bridged_socket(ws_conn).close_abrupt() is None + ws_conn.close.assert_called_once_with() + + def test_forward_connection_nests_ws_connect_and_counts(): from unittest import mock @@ -622,6 +648,44 @@ def close(self): assert from_sock.events == shut_then_close +def test_failed_connections_warn_once_until_one_succeeds(caplog): + import logging + from unittest import mock + + from google.cloud.managed_spark_connect.client import proxy + + p = proxy.ManagedSparkSessionProxy(0, "h") + failures = [] + + def fake_forward_connection(*args, on_backend_open=None, **kwargs): + if failures: + raise ConnectionError(failures.pop(0)) + on_backend_open(mock.Mock()) + + def levels(): + out = [ + r.levelno + for r in caplog.records + if r.getMessage().startswith("Proxy connection") + ] + caplog.clear() + return out + + with ( + caplog.at_level(logging.DEBUG, logger=proxy.__name__), + mock.patch.object(proxy, "forward_connection", fake_forward_connection), + ): + failures[:] = ["503", "503"] + p._handle_connection(1, mock.Mock(), None, None) + p._handle_connection(2, mock.Mock(), None, None) + assert levels() == [logging.WARNING, logging.DEBUG] + + p._handle_connection(3, mock.Mock(), None, None) # Succeeds. + failures[:] = ["503"] + p._handle_connection(4, mock.Mock(), None, None) + assert levels() == [logging.WARNING] + + def test_managed_spark_session_proxy_context_manager_exits_cleanly(): from google.cloud.managed_spark_connect.client import proxy diff --git a/tests/unit/test_reattach_patch.py b/tests/unit/test_reattach_patch.py index f29e0241..04104812 100644 --- a/tests/unit/test_reattach_patch.py +++ b/tests/unit/test_reattach_patch.py @@ -11,10 +11,15 @@ # 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 one-shot replay on INVALID_CURSOR.POSITION_NOT_AVAILABLE.""" +"""Unit tests for the one-shot replay on INVALID_CURSOR.POSITION_NOT_AVAILABLE +and for ReleaseExecute coalescing.""" +import gc import logging import os +import threading +import time +import types from unittest import mock import pytest @@ -23,12 +28,18 @@ UnknownException, ) from pyspark.sql.connect import proto +from pyspark.sql.connect.client.reattach import ( + ExecutePlanResponseReattachableIterator, +) +from pyspark.sql.connect.client.retries import Retrying from google.cloud.managed_spark_connect import ManagedSparkSession from google.cloud.managed_spark_connect import session as session_mod +from google.cloud.managed_spark_connect.client import reattach_patch from google.cloud.managed_spark_connect.client.reattach_patch import ( ENV_REPLAY_ON_POSITION_LOST, POSITION_LOST_CONDITION, + install_release_coalescing, install_replay_on_position_lost, ) from google.cloud.managed_spark_connect.instrumentation import ( @@ -104,6 +115,14 @@ def _no_replay_env(monkeypatch): monkeypatch.delenv(ENV_REPLAY_ON_POSITION_LOST, raising=False) +@pytest.fixture(autouse=True) +def _fresh_release_pool(): + """pyspark's release pool is class-level state; recreate it per test.""" + ExecutePlanResponseReattachableIterator._release_thread_pool_instance = None + yield + ExecutePlanResponseReattachableIterator._release_thread_pool_instance = None + + def test_root_plan_replays_once_under_a_new_operation_id(caplog): inst = Instrumentation([InMemorySink()], {"run_id": "r1"}) bars = _BarCloser() @@ -392,6 +411,7 @@ def fake_fetch(self, req, observations, self_destruct=False): .getOrCreate() ) try: + assert session.client._stub in reattach_patch._COALESCED_STUBS req = _root_request("op-first") res = session.client._execute_and_fetch(req, {}) assert res[0] == "ok-table" @@ -438,3 +458,216 @@ def fake_fetch(self, req, observations, self_destruct=False): finally: ManagedSparkSession._set_tqdm_factory(None) ManagedSparkSession._execution_progress_bar.clear() + + +# -------------------------------------------------------------------------- +# ReleaseExecute coalescing +# -------------------------------------------------------------------------- + + +class _StreamStub: + """A SparkConnectService stub serving ``n`` responses and then + ``result_complete``. ``ReleaseExecute`` sets ``entered``, waits for + ``gate`` (open by default), takes ``release_delay_s`` and records + ``"ALL"`` or the ``until`` response id. ``interrupt_after`` raises + KeyboardInterrupt from the stream after that many responses, as a Ctrl-C + does while pyspark is blocked reading it. Methods take extra keyword + arguments such as the ``timeout`` that pyspark 4.3 passes.""" + + def __init__(self, n, release_delay_s=0.0, interrupt_after=None): + self._n = n + self._release_delay_s = release_delay_s + self._interrupt_after = interrupt_after + self._lock = threading.Lock() + self.releases = [] + self.release_all = threading.Event() + self.entered = threading.Event() + self.gate = threading.Event() + self.gate.set() + + def ExecutePlan(self, req, metadata=None, **kwargs): + for i in range(1, self._n + 1): + if i - 1 == self._interrupt_after: + raise KeyboardInterrupt() + yield proto.ExecutePlanResponse( + session_id=req.session_id, + operation_id=req.operation_id, + response_id=f"r-{i:03d}", + ) + done = proto.ExecutePlanResponse( + session_id=req.session_id, + operation_id=req.operation_id, + response_id="r-done", + ) + done.result_complete.CopyFrom( + proto.ExecutePlanResponse.ResultComplete() + ) + yield done + + def ReleaseExecute(self, req, metadata=None, **kwargs): + self.entered.set() + assert self.gate.wait(10) + time.sleep(self._release_delay_s) + with self._lock: + if req.HasField("release_until"): + self.releases.append(req.release_until.response_id) + else: + self.releases.append("ALL") + self.release_all.set() + return proto.ReleaseExecuteResponse() + + def until_ids(self): + with self._lock: + return [r for r in self.releases if r != "ALL"] + + +def _iterator(stub, coalesce=True): + if coalesce: + install_release_coalescing(types.SimpleNamespace(_stub=stub)) + return ExecutePlanResponseReattachableIterator( + _root_request(), stub, lambda: Retrying([]), [] + ) + + +def _wait_idle(it): + """Wait until no coalesced release is scheduled for ``it``.""" + state = it.__dict__.get("_msc_release") + deadline = time.monotonic() + 10 + while state is not None and state.busy: + assert time.monotonic() < deadline, "release drain never finished" + time.sleep(0.01) + + +def _shut_down_release_pool(it): + """Shut down the pool ``it`` sends releases on, once the RPCs queued on + it are done: the class-level pool up to pyspark 4.1, one that 4.2+ keeps + on the iterator itself.""" + it._release_thread_pool.shutdown(wait=True) + ExecutePlanResponseReattachableIterator._release_thread_pool_instance = None + + +def _finish_like_pyspark(it): + """End ``it`` as pyspark's client does: 4.1+ closes it in a finally; + 4.0 just drops it, and its ``__del__`` sends ReleaseExecute(all).""" + if not hasattr(type(it), "__del__"): + it.close() + + +def test_release_until_sends_one_rpc_at_a_time_with_the_latest_id(): + stub = _StreamStub(100, release_delay_s=0.05) + it = _iterator(stub) + + assert len(list(it)) == 101 + + assert stub.release_all.wait(10) + _wait_idle(it) + # pyspark would send 100. At most the first goes out before the stream + # outruns the RPC, then the latest, and result_complete makes the rest + # moot. Exactly which depends on thread scheduling. + assert len(stub.until_ids()) <= 3 + + +def test_release_until_sends_the_latest_pending_id_after_the_rpc_in_flight(): + stub = _StreamStub(10) + stub.gate.clear() + it = _iterator(stub) + + next(it) + assert stub.entered.wait(10) # ReleaseUntil(r-001) is in flight. + for _ in range(4): + next(it) + stub.gate.set() + _wait_idle(it) + + assert stub.until_ids() == ["r-001", "r-005"] + + +def test_iteration_does_not_wait_for_a_slow_release_rpc(): + stub = _StreamStub(50, release_delay_s=1.0) + it = _iterator(stub) + + start = time.monotonic() + assert len(list(it)) == 51 + assert time.monotonic() - start < 0.5 + + +def test_early_break_still_sends_release_all(): + stub = _StreamStub(100, release_delay_s=0.01) + it = _iterator(stub) + + for i, _ in enumerate(it): + if i == 2: + break + _finish_like_pyspark(it) + del it + gc.collect() + + assert stub.release_all.wait(10) + + +def test_keyboard_interrupt_while_reading_still_sends_release_all(): + stub = _StreamStub(100, interrupt_after=3) + + def consume(): + it = _iterator(stub) + try: + for _ in it: + pass + finally: + _finish_like_pyspark(it) + + with pytest.raises(KeyboardInterrupt): + consume() + gc.collect() + + assert stub.release_all.wait(10) + + +def test_other_clients_iterators_use_pyspark_release_until(): + install_release_coalescing(types.SimpleNamespace(_stub=_StreamStub(1))) + stub = _StreamStub(10, release_delay_s=0.05) + it = _iterator(stub, coalesce=False) + + for _ in range(5): + next(it) + _shut_down_release_pool(it) # Waits for the RPCs. + + assert "_msc_release" not in it.__dict__ + # One RPC per response, each its own task on the multi-threaded pool, so + # they can finish in any order. + assert sorted(stub.until_ids()) == [ + "r-001", + "r-002", + "r-003", + "r-004", + "r-005", + ] + + +def test_release_until_tolerates_a_shut_down_pool(): + stub = _StreamStub(10) + it = _iterator(stub) + # The pool is created eagerly up to pyspark 4.1, on first use from 4.2. + pool = it._release_thread_pool + pool.shutdown() + + next(it) # Would raise "cannot schedule new futures after shutdown". + next(it) + + assert stub.until_ids() == [] + assert not it.__dict__["_msc_release"].busy + # pyspark's own _release_all, when `it` is dropped, only skips a pool + # that is gone. + ExecutePlanResponseReattachableIterator._release_thread_pool_instance = None + + +def test_release_failures_are_logged_as_warnings(caplog): + stub = _StreamStub(10) + stub.ReleaseExecute = mock.Mock(side_effect=RuntimeError("boom")) + it = _iterator(stub) + + with caplog.at_level(logging.WARNING, logger=reattach_patch.__name__): + next(it) + _wait_idle(it) + + assert "ReleaseExecute failed with exception: boom." in caplog.messages diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index d4a2d2bc..d6ea70cf 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -1537,20 +1537,47 @@ def test_display_session_link_on_creation_plain_text( "test_session?project=test-project" ) + def _attach_in_interactive_terminal(self): + """No in-process session, and plain-text output (an interactive + terminal) for _get_exiting_active_session.""" + for patcher in ( + mock.patch.object(ManagedSparkSession, "_default_session", None), + mock.patch( + "google.cloud.managed_spark_connect.environment.is_interactive", + return_value=True, + ), + mock.patch( + "google.cloud.managed_spark_connect.environment.is_interactive_terminal", + return_value=True, + ), + ): + patcher.start() + self.addCleanup(patcher.stop) + def _get_existing_session_output(self): ManagedSparkSession._active_s8s_session_id = "test_session" + self._attach_in_interactive_terminal() self.addCleanup( setattr, ManagedSparkSession, "_active_s8s_session_id", None ) + fake_session = mock.Mock() with mock.patch.object( - ManagedSparkSession, "getActiveSession" - ) as mock_get_active_session, mock.patch( - "google.cloud.managed_spark_connect.session.get_active_s8s_session_response" - ) as mock_get_response, mock.patch( + ManagedSparkSession, "getActiveSession", return_value=None + ), mock.patch.object( + ManagedSparkSession.Builder, + "_create_spark_connect_session_from_s8s", + return_value=fake_session, + ), mock.patch( + "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", + return_value=mock.Mock( + state=Session.State.ACTIVE, + runtime_info=mock.Mock( + endpoints={"Spark Connect Server": "sc://localhost:443"} + ), + ), + ), mock.patch( "builtins.print" ) as mock_print: - mock_get_response.return_value = mock.Mock() - mock_get_active_session.return_value = mock.Mock() ManagedSparkSession.builder._get_exiting_active_session() return "\n".join( str(call.args[0]) for call in mock_print.call_args_list @@ -1575,14 +1602,24 @@ def test_get_exiting_active_session_prints_only_text_links(self): def test_get_exiting_active_session_prints_versions(self): self._patch_ide(self._AGY_IDE) self.mock_server_spark_version.return_value = "4.0.2" - response = Session() + response = mock.Mock( + state=Session.State.ACTIVE, + runtime_info=mock.Mock( + endpoints={"Spark Connect Server": "sc://localhost:443"} + ), + ) ManagedSparkSession._active_s8s_session_id = "test_session" + self._attach_in_interactive_terminal() self.addCleanup( setattr, ManagedSparkSession, "_active_s8s_session_id", None ) active = mock.Mock() with mock.patch.object( - ManagedSparkSession, "getActiveSession", return_value=active + ManagedSparkSession, "getActiveSession", return_value=None + ), mock.patch.object( + ManagedSparkSession.Builder, + "_create_spark_connect_session_from_s8s", + return_value=active, ), mock.patch( "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", return_value=response, @@ -1609,7 +1646,7 @@ def test_get_exiting_active_session_connects_before_versions(self): self.mock_server_spark_version.return_value = "4.0.2" response = Session() ManagedSparkSession._active_s8s_session_id = "test_session" - ManagedSparkSession._default_session = None + self._attach_in_interactive_terminal() self.addCleanup( setattr, ManagedSparkSession, "_active_s8s_session_id", None ) @@ -1624,7 +1661,7 @@ def test_get_exiting_active_session_connects_before_versions(self): builder, "_wait_for_session_available", return_value=response ), mock.patch.object( builder, - "_Builder__create_spark_connect_session_from_s8s", + "_create_spark_connect_session_from_s8s", return_value=connected, ), mock.patch( "google.cloud.managed_spark_connect.environment.is_interactive", @@ -3091,9 +3128,7 @@ def test_get_or_create_dataproc_batch_returns_local_session( @mock.patch("google.auth.default") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_custom_session_id_active_with_endpoint_single_rpc( self, mock_session_controller_client, @@ -3135,9 +3170,7 @@ def test_custom_session_id_active_with_endpoint_single_rpc( @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" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_custom_session_id_creating_waits_for_endpoint( self, mock_session_controller_client, @@ -3190,9 +3223,7 @@ def test_custom_session_id_creating_waits_for_endpoint( @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" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_custom_session_id_creating_with_prepopulated_endpoint_waits_until_active( self, mock_session_controller_client, @@ -3248,9 +3279,7 @@ def test_custom_session_id_creating_with_prepopulated_endpoint_waits_until_activ @mock.patch("google.auth.default") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_custom_session_id_switches_between_distinct_active_sessions( self, mock_session_controller_client, @@ -3318,6 +3347,9 @@ def test_custom_session_id_switches_between_distinct_active_sessions( mock_client.terminate_session.assert_not_called() finally: self.stopSession(mock_client, s2) + # s1 is no longer the default session, so this only closes its + # client. + s1.stop() class ServerSparkVersionTest(unittest.TestCase): @@ -3419,9 +3451,7 @@ def test_dataproc_session_id_builder_method(self): builder.sessionId("123-invalid") self.assertIn("Invalid session ID", str(context.exception)) - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_session_reuse_with_custom_id(self, mock_session_controller_client): """Test that sessions are reused when custom ID is provided.""" mock_client = mock_session_controller_client.return_value @@ -3446,9 +3476,7 @@ def test_session_reuse_with_custom_id(self, mock_session_controller_client): mock_client.get_session.assert_called_once() @mock.patch("google.cloud.managed_spark_connect.session.logger") - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_session_delete_terminated_and_warn( self, mock_session_controller_client, mock_logger ): @@ -3485,9 +3513,7 @@ def test_session_delete_terminated_and_warn( self.assertIn("Cloud Logging", warning_msg) @mock.patch("google.cloud.managed_spark_connect.session.logger") - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_session_delete_failed_and_warn( self, mock_session_controller_client, mock_logger ): @@ -3520,9 +3546,7 @@ def test_session_delete_failed_and_warn( @mock.patch("time.sleep", return_value=None) @mock.patch("google.cloud.managed_spark_connect.session.logger") - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) + @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_session_terminating_waits_then_deletes_and_warns( self, mock_session_controller_client, mock_logger, mock_sleep ): @@ -3563,9 +3587,6 @@ def test_session_terminating_waits_then_deletes_and_warns( @pytest.mark.slow @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") - @mock.patch( - "google.cloud.managed_spark_connect.session.SessionControllerClient" - ) @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( "google.cloud.managed_spark_connect.session.is_s8s_session_active" @@ -3576,7 +3597,6 @@ def test_get_or_create_with_terminated_custom_id_deletes_and_recreates( mock_logger, mock_is_s8s_session_active, mock_client_config, - mock_session_controller_client_in_session, mock_dataproc_session_controller_client, mock_credentials, ): @@ -3589,8 +3609,7 @@ def test_get_or_create_with_terminated_custom_id_deletes_and_recreates( cred.token = "token" mock_credentials.return_value = (cred, "") - mock_client = mock_session_controller_client_in_session.return_value - mock_dataproc_session_controller_client.return_value = mock_client + mock_client = mock_dataproc_session_controller_client.return_value terminated_session = Session() terminated_session.state = Session.State.TERMINATED