From cd7d03e58f0288aafa67780e8fa0480e49c57b80 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Mon, 19 Oct 2026 12:45:58 -0700 Subject: [PATCH 1/4] fix(grpc): set default GRPC_VERBOSITY=ERROR, pre-warm credentials, and close proxy sockets cleanly - Set GRPC_VERBOSITY=ERROR (unless already set) at the top of google/cloud/managed_spark_connect/__init__.py when grpc has not been imported yet, since gRPC only reads it on import. User-configured GRPC_VERBOSITY and GRPC_ENABLE_FORK_SUPPORT are left untouched. Documented in the README. - Pre-warm refreshed ADC credentials on the calling thread in ProxiedChannel.__init__(), with the session's recorder, before starting ManagedSparkSessionProxy or the gRPC channel, so subprocess credential helpers (gcloud) do not fork while a gRPC thread holds an active C-core ExecCtx. - Track active client sockets, backend bridged sockets, and worker threads in ManagedSparkSessionProxy; stop() closes all of them and joins the threads against one shared deadline. A backend that connects after stop() is closed and its connection abandoned. Backend sockets are closed with the new bridged_socket.close_abrupt(), which doesn't wait for a WebSocket Close handshake. A failed connection logs a WARNING once, then DEBUG until a connection succeeds again. - Add ManagedSparkSession._close_client(), the one teardown path for stop(): close the StreamingQueryManager and the client (logging a failure as a warning), unregister PySpark's SparkConnectClient._cleanup_ml from atexit, and remove the client hooks, method wrappers and progress handler, so closed sessions and gRPC channels (plus their C-core eventpoll/eventfd FDs) can be garbage-collected instead of remaining pinned for the process lifetime. - Skip pyspark's SparkSession.__del__ once _close_client() has run. Otherwise a stopped session collected by the GC on a ReleaseExecute pool thread shut down pyspark's shared release pool from inside it, leaving the pool unusable for every later query in the process. - Ensure SPARK_CONNECT_MODE_ENABLED=1 is set in ManagedSparkSession.__init__() and Builder.__connect() in addition to Builder.__create(). - Add pre-warm, teardown, and GC regression tests in tests/unit/test_grpc_fork_noise.py, and a tests/unit/conftest.py that keeps unit tests from looking up real credentials. --- README.md | 1 + .../cloud/managed_spark_connect/__init__.py | 7 + .../managed_spark_connect/client/core.py | 38 +- .../managed_spark_connect/client/proxy.py | 175 ++++-- google/cloud/managed_spark_connect/session.py | 181 ++++-- tests/unit/conftest.py | 38 ++ tests/unit/test_grpc_fork_noise.py | 586 ++++++++++++++++++ tests/unit/test_instrumentation.py | 2 +- tests/unit/test_instrumentation_events.py | 14 +- tests/unit/test_proxy.py | 64 ++ 10 files changed, 993 insertions(+), 113 deletions(-) create mode 100644 tests/unit/conftest.py create mode 100644 tests/unit/test_grpc_fork_noise.py diff --git a/README.md b/README.md index 54c4e055..d34b0482 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,7 @@ 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. | ### Reusing Named Sessions Across Notebooks diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index dcee82e2..eaded123 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -13,9 +13,16 @@ # 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") + def _in_ipython() -> bool: # Not environment.is_interactive(), which is also true in a plain diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index b06a9413..b40fcc4d 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -410,29 +410,43 @@ def __init__(self, target_host, is_active_callback, instrumentation=None): ) 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: diff --git a/google/cloud/managed_spark_connect/client/proxy.py b/google/cloud/managed_spark_connect/client/proxy.py index 35361ae4..48f56ad0 100755 --- a/google/cloud/managed_spark_connect/client/proxy.py +++ b/google/cloud/managed_spark_connect/client/proxy.py @@ -180,6 +180,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 +528,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 +556,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 +583,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 +631,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 +674,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 +748,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/session.py b/google/cloud/managed_spark_connect/session.py index 2adfa812..acd6af71 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -658,6 +658,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 @@ -1649,6 +1650,7 @@ def __init__( ``instrumentation`` package). Exposed as ``self.instrumentation``. """ + os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" super().__init__(connection, user_id) self.instrumentation = instrumentation @@ -2089,76 +2091,135 @@ 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 _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 + 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. - span.set( - terminate=should_terminate, - session_id=ManagedSparkSession._active_s8s_session_id, + Unregistering pyspark's atexit hook lets a stopped session be garbage + collected. Removing the client hooks, method wrappers and progress + handler 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 ) - 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" - ) + 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) + for name in ( + "_execute_plan_request_with_metadata", + "_execute", + "_execute_and_fetch_as_iterator", + ): + client.__dict__.pop(name, None) + client.clear_progress_handlers() + self.__dict__.pop("clearProgressHandlers", None) + + def _stop(self, terminate: Optional[bool], span) -> None: + with ManagedSparkSession._lock: + try: + 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 - self._remove_stopped_session_from_file() + 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 - # Clean up SparkSession._instantiatedSession if it points to this session - try: - from pyspark.sql import SparkSession as PySparkSQLSession + 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() - 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 + # Clean up SparkSession._instantiatedSession if it points to this session + try: + from pyspark.sql import SparkSession as PySparkSQLSession - 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 + 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() def terminate_s8s_session( diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 00000000..9cc76556 --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,38 @@ +# 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. +from unittest import mock + +from google.auth.exceptions import DefaultCredentialsError +import pytest + +from google.cloud.managed_spark_connect.client import proxy + + +@pytest.fixture(autouse=True) +def _no_application_default_credentials(): + """Unit tests never look up real credentials. + + Opening a channel pre-warms credentials, which would otherwise read the + developer's gcloud config or probe the metadata server. Tests that need + credentials patch ``google.auth.default`` themselves. + """ + + def no_credentials(*args, **kwargs): + # A new exception each time: a shared instance would keep the + # traceback, and every frame in it, alive between calls. + raise DefaultCredentialsError("no credentials in unit tests") + + with mock.patch("google.auth.default", side_effect=no_credentials): + yield + proxy.clear_credentials_cache() diff --git a/tests/unit/test_grpc_fork_noise.py b/tests/unit/test_grpc_fork_noise.py new file mode 100644 index 00000000..2f91c798 --- /dev/null +++ b/tests/unit/test_grpc_fork_noise.py @@ -0,0 +1,586 @@ +# 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 (the conftest default) only skips the pre-warm.""" + 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) diff --git a/tests/unit/test_instrumentation.py b/tests/unit/test_instrumentation.py index 3ddd5ccb..fcfb829e 100644 --- a/tests/unit/test_instrumentation.py +++ b/tests/unit/test_instrumentation.py @@ -735,7 +735,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 diff --git a/tests/unit/test_instrumentation_events.py b/tests/unit/test_instrumentation_events.py index 9d49e9c8..0d418c68 100644 --- a/tests/unit/test_instrumentation_events.py +++ b/tests/unit/test_instrumentation_events.py @@ -53,6 +53,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"}) @@ -837,7 +843,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 +856,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 +881,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 +1076,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()]) diff --git a/tests/unit/test_proxy.py b/tests/unit/test_proxy.py index 6db96ec4..d869566a 100644 --- a/tests/unit/test_proxy.py +++ b/tests/unit/test_proxy.py @@ -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 From c146c04ad596d884beddd826516f793d10d303a8 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Mon, 19 Oct 2026 12:56:38 -0700 Subject: [PATCH 2/4] fix(reattach): coalesce concurrent ReleaseUntil RPCs pyspark sends one ReleaseExecute(until) per response, each submitted to its shared release thread pool. A fast result stream behind a slow round trip through the proxy queues up hundreds of RPCs that are stale by the time they run. install_release_coalescing(client), called for every ManagedSparkSession, keeps at most one ReleaseExecute(until) in flight per iterator: while one is running, later responses only move the pending release point, and the next RPC sends the latest one. Releases stop once the result is complete, where pyspark's own ReleaseAll takes over. The RPC is built and retried the way pyspark's _release_until does it, and failures are still logged as warnings. A shut-down release pool skips the release instead of failing the caller's iteration. The patch replaces ExecutePlanResponseReattachableIterator._release_until once per process but only changes iterators whose stub belongs to a client it was installed for; plain pyspark clients in the same process keep pyspark's behaviour. ReleaseAll is untouched, so an early break or a KeyboardInterrupt still interrupts the operation on the server. --- README.md | 6 + docs/instrumentation.md | 4 + .../client/reattach_patch.py | 103 ++++++++ google/cloud/managed_spark_connect/session.py | 2 + tests/unit/test_reattach_patch.py | 235 +++++++++++++++++- 5 files changed, 349 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index d34b0482..248c28d6 100644 --- a/README.md +++ b/README.md @@ -178,6 +178,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/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/session.py b/google/cloud/managed_spark_connect/session.py index acd6af71..b38d49f5 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 @@ -1665,6 +1666,7 @@ 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 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 From e5500afdc575002d2685f319306617b97ff3a2c0 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Mon, 19 Oct 2026 12:25:38 -0700 Subject: [PATCH 3/4] fix(logging): remove root basicConfig and gate session/query output on interactivity and builder knobs - Remove logging.basicConfig(level=logging.INFO) from session.py and attach a NullHandler to the google.cloud.managed_spark_connect package logger at import (new _logging module), with optional MANAGED_SPARK_CONNECT_LOG_LEVEL configuration that only configures the package logger. If it attaches its own StreamHandler, the package logger stops propagating, so records are not printed twice once the application configures the root logger. - Gate the session-creation tqdm progress bar thread on interactive/forced progress bar settings so non-interactive sessions spawn no thread and emit no trailing newline. - Demote 'Managed Spark Session deleted' print() statements in get_active_s8s_session_response() and custom-ID session reuse logger.info() calls in _get_session_by_id() to logger.debug(). - Route first-time session attach and inactive-session messages in _get_exiting_active_session() through _output_element_or_message() (with the session ID HTML-escaped) and keep repeated getOrCreate() calls returning the in-process singleton silent. - Add MANAGED_SPARK_CONNECT_SHOW_OPERATION_LINKS / Builder.showOperationLinks(), MANAGED_SPARK_CONNECT_PROGRESS_BAR / Builder.showProgressBar(), and MANAGED_SPARK_CONNECT_PROGRESS_BAR_LEAVE / Builder.progressBarLeave() controls, defaulting execution progress bar leave=False. Links and bars still only appear in interactive environments; the settings can hide them there. They apply to a newly created or attached session, and are documented in the README. - Normalize runtime_version to major.minor before looking up in RUNTIME_PYTHON_MAP in _check_python_version_compatibility(). - Demote unpinned package recommendation in pypi_artifacts.py to DEBUG, demote ProxiedChannel._check_active() inactive-session log to DEBUG before raising RuntimeError, add stacklevel=2 to deprecated package warning in __init__.py, and replace assert all(...) with early return in _display_operation_link(). - When a ManagedSparkSession is created, silence PySpark's SQLQueryContextLogger and DataFrameQueryContextLogger stderr JSON dumps (for the whole process) unless SPARK_CONNECT_LOG_LEVEL or MANAGED_SPARK_CONNECT_LOG_LEVEL is set. - Add comprehensive unit tests in tests/unit/test_logging_hygiene.py. --- README.md | 10 + .../cloud/managed_spark_connect/__init__.py | 7 +- .../cloud/managed_spark_connect/_logging.py | 78 ++ .../managed_spark_connect/client/core.py | 2 +- .../managed_spark_connect/pypi_artifacts.py | 2 +- google/cloud/managed_spark_connect/session.py | 233 ++++-- tests/unit/test_logging_hygiene.py | 775 ++++++++++++++++++ tests/unit/test_session.py | 55 +- 8 files changed, 1086 insertions(+), 76 deletions(-) create mode 100644 google/cloud/managed_spark_connect/_logging.py create mode 100644 tests/unit/test_logging_hygiene.py diff --git a/README.md b/README.md index 248c28d6..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) @@ -135,6 +140,11 @@ These environment variables change how the client behaves. | `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 diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index eaded123..791be7a7 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -23,6 +23,10 @@ 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 @@ -99,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 b40fcc4d..f6c1dd39 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -471,7 +471,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/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 b38d49f5..ca79de08 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -93,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 @@ -257,12 +256,40 @@ def fetch(): #: 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 _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 _resolve_session_controller_client_cls(): return ( SessionControllerClient @@ -485,6 +512,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): @@ -496,6 +526,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.""" @@ -685,6 +744,9 @@ 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 @@ -735,31 +797,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 + if _progress_bar_enabled(self._show_progress_bar): - pbar.close() - # Print new line after the progress bar - print() - - 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" @@ -788,7 +853,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: @@ -802,8 +868,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() ) @@ -823,27 +887,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( @@ -851,6 +906,11 @@ 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" @@ -1063,27 +1123,48 @@ def _get_exiting_active_session( 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() @@ -1405,6 +1486,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 @@ -1545,7 +1629,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) @@ -1633,6 +1717,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. @@ -1649,12 +1736,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), @@ -1760,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 @@ -1773,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 @@ -1805,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", ) @@ -1938,18 +2042,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}/" @@ -2391,7 +2496,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: diff --git a/tests/unit/test_logging_hygiene.py b/tests/unit/test_logging_hygiene.py new file mode 100644 index 00000000..933e418a --- /dev/null +++ b/tests/unit/test_logging_hygiene.py @@ -0,0 +1,775 @@ +# 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, + "_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, + "_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._supports_error_guard = True + 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_session.py b/tests/unit/test_session.py index d4a2d2bc..b53516c8 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, + "_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, + "_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 ) From bff2b7de13f840d7d9af09eb631d0d69d5620ad3 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Thu, 29 Oct 2026 10:54:55 -0700 Subject: [PATCH 4/4] refactor: remove test-detection branches and fix non-singleton stop() Production code no longer checks whether tests patched it: - Drop _ORIG_SESSION_CONTROLLER_CLIENT, _resolve_session_controller_client_cls and the side_effect check: Dataproc clients always come from dataproc_v1.SessionControllerClient and are always cached. Tests patch that name. - Drop the side_effect check in get_default_credentials and the isinstance(valid, bool) check in get_refreshed_credentials: credentials are refreshed when they are not valid or have no token. - Drop _ORIG_GET_ACTIVE_S8S_SESSION_RESPONSE and _ORIG_IS_S8S_SESSION_ACTIVE. The liveness guard asks the cache, then is_s8s_session_active. The cache key no longer includes the client class. - Drop the _supports_error_guard function attribute: the channel runs the liveness guard after a failed RPC for every callback. ManagedSparkSession.stop() on a session that is not the default one (one attached directly, or replaced by a later getOrCreate()) used to terminate the default session's server session and clear its class-level state. It now only closes its own client. ManagedSparkSession now declares its own _default_session and getOrCreate() always sets it, so a plain pyspark session holding pyspark's default can't keep the managed session from owning, and stop() from terminating, its server session. Likewise, ManagedSparkSession.getActiveSession() ignores a plain pyspark session in pyspark's shared per-thread active slot, which getOrCreate() would otherwise have reused as if it were managed. The session's own wrappers of _execute, _execute_and_fetch_as_iterator and _execute_plan_request_with_metadata are installed with install_client_hooks like the others, so closing the session restores all of them. Tests: unit tests get fake valid credentials from tests/unit/conftest.py, which now also holds the shared session_response()/dataproc_mocks fixtures; the Builder's private _create_spark_connect_session_from_s8s is no longer name-mangled; a test session that was never stopped is stopped. --- .../managed_spark_connect/client/core.py | 5 - .../managed_spark_connect/client/proxy.py | 11 +- google/cloud/managed_spark_connect/session.py | 182 ++++++++---------- tests/local/test_sparksql_local.py | 7 +- tests/unit/conftest.py | 119 ++++++++++-- tests/unit/test_core.py | 5 - tests/unit/test_dataproc.py | 52 +---- tests/unit/test_grpc_fork_noise.py | 149 +++++++++++++- tests/unit/test_instrumentation.py | 91 +-------- tests/unit/test_instrumentation_events.py | 90 ++------- tests/unit/test_logging_hygiene.py | 5 +- tests/unit/test_polling.py | 12 +- tests/unit/test_proxy.py | 2 +- tests/unit/test_session.py | 48 ++--- 14 files changed, 382 insertions(+), 396 deletions(-) diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py index f6c1dd39..73a48106 100644 --- a/google/cloud/managed_spark_connect/client/core.py +++ b/google/cloud/managed_spark_connect/client/core.py @@ -401,9 +401,6 @@ 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 @@ -452,8 +449,6 @@ 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: diff --git a/google/cloud/managed_spark_connect/client/proxy.py b/google/cloud/managed_spark_connect/client/proxy.py index 48f56ad0..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: diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index ca79de08..4702ef28 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -249,7 +249,6 @@ 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] = {} @@ -290,14 +289,6 @@ def _progress_bar_enabled(explicit: Optional[bool]) -> bool: ) and _display_flag(explicit, PROGRESS_BAR_ENV, True) -def _resolve_session_controller_client_cls(): - return ( - SessionControllerClient - if SessionControllerClient is not _ORIG_SESSION_CONTROLLER_CLIENT - else dataproc_v1.SessionControllerClient - ) - - def _get_client_option(client_options: Any, name: str) -> Any: if client_options is None: return None @@ -318,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") @@ -482,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`). @@ -708,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 @@ -752,7 +742,16 @@ def __connect(self, session_response, session_name): # 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 @@ -915,7 +914,7 @@ def create_session_pbar(): 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 @@ -1120,7 +1119,7 @@ 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( @@ -1222,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 @@ -1762,46 +1766,33 @@ def __init__( 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_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 - 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 + return _execute_plan_request_with_metadata - self.client._execute_plan_request_with_metadata = ( - execute_plan_request_wrapped_method - ) + 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) - 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) + return _execute - self.client._execute = MethodType(execute_wrapped_method, self.client) + 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) - 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 - ) - - 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): @@ -1815,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 @@ -1984,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. @@ -2215,9 +2225,10 @@ 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, method wrappers and progress - handler breaks the session <-> client reference cycles, so that - mostly happens by reference counting rather than in the cyclic GC. + 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. @@ -2248,19 +2259,17 @@ def _close_client(self) -> None: if self.instrumentation is not None: uninstall_client_wrappers(client) uninstall_client_hooks(client) - for name in ( - "_execute_plan_request_with_metadata", - "_execute", - "_execute_and_fetch_as_iterator", - ): - client.__dict__.pop(name, None) client.clear_progress_handlers() self.__dict__.pop("clearProgressHandlers", None) def _stop(self, terminate: Optional[bool], span) -> None: with ManagedSparkSession._lock: try: - self._stop_server_session(terminate, span) + # 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: @@ -2402,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) @@ -2413,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()) @@ -2438,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) @@ -2508,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: @@ -2520,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, @@ -2542,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 ): @@ -2582,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 index 9cc76556..1a675c07 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -11,28 +11,123 @@ # 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 -from google.auth.exceptions import DefaultCredentialsError 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 _no_application_default_credentials(): +def _fake_application_default_credentials(): """Unit tests never look up real credentials. - Opening a channel pre-warms credentials, which would otherwise read the - developer's gcloud config or probe the metadata server. Tests that need - credentials patch ``google.auth.default`` themselves. + 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. """ - - def no_credentials(*args, **kwargs): - # A new exception each time: a shared instance would keep the - # traceback, and every frame in it, alive between calls. - raise DefaultCredentialsError("no credentials in unit tests") - - with mock.patch("google.auth.default", side_effect=no_credentials): + 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 index 2f91c798..64e43b4e 100644 --- a/tests/unit/test_grpc_fork_noise.py +++ b/tests/unit/test_grpc_fork_noise.py @@ -245,8 +245,14 @@ def fake_default(scopes=None): def test_proxied_channel_opens_when_credential_prewarm_fails(): - """No credentials yet (the conftest default) only skips the pre-warm.""" - ch = core.ProxiedChannel("localhost", lambda: True) + """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: @@ -584,3 +590,142 @@ def release_pool(): 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 fcfb829e..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 @@ -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 0d418c68..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 @@ -604,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.""" @@ -1271,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 @@ -1293,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 ): @@ -1320,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" @@ -1337,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 index 933e418a..cf98f4ca 100644 --- a/tests/unit/test_logging_hygiene.py +++ b/tests/unit/test_logging_hygiene.py @@ -372,7 +372,7 @@ def test_interactive_attach_prints_once_and_singleton_reuse_is_silent( ), mock.patch.object( ManagedSparkSession.Builder, - "_Builder__create_spark_connect_session_from_s8s", + "_create_spark_connect_session_from_s8s", return_value=fake_session, ), ): @@ -426,7 +426,7 @@ def test_non_interactive_first_time_attach_is_silent(capsys, monkeypatch): ), mock.patch.object( ManagedSparkSession.Builder, - "_Builder__create_spark_connect_session_from_s8s", + "_create_spark_connect_session_from_s8s", return_value=fake_session, ), mock.patch.object( @@ -701,7 +701,6 @@ def test_proxied_channel_check_active_logs_at_debug_not_warning(caplog): ch = object.__new__(ProxiedChannel) ch._is_active_callback = lambda: False - ch._supports_error_guard = True ch._guard_errors = core._OutageLog(core.logger) with caplog.at_level( 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 d869566a..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) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index b53516c8..d6ea70cf 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -1565,7 +1565,7 @@ def _get_existing_session_output(self): ManagedSparkSession, "getActiveSession", return_value=None ), mock.patch.object( ManagedSparkSession.Builder, - "_Builder__create_spark_connect_session_from_s8s", + "_create_spark_connect_session_from_s8s", return_value=fake_session, ), mock.patch( "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", @@ -1618,7 +1618,7 @@ def test_get_exiting_active_session_prints_versions(self): ManagedSparkSession, "getActiveSession", return_value=None ), mock.patch.object( ManagedSparkSession.Builder, - "_Builder__create_spark_connect_session_from_s8s", + "_create_spark_connect_session_from_s8s", return_value=active, ), mock.patch( "google.cloud.managed_spark_connect.session.get_active_s8s_session_response", @@ -1661,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", @@ -3128,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, @@ -3172,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, @@ -3227,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, @@ -3285,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, @@ -3355,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): @@ -3456,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 @@ -3483,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 ): @@ -3522,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 ): @@ -3557,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 ): @@ -3600,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" @@ -3613,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, ): @@ -3626,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