From f36658a19419113cb433194e8c0ef7a2c85b1df0 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:41:10 -0700 Subject: [PATCH 1/8] feat: rename Dataproc Spark Connect to Managed Spark Connect Renames the package, modules, and public classes from Dataproc-branded names to Managed Spark Connect (DataprocSparkSession -> ManagedSparkSession, dataproc_spark_connect -> managed_spark_connect, dataproc_magics -> managed_spark_magics, PyPI package dataproc-spark-connect -> managed-spark-connect). The sessionTemplate() builder method is renamed to runtimeProfile(), and the DATAPROC_SPARK_CONNECT_* env vars gain MANAGED_SPARK_CONNECT_* equivalents. Old import paths, class names, sessionTemplate(), and env var names are kept as deprecated aliases that emit a DeprecationWarning, so existing integrations keep working. References to the actual underlying GCP Dataproc API (google-cloud-dataproc, dataproc_v1, dataprocSessionConfig/ dataprocSessionId, dataproc.googleapis.com) are left unchanged since they name real external resources, not this library's branding. --- .env.example | 27 +- .github/workflows/integration-tests.yaml | 8 +- DEVELOPING.md | 6 +- README.md | 58 +- cloudbuild/cloudbuild.yaml | 4 +- google/cloud/dataproc_magics/__init__.py | 8 + google/cloud/dataproc_magics/magics.py | 65 +- .../cloud/dataproc_spark_connect/__init__.py | 23 +- .../dataproc_spark_connect/client/__init__.py | 1 + .../dataproc_spark_connect/client/core.py | 133 +- .../dataproc_spark_connect/client/proxy.py | 282 +--- .../dataproc_spark_connect/environment.py | 179 +- .../dataproc_spark_connect/exceptions.py | 18 +- .../dataproc_spark_connect/pypi_artifacts.py | 49 +- .../cloud/dataproc_spark_connect/session.py | 1373 +--------------- .../cloud/managed_spark_connect/__init__.py | 30 + .../managed_spark_connect/client/__init__.py | 14 + .../managed_spark_connect/client/core.py | 141 ++ .../managed_spark_connect/client/proxy.py | 269 +++ .../managed_spark_connect/environment.py | 190 +++ .../cloud/managed_spark_connect/exceptions.py | 27 + .../managed_spark_connect/pypi_artifacts.py | 48 + google/cloud/managed_spark_connect/session.py | 1440 +++++++++++++++++ google/cloud/managed_spark_magics/__init__.py | 19 + google/cloud/managed_spark_magics/magics.py | 76 + setup.py | 6 +- .../__init__.py | 0 .../test_magics.py | 22 +- tests/integration/test_session.py | 142 +- .../__init__.py | 0 .../test_magics.py | 18 +- tests/unit/test_deprecated_shims.py | 75 + tests/unit/test_environment.py | 158 +- tests/unit/test_init.py | 34 +- tests/unit/test_proxy.py | 2 +- tests/unit/test_pypi_artifacts.py | 2 +- tests/unit/test_session.py | 467 +++--- 37 files changed, 2906 insertions(+), 2508 deletions(-) mode change 100755 => 100644 google/cloud/dataproc_spark_connect/client/proxy.py create mode 100644 google/cloud/managed_spark_connect/__init__.py create mode 100644 google/cloud/managed_spark_connect/client/__init__.py create mode 100644 google/cloud/managed_spark_connect/client/core.py create mode 100755 google/cloud/managed_spark_connect/client/proxy.py create mode 100644 google/cloud/managed_spark_connect/environment.py create mode 100644 google/cloud/managed_spark_connect/exceptions.py create mode 100644 google/cloud/managed_spark_connect/pypi_artifacts.py create mode 100644 google/cloud/managed_spark_connect/session.py create mode 100644 google/cloud/managed_spark_magics/__init__.py create mode 100644 google/cloud/managed_spark_magics/magics.py rename tests/integration/{dataproc_magics => managed_spark_magics}/__init__.py (100%) rename tests/integration/{dataproc_magics => managed_spark_magics}/test_magics.py (89%) rename tests/unit/{dataproc_magics => managed_spark_magics}/__init__.py (100%) rename tests/unit/{dataproc_magics => managed_spark_magics}/test_magics.py (84%) create mode 100644 tests/unit/test_deprecated_shims.py diff --git a/.env.example b/.env.example index a3047e70..26e6520f 100644 --- a/.env.example +++ b/.env.example @@ -1,4 +1,4 @@ -# Google Cloud Configuration for Dataproc Spark Connect Integration Tests +# Google Cloud Configuration for Managed Spark Connect Integration Tests # Copy this file to .env and fill in your actual values # ============================================================================ @@ -8,7 +8,7 @@ # Your Google Cloud Project ID GOOGLE_CLOUD_PROJECT="your-project-id" -# Google Cloud Region where Dataproc sessions will be created +# Google Cloud Region where Managed Spark sessions will be created GOOGLE_CLOUD_REGION="us-central1" # Path to service account key file (if using SERVICE_ACCOUNT auth) @@ -19,35 +19,35 @@ GOOGLE_APPLICATION_CREDENTIALS="/path/to/your/service-account-key.json" # ============================================================================ # Authentication type (SERVICE_ACCOUNT or END_USER_CREDENTIALS). If not set, API default is used. -# DATAPROC_SPARK_CONNECT_AUTH_TYPE="SERVICE_ACCOUNT" -# DATAPROC_SPARK_CONNECT_AUTH_TYPE="END_USER_CREDENTIALS" +# MANAGED_SPARK_CONNECT_AUTH_TYPE="SERVICE_ACCOUNT" +# MANAGED_SPARK_CONNECT_AUTH_TYPE="END_USER_CREDENTIALS" # Service account email for workload authentication (optional) -# DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT="your-service-account@your-project.iam.gserviceaccount.com" +# MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT="your-service-account@your-project.iam.gserviceaccount.com" # ============================================================================ # SESSION CONFIGURATION # ============================================================================ # Session timeout in seconds (how long session stays active) -# DATAPROC_SPARK_CONNECT_TTL_SECONDS="3600" +# MANAGED_SPARK_CONNECT_TTL_SECONDS="3600" # Session idle timeout in seconds (how long session stays active when idle) -# DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS="900" +# MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS="900" # Automatically terminate session when Python process exits (true/false) -# DATAPROC_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT="false" +# MANAGED_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT="false" # Custom file path for storing active session information -# DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH="/tmp/dataproc_spark_connect_session" +# MANAGED_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH="/tmp/managed_spark_connect_session" # ============================================================================ # DATA SOURCE CONFIGURATION # ============================================================================ # Default data source for Spark SQL (currently only supports "bigquery") -# Only available for Dataproc runtime version 2.3 -# DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE="bigquery" +# Only available for Managed Spark runtime version 2.3 +# MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE="bigquery" # ============================================================================ # ADVANCED CONFIGURATION @@ -56,8 +56,7 @@ GOOGLE_APPLICATION_CREDENTIALS="/path/to/your/service-account-key.json" # Custom Dataproc API endpoint (uncomment if needed) # GOOGLE_CLOUD_DATAPROC_API_ENDPOINT="your-region-dataproc.googleapis.com" -# Subnet URI for Dataproc Spark Connect (full resource name format) +# Subnet URI for Managed Spark Connect (full resource name format) # Example: projects/your-project-id/regions/us-central1/subnetworks/your-subnet-name -# DATAPROC_SPARK_CONNECT_SUBNET="projects/your-project-id/regions/us-central1/subnetworks/your-subnet-name" - +# MANAGED_SPARK_CONNECT_SUBNET="projects/your-project-id/regions/us-central1/subnetworks/your-subnet-name" diff --git a/.github/workflows/integration-tests.yaml b/.github/workflows/integration-tests.yaml index b23a0067..8f5b7627 100644 --- a/.github/workflows/integration-tests.yaml +++ b/.github/workflows/integration-tests.yaml @@ -17,7 +17,7 @@ # Required GitHub Secrets: # - GCP_SA_KEY: Service account JSON key (project_id and client_email extracted automatically) # - GCP_REGION: Google Cloud Region (optional, defaults to us-central1) -# - GCP_SUBNET: Dataproc subnet URI +# - GCP_SUBNET: Managed Spark subnet URI # # See INTEGRATION_TESTS.md for setup instructions. @@ -71,10 +71,10 @@ jobs: CI: "true" # Extract from service account JSON automatically GOOGLE_CLOUD_PROJECT: ${{ fromJson(secrets.GCP_SA_KEY).project_id }} - DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT: ${{ fromJson(secrets.GCP_SA_KEY).client_email }} + MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT: ${{ fromJson(secrets.GCP_SA_KEY).client_email }} # Infrastructure-specific secrets GOOGLE_CLOUD_REGION: ${{ secrets.GCP_REGION || 'us-central1' }} - DATAPROC_SPARK_CONNECT_SUBNET: ${{ secrets.GCP_SUBNET }} - DATAPROC_SPARK_CONNECT_AUTH_TYPE: "SERVICE_ACCOUNT" + MANAGED_SPARK_CONNECT_SUBNET: ${{ secrets.GCP_SUBNET }} + MANAGED_SPARK_CONNECT_AUTH_TYPE: "SERVICE_ACCOUNT" run: | python -m pytest tests/integration/ -v --tb=short -x \ No newline at end of file diff --git a/DEVELOPING.md b/DEVELOPING.md index a1de08f4..c9a8a5ea 100644 --- a/DEVELOPING.md +++ b/DEVELOPING.md @@ -35,7 +35,7 @@ configuration details on the command line. For example: env \ GOOGLE_CLOUD_PROJECT='project-id' \ GOOGLE_CLOUD_REGION='us-central1' \ - DATAPROC_SPARK_CONNECT_SUBNET='subnet-id' \ + MANAGED_SPARK_CONNECT_SUBNET='subnet-id' \ pytest --tb=auto -v ``` @@ -70,7 +70,7 @@ use. This will be set automatically if you set it to `auto`. For example: env \ GOOGLE_CLOUD_PROJECT='project-id' \ GOOGLE_CLOUD_REGION='us-central1' \ - DATAPROC_SPARK_CONNECT_SUBNET='subnet-id' \ - DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT='service@account.test' \ + MANAGED_SPARK_CONNECT_SUBNET='subnet-id' \ + MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT='service@account.test' \ pytest -n auto --tb=auto -v ``` diff --git a/README.md b/README.md index a746e132..7d81c692 100644 --- a/README.md +++ b/README.md @@ -1,26 +1,26 @@ -# Dataproc Spark Connect Client +# Managed Spark Connect Client A wrapper of the Apache [Spark Connect](https://spark.apache.org/spark-connect/) client with additional functionalities that allow applications to communicate -with a remote Dataproc Spark Session using the Spark Connect protocol without +with a remote Managed Spark Session using the Spark Connect protocol without requiring additional steps. ## Install ```sh -pip install dataproc_spark_connect +pip install managed_spark_connect ``` ## Uninstall ```sh -pip uninstall dataproc_spark_connect +pip uninstall managed_spark_connect ``` ## Setup This client requires permissions to -manage [Dataproc Sessions and Session Templates](https://cloud.google.com/dataproc-serverless/docs/concepts/iam). +manage [Managed Spark Sessions and Runtime Profiles](https://cloud.google.com/dataproc-serverless/docs/concepts/iam). If you are running the client outside of Google Cloud, you need to provide authentication credentials. Set the `GOOGLE_APPLICATION_CREDENTIALS` environment @@ -36,42 +36,42 @@ in your code using the builder API: ## Usage -1. Install the latest version of Dataproc Spark Connect: +1. Install the latest version of Managed Spark Connect: ```sh - pip install -U dataproc-spark-connect + pip install -U managed-spark-connect ``` 2. Add the required imports into your PySpark application or notebook and start a Spark session using the fluent API: ```python - from google.cloud.dataproc_spark_connect import DataprocSparkSession - spark = DataprocSparkSession.builder.getOrCreate() + from google.cloud.managed_spark_connect import ManagedSparkSession + spark = ManagedSparkSession.builder.getOrCreate() ``` 3. You can configure Spark properties using the `.config()` method: ```python - from google.cloud.dataproc_spark_connect import DataprocSparkSession - spark = DataprocSparkSession.builder.config('spark.executor.memory', '4g').config('spark.executor.cores', '2').getOrCreate() + from google.cloud.managed_spark_connect import ManagedSparkSession + spark = ManagedSparkSession.builder.config('spark.executor.memory', '4g').config('spark.executor.cores', '2').getOrCreate() ``` 4. For advanced configuration, you can use the `Session` class to customize settings like subnetwork or other environment configurations: ```python - from google.cloud.dataproc_spark_connect import DataprocSparkSession + from google.cloud.managed_spark_connect import ManagedSparkSession from google.cloud.dataproc_v1 import Session session_config = Session() session_config.environment_config.execution_config.subnetwork_uri = '' session_config.runtime_config.version = '3.0' - spark = DataprocSparkSession.builder.projectId('my-project').location('us-central1').dataprocSessionConfig(session_config).getOrCreate() + spark = ManagedSparkSession.builder.projectId('my-project').location('us-central1').dataprocSessionConfig(session_config).getOrCreate() ``` ### Builder Configuration -The `DataprocSparkSession.builder` provides a fluent API to configure the session. Below is a list of available methods: +The `ManagedSparkSession.builder` provides a fluent API to configure the session. Below is a list of available methods: | Method | Description | |--------|-------------| @@ -83,9 +83,9 @@ The `DataprocSparkSession.builder` provides a fluent API to configure the sessio | `labels(labels)` | Adds multiple labels to the session. | | `location(location)` | Sets the Google Cloud region. | | `projectId(project_id)` | Sets the Google Cloud project ID. | -| `runtimeVersion(version)` | Sets the Dataproc runtime version (e.g., "3.0"). | +| `runtimeProfile(profile)` | Sets the Runtime Profile to use. | +| `runtimeVersion(version)` | Sets the Managed Spark runtime version (e.g., "3.0"). | | `serviceAccount(account)` | Sets the service account for the session. | -| `sessionTemplate(template)` | Sets the session template to use. | | `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. | @@ -98,9 +98,9 @@ To create or connect to a named session: 1. Create a session with a custom ID in your first notebook: ```python - from google.cloud.dataproc_spark_connect import DataprocSparkSession + from google.cloud.managed_spark_connect import ManagedSparkSession session_id = 'my-ml-pipeline-session' - spark = DataprocSparkSession.builder.dataprocSessionId(session_id).getOrCreate() + spark = ManagedSparkSession.builder.dataprocSessionId(session_id).getOrCreate() df = spark.createDataFrame([(1, 'data')], ['id', 'value']) df.show() ``` @@ -108,9 +108,9 @@ To create or connect to a named session: 2. Reuse the same session in another notebook by specifying the same session ID: ```python - from google.cloud.dataproc_spark_connect import DataprocSparkSession + from google.cloud.managed_spark_connect import ManagedSparkSession session_id = 'my-ml-pipeline-session' - spark = DataprocSparkSession.builder.dataprocSessionId(session_id).getOrCreate() + spark = ManagedSparkSession.builder.dataprocSessionId(session_id).getOrCreate() df = spark.createDataFrame([(2, 'more-data')], ['id', 'value']) df.show() ``` @@ -127,7 +127,7 @@ The package supports the [sparksql-magic](https://github.com/cryeo/sparksql-magi **Installation**: To use magic commands, install the required dependencies manually: ```bash -pip install dataproc-spark-connect +pip install managed-spark-connect pip install IPython sparksql-magic ``` @@ -163,9 +163,21 @@ Available options: See [sparksql-magic](https://github.com/cryeo/sparksql-magic) for more examples. -**Note**: Magic commands are optional. If you only need basic DataprocSparkSession functionality without Jupyter magic support, install only the base package: +**Note**: Magic commands are optional. If you only need basic ManagedSparkSession functionality without Jupyter magic support, install only the base package: ```bash -pip install dataproc-spark-connect +pip install managed-spark-connect +``` + +## Migrating from dataproc-spark-connect + +The `dataproc-spark-connect` package and the `google.cloud.dataproc_spark_connect` module have been renamed to `managed-spark-connect` / `google.cloud.managed_spark_connect`, and `DataprocSparkSession` has been renamed to `ManagedSparkSession`. The old import path and class name still work but emit a `DeprecationWarning` — update your imports when convenient: + +```python +# Before +from google.cloud.dataproc_spark_connect import DataprocSparkSession + +# After +from google.cloud.managed_spark_connect import ManagedSparkSession ``` ## Developing diff --git a/cloudbuild/cloudbuild.yaml b/cloudbuild/cloudbuild.yaml index 41890436..bb56bfb1 100644 --- a/cloudbuild/cloudbuild.yaml +++ b/cloudbuild/cloudbuild.yaml @@ -3,9 +3,9 @@ steps: # distribution artifacts. - name: 'gcr.io/cloud-builders/docker' id: 'build-container-image' - args: ['build', '--tag=gcr.io/${PROJECT_ID}/dataproc-spark-connect/dataproc-spark-connect-presubmit:${BUILD_ID}', -f, 'cloudbuild/Dockerfile', '.'] + args: ['build', '--tag=gcr.io/${PROJECT_ID}/managed-spark-connect/managed-spark-connect-presubmit:${BUILD_ID}', -f, 'cloudbuild/Dockerfile', '.'] # Run all unit tests - - name: 'gcr.io/${PROJECT_ID}/dataproc-spark-connect/dataproc-spark-connect-presubmit:${BUILD_ID}' + - name: 'gcr.io/${PROJECT_ID}/managed-spark-connect/managed-spark-connect-presubmit:${BUILD_ID}' id: 'run-unit-tests' waitFor: ['build-container-image'] entrypoint: 'pytest' diff --git a/google/cloud/dataproc_magics/__init__.py b/google/cloud/dataproc_magics/__init__.py index a348eb82..a7001f09 100644 --- a/google/cloud/dataproc_magics/__init__.py +++ b/google/cloud/dataproc_magics/__init__.py @@ -11,9 +11,17 @@ # 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. +"""Deprecated: this package has been renamed to ``google.cloud.managed_spark_magics``.""" +import warnings from .magics import DataprocMagics +warnings.warn( + "google.cloud.dataproc_magics is deprecated, use google.cloud.managed_spark_magics instead.", + DeprecationWarning, + stacklevel=2, +) + def load_ipython_extension(ipython): ipython.register_magics(DataprocMagics) diff --git a/google/cloud/dataproc_magics/magics.py b/google/cloud/dataproc_magics/magics.py index 278cc817..014d519f 100644 --- a/google/cloud/dataproc_magics/magics.py +++ b/google/cloud/dataproc_magics/magics.py @@ -11,66 +11,7 @@ # 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. +"""Deprecated: use ``google.cloud.managed_spark_magics.magics`` instead.""" +from google.cloud.managed_spark_magics.magics import ManagedSparkMagics -"""Dataproc magic implementations.""" - -import shlex -from IPython.core.magic import (Magics, magics_class, line_magic) -from google.cloud.dataproc_spark_connect import DataprocSparkSession - - -@magics_class -class DataprocMagics(Magics): - - def __init__( - self, - shell, - **kwargs, - ): - super().__init__(shell, **kwargs) - - @line_magic - def dpip(self, line): - """ - Custom magic to install pip packages as Spark Connect artifacts. - Usage: %dpip install pandas numpy - """ - try: - args = shlex.split(line) - - if not args or args[0] != "install": - raise RuntimeError( - "Usage: %dpip install ..." - ) - - packages = args[1:] # remove `install` - - if not packages: - raise RuntimeError("Error: No packages specified.") - - if any(pkg.startswith("-") for pkg in packages): - raise RuntimeError("Error: Flags are not currently supported.") - - sessions = [ - (key, value) - for key, value in self.shell.user_ns.items() - if isinstance(value, DataprocSparkSession) - ] - - if not sessions: - raise RuntimeError( - "Error: No active Dataproc Spark Session found. Please create one first." - ) - if len(sessions) > 1: - raise RuntimeError( - "Error: Found more than one active Dataproc Spark Sessions." - ) - - ((name, session),) = sessions - print(f"Active session found: {name}") - print(f"Installing packages: {packages}") - session.addArtifacts(*packages, pypi=True) - - print("Finished installing packages.") - except Exception as e: - raise RuntimeError(f"Failed to install packages: {e}") from e +DataprocMagics = ManagedSparkMagics diff --git a/google/cloud/dataproc_spark_connect/__init__.py b/google/cloud/dataproc_spark_connect/__init__.py index 008be626..862d30dc 100644 --- a/google/cloud/dataproc_spark_connect/__init__.py +++ b/google/cloud/dataproc_spark_connect/__init__.py @@ -11,19 +11,16 @@ # 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 importlib.metadata +"""Deprecated: this package has been renamed to ``google.cloud.managed_spark_connect``.""" import warnings -from .session import DataprocSparkSession +from google.cloud.managed_spark_connect import ManagedSparkSession -old_package_name = "google-spark-connect" -current_package_name = "dataproc-spark-connect" -try: - importlib.metadata.distribution(old_package_name) - 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}'." - ) -except: - pass +DataprocSparkSession = ManagedSparkSession + +warnings.warn( + "google.cloud.dataproc_spark_connect is deprecated, use google.cloud.managed_spark_connect instead. " + "DataprocSparkSession has been renamed to ManagedSparkSession.", + DeprecationWarning, + stacklevel=2, +) diff --git a/google/cloud/dataproc_spark_connect/client/__init__.py b/google/cloud/dataproc_spark_connect/client/__init__.py index 1902a49b..da634080 100644 --- a/google/cloud/dataproc_spark_connect/client/__init__.py +++ b/google/cloud/dataproc_spark_connect/client/__init__.py @@ -11,4 +11,5 @@ # 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. +"""Deprecated: use ``google.cloud.managed_spark_connect.client`` instead.""" from .core import DataprocChannelBuilder diff --git a/google/cloud/dataproc_spark_connect/client/core.py b/google/cloud/dataproc_spark_connect/client/core.py index 02cd2e70..7fedc5fc 100644 --- a/google/cloud/dataproc_spark_connect/client/core.py +++ b/google/cloud/dataproc_spark_connect/client/core.py @@ -11,131 +11,10 @@ # 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 logging +"""Deprecated: use ``google.cloud.managed_spark_connect.client.core`` instead.""" +from google.cloud.managed_spark_connect.client.core import ( + ManagedSparkChannelBuilder, + ProxiedChannel, +) -import google -import grpc -from pyspark.sql.connect.client import DefaultChannelBuilder - -from . import proxy - -logger = logging.getLogger(__name__) - - -class DataprocChannelBuilder(DefaultChannelBuilder): - """ - This is a helper class that is used to create a GRPC channel based on the given - connection string per the documentation of Spark Connect. - - This implementation of ChannelBuilder uses `secure_authorized_channel` from the - `google.auth.transport.grpc` package for authenticating secure channel. - - Examples - -------- - >>> cb = ChannelBuilder("sc://localhost") - ... cb.endpoint - - >>> cb = ChannelBuilder("sc://localhost/;use_ssl=true;token=aaa") - ... cb.secure - True - """ - - def __init__(self, url, is_active_callback=None): - self._is_active_callback = is_active_callback - super().__init__(url) - - def toChannel(self) -> grpc.Channel: - """ - Applies the parameters of the connection string and creates a new - GRPC channel according to the configuration. Passes optional channel options to - construct the channel. - - Returns - ------- - GRPC Channel instance. - """ - # TODO: Replace with a direct channel once all compatibility issues with - # grpc have been resolved. - return self._proxied_channel() - - def _proxied_channel(self) -> grpc.Channel: - return ProxiedChannel(self.host, self._is_active_callback) - - def _direct_channel(self) -> grpc.Channel: - destination = f"{self.host}:{self.port}" - - credentials, project = google.auth.default( - scopes=["https://www.googleapis.com/auth/cloud-platform"] - ) - # Get an HTTP request function to refresh credentials. - request = google.auth.transport.requests.Request() - # Create a channel. - - return google.auth.transport.grpc.secure_authorized_channel( - credentials, - request, - destination, - None, - None, - options=self._channel_options, - ) - - -class ProxiedChannel(grpc.Channel): - - def __init__(self, target_host, is_active_callback): - self._is_active_callback = is_active_callback - self._proxy = proxy.DataprocSessionProxy(0, target_host) - self._proxy.start() - self._proxied_connect_url = f"sc://localhost:{self._proxy.port}" - self._wrapped = DefaultChannelBuilder( - self._proxied_connect_url - ).toChannel() - - def __enter__(self): - return self - - def __exit__(self, *args): - ret = self._wrapped.__exit__(*args) - self._proxy.stop() - return ret - - def close(self): - ret = self._wrapped.close() - self._proxy.stop() - return ret - - def _wrap_method(self, wrapped_method): - if self._is_active_callback is None: - return wrapped_method - - def checked_method(*margs, **mkwargs): - if ( - self._is_active_callback is not None - and not self._is_active_callback() - ): - logger.warning(f"Session is no longer active") - raise RuntimeError( - "Session not active. Please create a new session" - ) - return wrapped_method(*margs, **mkwargs) - - return checked_method - - def stream_stream(self, *args, **kwargs): - return self._wrap_method(self._wrapped.stream_stream(*args, **kwargs)) - - def stream_unary(self, *args, **kwargs): - return self._wrap_method(self._wrapped.stream_unary(*args, **kwargs)) - - def subscribe(self, *args, **kwargs): - return self._wrap_method(self._wrapped.subscribe(*args, **kwargs)) - - def unary_stream(self, *args, **kwargs): - return self._wrap_method(self._wrapped.unary_stream(*args, **kwargs)) - - def unary_unary(self, *args, **kwargs): - return self._wrap_method(self._wrapped.unary_unary(*args, **kwargs)) - - def unsubscribe(self, *args, **kwargs): - return self._wrap_method(self._wrapped.unsubscribe(*args, **kwargs)) +DataprocChannelBuilder = ManagedSparkChannelBuilder diff --git a/google/cloud/dataproc_spark_connect/client/proxy.py b/google/cloud/dataproc_spark_connect/client/proxy.py old mode 100755 new mode 100644 index 2a1b3bf0..5449e384 --- a/google/cloud/dataproc_spark_connect/client/proxy.py +++ b/google/cloud/dataproc_spark_connect/client/proxy.py @@ -1,269 +1,13 @@ -#!/bin/env python - -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import contextlib -import logging -import socket -import threading - -import websockets.sync.client as websocketclient - -from google import auth as googleauth -from google.auth.transport import requests as googleauthrequests - -parser = argparse.ArgumentParser() -parser.add_argument("port") -parser.add_argument("target_host") - - -logger = logging.getLogger(__name__) - - -class bridged_socket(object): - """Socket-like object that uses a websocket-over-TCP Bridge transport. - - See: https://github.com/google/inverting-proxy/tree/master/utils/tcpbridge - """ - - def __init__(self, websocket_conn): - self._conn = websocket_conn - - def recv(self, buff_size): - # N.B. The websockets [recv method](https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.recv) - # does not support the buff_size parameter, but it does add a `timeout` keyword parameter not supported by normal - # socket objects. - # - # We set that timeout to 60 seconds to prevent any scenarios where we wind up stuck waiting for a message from a websocket connection - # that never comes. - msg = self._conn.recv(timeout=60) - return bytes.fromhex(msg) - - def send(self, msg_bytes): - msg = bytes.hex(msg_bytes) - self._conn.send(msg) - - def close(self): - return self._conn.close() - - -def connect_tcp_bridge(hostname): - """Create a socket-like connection to the given hostname using websocket. - - The backend server connected to over the websocket connection must be - running the TCP-bridge backend corresponding to this frontend. - - Args: - hostname: The hostname of the server running the TCP-bridge backend. - - Returns: - A socket-like object with `recv` and `send` methods. - """ - path = "tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96" - creds, _ = googleauth.default( - scopes=["https://www.googleapis.com/auth/cloud-platform"] - ) - creds.refresh(googleauthrequests.Request()) - - return websocketclient.connect( - f"wss://{hostname}/{path}", - additional_headers={"Authorization": f"Bearer {creds.token}"}, - open_timeout=30, - ) - - -def forward_bytes(name, from_sock, to_sock): - """Continuously stream bytes from the `from_sock` to the `to_sock`. - - This method terminates when either the `from_sock` is closed (causing - it to return a Falsy value from its `recv` method), or the first time - it hits an exception. - - This method is intended to be run in a separate thread of execution. - - Args: - name: forwarding thread name - from_sock: A socket-like object to stream bytes from. - to_sock: A socket-like object to stream bytes to. - """ - while True: - try: - bs = from_sock.recv(1024) - if not bs: - to_sock.close() - return - attempt = 0 - while bs and (attempt < 10): - attempt += 1 - try: - to_sock.send(bs) - bs = None - except TimeoutError: - # On timeouts during a send, we retry just the send - # to make sure we don't lose any bytes. - pass - if bs: - raise Exception(f"Failed to forward bytes for {name}") - except TimeoutError: - # On timeouts during a receive, we retry the entire flow. - pass - except Exception as ex: - logger.debug(f"[{name}] Exception forwarding bytes: {ex}") - to_sock.close() - return - - -def connect_sockets(conn_number, from_sock, to_sock): - """Create a connection between the two given ports. - - This method continuously streams bytes in both directions between the - given `from_sock` and `to_sock` socket-like objects. - - The caller is responsible for creating and closing the supplied sockets. - """ - forward_name = f"{conn_number}-forward" - t1 = threading.Thread( - name=forward_name, - target=forward_bytes, - args=[forward_name, from_sock, to_sock], - daemon=True, - ) - t1.start() - backward_name = f"{conn_number}-backward" - t2 = threading.Thread( - name=backward_name, - target=forward_bytes, - args=[backward_name, to_sock, from_sock], - daemon=True, - ) - t2.start() - t1.join() - t2.join() - - -def forward_connection(conn_number, conn, addr, target_host): - """Create a connection to the target and forward `conn` to it. - - This method creates a socket-like object holding a connection to the given - target host, and then continuously streams bytes in both directions between - `conn` and that newly created connection. - - Both the supplied incoming connection (`conn`) and the created outgoing - connection are automatically closed when this method terminates. - - This method should be run inside a daemon thread so that it will not - block program termination. - """ - with conn: - with connect_tcp_bridge(target_host) as websocket_conn: - backend_socket = bridged_socket(websocket_conn) - # 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(10) - connect_sockets(conn_number, conn, backend_socket) - - -class DataprocSessionProxy(object): - """A TCP proxy for forwarding requests to Dataproc Serverless Sessions. - - Spark Connect clients connect to this proxy using the h2c (without-SSL) - protocol, and this proxy adds SSL by tunneling those connections over - an HTTPS/WebSocket connection to the backend server. - - The tunneled requests are authenticated using the Google Application - Default Credentials. - """ - - def __init__(self, port, target_host): - self._port = port - self._target_host = target_host - self._started = False - self._killed = False - self._conn_number = 0 - - @property - def port(self): - """The local port the proxy is listening on""" - return self._port - - def start(self, daemon=True): - """Start the proxy. - - By the time this method returns the proxy has already started listening - on its local port will accept incoming connections. - """ - if self._started: - raise Exception("Dataproc session proxy already started") - self._started = True - s = threading.Semaphore(value=0) - t = threading.Thread(target=self._run, args=[s], daemon=daemon) - t.start() - s.acquire() - - def _run(self, s): - with socket.create_server(("127.0.0.1", self._port)) as frontend_socket: - if self._port == 0: - self._port = frontend_socket.getsockname()[1] - s.release() - while not self._killed: - conn, addr = frontend_socket.accept() - 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], - daemon=True, - ).start() - - def stop(self): - """Stop the proxy.""" - self._killed = True - - -@contextlib.contextmanager -def dataproc_session_proxy(port, target_host): - """Context manager for creating a Dataproc Session proxy. - - Usage: - with dataproc_session_proxy(0, backend_hostname) as p: - local_port = p.port - ... - - Args: - port: The local port to listen on. Use `0` to pick a free port. - target_host: The backend to proxy connections to. - - Returns: - A context manager wrapping a DataprocSessionProxy instance. - """ - proxy = DataprocSessionProxy(port, target_host) - try: - proxy.start(daemon=False) - yield proxy - finally: - proxy.stop() - - -if __name__ == "__main__": - args = parser.parse_args() - with dataproc_session_proxy(int(args.port), args.target_host) as p: - print(f"Proxy listening on port {p.port}") - try: - while True: - pass - except KeyboardInterrupt: - pass +"""Deprecated: use ``google.cloud.managed_spark_connect.client.proxy`` instead.""" + +from google.cloud.managed_spark_connect.client.proxy import ( + ManagedSparkSessionProxy, + connect_sockets, + connect_tcp_bridge, + forward_bytes, + forward_connection, + managed_spark_session_proxy, +) + +DataprocSessionProxy = ManagedSparkSessionProxy +dataproc_session_proxy = managed_spark_session_proxy diff --git a/google/cloud/dataproc_spark_connect/environment.py b/google/cloud/dataproc_spark_connect/environment.py index e19dd973..e04708cc 100644 --- a/google/cloud/dataproc_spark_connect/environment.py +++ b/google/cloud/dataproc_spark_connect/environment.py @@ -11,180 +11,5 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -import os -import sys -from typing import Callable, Tuple, List - - -def is_antigravity() -> bool: - """True if running inside the Antigravity environment.""" - return "antigravity" in os.getenv("__CFBundleIdentifier", "").lower() - - -def is_vscode() -> bool: - """True if running inside VS Code at all.""" - return os.getenv("VSCODE_PID") is not None - - -def is_jupyter() -> bool: - """True if running in a Jupyter environment.""" - return os.getenv("JPY_PARENT_PID") is not None - - -def is_colab_enterprise() -> bool: - """True if running in Colab Enterprise (Vertex AI).""" - return os.getenv("VERTEX_PRODUCT") == "COLAB_ENTERPRISE" - - -def is_colab() -> bool: - """True if running in Google Colab.""" - return os.getenv("COLAB_RELEASE_TAG") is not None - - -def is_workbench() -> bool: - """True if running in Vertex Workbench Instance (managed Jupyter).""" - return os.getenv("VERTEX_PRODUCT") == "WORKBENCH_INSTANCE" - - -def is_kaggle() -> bool: - """True if running in Kaggle Notebooks.""" - return os.getenv("KAGGLE_KERNEL_RUN_TYPE") is not None - - -def is_databricks() -> bool: - """True if running in Databricks.""" - return os.getenv("DATABRICKS_RUNTIME_VERSION") is not None - - -def is_sagemaker() -> bool: - """True if running in AWS SageMaker.""" - return os.getenv("SAGEMAKER_INTERNAL_IMAGE_URI") is not None - - -def is_deepnote() -> bool: - """True if running in Deepnote.""" - return os.getenv("DEEPNOTE_PROJECT_ID") is not None - - -def is_datalore() -> bool: - """True if running in JetBrains Datalore.""" - return os.getenv("DATALORE_USER") is not None - - -def is_spyder() -> bool: - """True if running inside Spyder IDE.""" - return any(k.startswith("SPYDER") for k in os.environ) - - -def is_cloud_shell() -> bool: - """True if running in Google Cloud Shell.""" - return os.getenv("CLOUD_SHELL") is not None - - -def is_codespaces() -> bool: - """True if running in GitHub Codespaces.""" - return os.getenv("CODESPACES") is not None - - -def is_jetbrains_ide() -> bool: - """True if running inside JetBrains IDE.""" - return ( - "jetbrains" in os.getenv("TERMINAL_EMULATOR", "").lower() - or "PYCHARM_HOSTED" in os.environ - ) - - -def is_hex() -> bool: - """True if running in Hex.""" - return os.getenv("HEX_PROJECT_ID") is not None - - -def is_polynote() -> bool: - """True if running in Polynote.""" - return os.getenv("POLYNOTE_VERSION") is not None - - -def is_eclipse() -> bool: - """True if running inside Eclipse IDE.""" - return "ECLIPSE_HOME" in os.environ or any( - k.startswith("ECLIPSE") for k in os.environ - ) - - -def is_interactive() -> bool: - try: - from IPython import get_ipython - - if get_ipython() is not None: - return True - except ImportError: - pass - - return hasattr(sys, "ps1") or bool(sys.flags.interactive) - - -def is_terminal() -> bool: - return sys.stdin.isatty() - - -def is_interactive_terminal() -> bool: - return is_interactive() and is_terminal() - - -def is_dataproc_batch() -> bool: - return os.getenv("DATAPROC_WORKLOAD_TYPE") == "batch" - - -def get_client_environment_label() -> str: - """ - Map current environment to a standardized client label. - - Priority order: - 1. Colab Enterprise ("colab-enterprise") - 2. Colab ("colab") - 3. Vertex Workbench Instance ("workbench-jupyter") - 4. Kaggle ("kaggle") - 5. AWS SageMaker ("sagemaker") - 6. Databricks ("databricks") - 7. Deepnote ("deepnote") - 8. JetBrains Datalore ("datalore") - 9. GitHub Codespaces ("codespaces") - 10. Google Cloud Shell ("cloud-shell") - 11. Hex ("hex") - 12. Polynote ("polynote") - 13. Antigravity ("antigravity") - 14. VS Code ("vscode") - 15. JetBrains IDE ("jetbrains") - 16. Spyder ("spyder") - 17. Eclipse ("eclipse") - 18. Jupyter ("jupyter") - 19. Unknown ("unknown") - """ - checks: List[Tuple[Callable[[], bool], str]] = [ - (is_colab_enterprise, "colab-enterprise"), - (is_colab, "colab"), - (is_workbench, "workbench-jupyter"), - (is_kaggle, "kaggle"), - (is_sagemaker, "sagemaker"), - (is_databricks, "databricks"), - (is_deepnote, "deepnote"), - (is_datalore, "datalore"), - (is_codespaces, "codespaces"), - (is_cloud_shell, "cloud-shell"), - (is_hex, "hex"), - (is_polynote, "polynote"), - (is_antigravity, "antigravity"), - (is_vscode, "vscode"), - (is_jetbrains_ide, "jetbrains"), - (is_spyder, "spyder"), - (is_eclipse, "eclipse"), - (is_jupyter, "jupyter"), - ] - for detector, label in checks: - try: - if detector(): - return label - except Exception: - pass - return "unknown" +"""Deprecated: use ``google.cloud.managed_spark_connect.environment`` instead.""" +from google.cloud.managed_spark_connect.environment import * # noqa: F401,F403 diff --git a/google/cloud/dataproc_spark_connect/exceptions.py b/google/cloud/dataproc_spark_connect/exceptions.py index 3e5c8e90..d958a7cb 100644 --- a/google/cloud/dataproc_spark_connect/exceptions.py +++ b/google/cloud/dataproc_spark_connect/exceptions.py @@ -11,17 +11,9 @@ # 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. +"""Deprecated: use ``google.cloud.managed_spark_connect.exceptions`` instead.""" +from google.cloud.managed_spark_connect.exceptions import ( + ManagedSparkConnectException, +) - -class DataprocSparkConnectException(Exception): - """A custom exception class to only print the error messages. - This would be used for exceptions where the stack trace - doesn't provide any additional information. - """ - - def __init__(self, message): - self.message = message - super().__init__(message) - - def _render_traceback_(self): - return [self.message] +DataprocSparkConnectException = ManagedSparkConnectException diff --git a/google/cloud/dataproc_spark_connect/pypi_artifacts.py b/google/cloud/dataproc_spark_connect/pypi_artifacts.py index b31d3864..f6f27812 100644 --- a/google/cloud/dataproc_spark_connect/pypi_artifacts.py +++ b/google/cloud/dataproc_spark_connect/pypi_artifacts.py @@ -1,48 +1,3 @@ -import json -import logging -import os -import tempfile +"""Deprecated: use ``google.cloud.managed_spark_connect.pypi_artifacts`` instead.""" -from packaging.requirements import Requirement - -logger = logging.getLogger(__name__) - - -class PyPiArtifacts: - """ - This is a helper class to serialize the PYPI package installation request with a "magic" file name - that Spark Connect server understands - """ - - @staticmethod - 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") - return reqs - - def __init__(self, packages: set[str]): - self.requirements = PyPiArtifacts.__try_parsing_package(packages) - - def write_packages_config(self, s8s_session_uuid: str) -> str: - """ - Can't use the same file-name as Spark throws exception that file already exists - Keep the filename/format in sync with server - """ - dependencies = { - "version": "0.5", - "packageType": "PYPI", - "packages": [str(req) for req in self.requirements], - } - - file_path = os.path.join( - tempfile.gettempdir(), - s8s_session_uuid, - "add-artifacts-1729-" + self.__str__() + ".json", - ) - - os.makedirs(os.path.dirname(file_path), exist_ok=True) - with open(file_path, "w") as json_file: - json.dump(dependencies, json_file, indent=4) - logger.debug("Dumping dependencies request in file: " + file_path) - return file_path +from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts diff --git a/google/cloud/dataproc_spark_connect/session.py b/google/cloud/dataproc_spark_connect/session.py index 7ad24fe8..4e838d96 100644 --- a/google/cloud/dataproc_spark_connect/session.py +++ b/google/cloud/dataproc_spark_connect/session.py @@ -11,1372 +11,11 @@ # 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 atexit -import datetime -import functools -import json -import logging -import os -import random -import re -import string -import threading -import time -import uuid -import tqdm -from packaging import version -from types import MethodType -from typing import Any, cast, ClassVar, Dict, Iterable, Optional, Union - -from google.api_core import retry -from google.api_core.client_options import ClientOptions -from google.api_core.exceptions import ( - Aborted, - FailedPrecondition, - InvalidArgument, - NotFound, - PermissionDenied, -) -from google.api_core.future.polling import POLLING_PREDICATE -from google.auth.exceptions import DefaultCredentialsError -from google.cloud.dataproc_spark_connect.client import DataprocChannelBuilder -from google.cloud.dataproc_spark_connect.exceptions import DataprocSparkConnectException -from google.cloud.dataproc_spark_connect.pypi_artifacts import PyPiArtifacts -from google.cloud.dataproc_v1 import ( - AuthenticationConfig, - CreateSessionRequest, - DeleteSessionRequest, - GetSessionRequest, - Session, - SessionControllerClient, - TerminateSessionRequest, -) -from google.cloud.dataproc_v1.types import sessions -from google.cloud.dataproc_spark_connect import environment -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 -SYSTEM_LABELS = { - "dataproc-session-client", - "goog-colab-notebook-id", -} - -_DATAPROC_SESSIONS_BASE_URL = ( - "https://console.cloud.google.com/dataproc/interactive" +"""Deprecated: use ``google.cloud.managed_spark_connect.session`` instead.""" +from google.cloud.managed_spark_connect.session import ( + ManagedSparkSession, + _is_valid_label_value, + _is_valid_session_id, ) - -def _is_valid_label_value(value: str) -> bool: - """ - Validates if a string complies with Google Cloud label value format. - Only lowercase letters, numbers, and dashes are allowed. - The value must start with lowercase letter or number and end with a lowercase letter or number. - Maximum length is 63 characters. - """ - if not value: - return False - - # Check maximum length (63 characters for GCP label values) - if len(value) > 63: - return False - - # Check if the value matches the pattern: starts and ends with alphanumeric, - # contains only lowercase letters, numbers, and dashes - pattern = r"^[a-z0-9]([a-z0-9-]*[a-z0-9])?$" - return bool(re.match(pattern, value)) - - -def _is_valid_session_id(session_id: str) -> bool: - """ - Validates if a string complies with Google Cloud session ID format. - - Must be 4-63 characters - - Only lowercase letters, numbers, and dashes are allowed - - Must start with a lowercase letter - - Cannot end with a dash - """ - if not session_id: - return False - - # The pattern is sufficient for validation and already enforces length constraints. - pattern = r"^[a-z][a-z0-9-]{2,61}[a-z0-9]$" - return bool(re.match(pattern, session_id)) - - -class DataprocSparkSession(SparkSession): - """The entry point to programming Spark with the Dataset and DataFrame API. - - A DataprocRemoteSparkSession can be used to create :class:`DataFrame`, register :class:`DataFrame` as - tables, execute SQL over tables, cache tables, and read parquet files. - - Examples - -------- - - Create a Spark session with Dataproc Spark Connect. - - >>> spark = ( - ... DataprocSparkSession.builder - ... .appName("Word Count") - ... .dataprocSessionConfig(Session()) - ... .getOrCreate() - ... ) # doctest: +SKIP - """ - - _DEFAULT_RUNTIME_VERSION = "3.0" - _MIN_RUNTIME_VERSION = "3.0" - - _active_s8s_session_uuid: ClassVar[Optional[str]] = None - _project_id = None - _region = None - _client_options = None - _active_s8s_session_id: ClassVar[Optional[str]] = None - _active_session_uses_custom_id: ClassVar[bool] = False - _execution_progress_bar = dict() - - class Builder(SparkSession.Builder): - - def __init__(self): - self._options: Dict[str, Any] = {} - self._channel_builder: Optional[DataprocChannelBuilder] = None - self._dataproc_config: Optional[Session] = None - self._custom_session_id: Optional[str] = None - self._project_id = os.getenv("GOOGLE_CLOUD_PROJECT") - self._region = os.getenv("GOOGLE_CLOUD_REGION") - self._client_options = ClientOptions( - api_endpoint=os.getenv( - "GOOGLE_CLOUD_DATAPROC_API_ENDPOINT", - f"{self._region}-dataproc.googleapis.com", - ) - ) - self._session_controller_client: Optional[ - SessionControllerClient - ] = None - - @property - def session_controller_client(self) -> SessionControllerClient: - """Get or create a SessionControllerClient instance.""" - if self._session_controller_client is None: - self._session_controller_client = SessionControllerClient( - client_options=self._client_options - ) - return self._session_controller_client - - def projectId(self, project_id): - self._project_id = project_id - return self - - def location(self, location): - self._region = location - self._client_options.api_endpoint = os.getenv( - "GOOGLE_CLOUD_DATAPROC_API_ENDPOINT", - f"{self._region}-dataproc.googleapis.com", - ) - return self - - def dataprocSessionId(self, session_id: str): - """ - Set a custom session ID for creating or reusing sessions. - - The session ID must: - - Be 4-63 characters long - - Start with a lowercase letter - - Contain only lowercase letters, numbers, and hyphens - - Not end with a hyphen - - Args: - session_id: The custom session ID to use - - Returns: - This Builder instance for method chaining - - Raises: - ValueError: If the session ID format is invalid - """ - if not _is_valid_session_id(session_id): - raise ValueError( - f"Invalid session ID: '{session_id}'. " - "Session ID must be 4-63 characters, start with a lowercase letter, " - "contain only lowercase letters, numbers, and hyphens, " - "and not end with a hyphen." - ) - self._custom_session_id = session_id - return self - - def dataprocSessionConfig(self, dataproc_config: Session): - self._dataproc_config = dataproc_config - for k, v in dataproc_config.runtime_config.properties.items(): - self._options[cast(str, k)] = to_str(v) - return self - - @property - def dataproc_config(self): - with self._lock: - self._dataproc_config = self._dataproc_config or Session() - return self._dataproc_config - - def runtimeVersion(self, version: str): - self.dataproc_config.runtime_config.version = version - return self - - def serviceAccount(self, account: str): - self.dataproc_config.environment_config.execution_config.service_account = ( - account - ) - return self - - def subnetwork(self, subnet: str): - self.dataproc_config.environment_config.execution_config.subnetwork_uri = ( - subnet - ) - return self - - def ttl(self, duration: datetime.timedelta): - """Set the time-to-live (TTL) for the session using a timedelta object.""" - return self.ttlSeconds(int(duration.total_seconds())) - - def ttlSeconds(self, seconds: int): - """Set the time-to-live (TTL) for the session in seconds.""" - self.dataproc_config.environment_config.execution_config.ttl = { - "seconds": seconds - } - return self - - def idleTtl(self, duration: datetime.timedelta): - """Set the idle time-to-live (idle TTL) for the session using a timedelta object.""" - return self.idleTtlSeconds(int(duration.total_seconds())) - - def idleTtlSeconds(self, seconds: int): - """Set the idle time-to-live (idle TTL) for the session in seconds.""" - self.dataproc_config.environment_config.execution_config.idle_ttl = { - "seconds": seconds - } - return self - - def sessionTemplate(self, template: str): - self.dataproc_config.session_template = template - return self - - def label(self, key: str, value: str): - """Add a single label to the session.""" - return self.labels({key: value}) - - def labels(self, labels: Dict[str, str]): - # Filter out system labels and warn user - filtered_labels = {} - for key, value in labels.items(): - if key in SYSTEM_LABELS: - logger.warning( - f"Label '{key}' is a system label and cannot be overridden by user. Ignoring." - ) - else: - filtered_labels[key] = value - - self.dataproc_config.labels.update(filtered_labels) - return self - - def remote(self, url: Optional[str] = None) -> "SparkSession.Builder": - if url: - raise NotImplemented( - "DataprocSparkSession does not support connecting to an existing remote server" - ) - else: - return self - - def create(self) -> "DataprocSparkSession": - raise NotImplemented( - "DataprocSparkSession allows session creation only through getOrCreate" - ) - - def __create_spark_connect_session_from_s8s( - self, session_response, session_name - ) -> "DataprocSparkSession": - DataprocSparkSession._active_s8s_session_uuid = ( - session_response.uuid - ) - DataprocSparkSession._project_id = self._project_id - DataprocSparkSession._region = self._region - DataprocSparkSession._client_options = self._client_options - spark_connect_url = session_response.runtime_info.endpoints.get( - "Spark Connect Server" - ) - url = f"{spark_connect_url}/;session_id={session_response.uuid};use_ssl=true" - logger.debug(f"Spark Connect URL: {url}") - self._channel_builder = DataprocChannelBuilder( - url, - is_active_callback=lambda: is_s8s_session_active( - session_name, self._client_options - ), - ) - - assert self._channel_builder is not None - session = DataprocSparkSession(connection=self._channel_builder) - - # Register handler for Cell Execution Progress bar - session._register_progress_execution_handler() - - DataprocSparkSession._set_default_and_active_session(session) - - return session - - def __create(self) -> "DataprocSparkSession": - with self._lock: - - if self._options.get("spark.remote", False): - raise NotImplemented( - "DataprocSparkSession does not support connecting to an existing Spark Connect remote server" - ) - - from google.cloud.dataproc_v1 import SessionControllerClient - - dataproc_config: Session = self._get_dataproc_config() - - # Check runtime version compatibility before creating session - self._check_runtime_compatibility(dataproc_config) - - # Use custom session ID if provided, otherwise generate one - session_id = ( - self._custom_session_id - if self._custom_session_id - else self.generate_dataproc_session_id() - ) - - dataproc_config.name = f"projects/{self._project_id}/locations/{self._region}/sessions/{session_id}" - logger.debug( - f"Dataproc Session configuration:\n{dataproc_config}" - ) - - session_request = CreateSessionRequest() - session_request.session_id = session_id - session_request.session = dataproc_config - session_request.parent = ( - f"projects/{self._project_id}/locations/{self._region}" - ) - - logger.debug("Creating Dataproc Session") - DataprocSparkSession._active_s8s_session_id = session_id - # Track whether this session uses a custom ID (unmanaged) or auto-generated ID (managed) - DataprocSparkSession._active_session_uses_custom_id = ( - self._custom_session_id is not None - ) - s8s_creation_start_time = time.time() - - stop_create_session_pbar_event = threading.Event() - - 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 - if i >= iterations - 1: - # Sleep until session created - while not stop_create_session_pbar_event.is_set(): - time.sleep(1) - else: - time.sleep(1) - - 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" - - try: - if ( - os.getenv( - "DATAPROC_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT", - "false", - ) - == "true" - ): - atexit.register( - lambda: terminate_s8s_session( - self._project_id, - self._region, - session_id, - self._client_options, - ) - ) - operation = SessionControllerClient( - client_options=self._client_options - ).create_session(session_request) - self._display_session_link_on_creation(session_id) - self._display_view_session_details_button(session_id) - create_session_pbar_thread.start() - session_response: Session = operation.result( - polling=retry.Retry( - predicate=POLLING_PREDICATE, - initial=5.0, # seconds - maximum=5.0, # seconds - multiplier=1.0, - timeout=600, # seconds - ) - ) - stop_create_session_pbar_event.set() - create_session_pbar_thread.join() - self._print_session_created_message() - file_path = ( - DataprocSparkSession._get_active_session_file_path() - ) - if file_path is not None: - try: - session_data = { - "session_name": session_response.name, - "session_uuid": session_response.uuid, - } - os.makedirs( - os.path.dirname(file_path), exist_ok=True - ) - with open(file_path, "w") as json_file: - json.dump(session_data, json_file, indent=4) - except Exception as e: - logger.error( - 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() - DataprocSparkSession._active_s8s_session_id = None - DataprocSparkSession._active_session_uses_custom_id = False - raise DataprocSparkConnectException( - f"Error while creating Dataproc 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() - DataprocSparkSession._active_s8s_session_id = None - DataprocSparkSession._active_session_uses_custom_id = False - raise DataprocSparkConnectException( - "Credentials error while creating Dataproc 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() - DataprocSparkSession._active_s8s_session_id = None - DataprocSparkSession._active_session_uses_custom_id = False - raise RuntimeError( - f"Error while creating Dataproc Session" - ) from e - finally: - stop_create_session_pbar_event.set() - - logger.debug( - f"Dataproc Session created: {session_id} in {int(time.time() - s8s_creation_start_time)} seconds" - ) - return self.__create_spark_connect_session_from_s8s( - session_response, dataproc_config.name - ) - - def _wait_for_session_available( - self, session_name: str, timeout: int = 300 - ) -> Session: - start_time = time.time() - while time.time() - start_time < timeout: - try: - session = self.session_controller_client.get_session( - name=session_name - ) - if "Spark Connect Server" in session.runtime_info.endpoints: - return session - time.sleep(5) - except Exception as e: - logger.warning( - f"Error while polling for Spark Connect endpoint: {e}" - ) - time.sleep(5) - raise RuntimeError( - f"Spark Connect endpoint not available for session {session_name} after {timeout} seconds." - ) - - def _display_session_link_on_creation(self, session_id): - session_url = f"{_DATAPROC_SESSIONS_BASE_URL}/{self._region}/{session_id}?project={self._project_id}" - plain_message = f"Creating Dataproc Session: {session_url}" - if environment.is_colab_enterprise(): - html_element = f""" -
-

Creating Dataproc Spark Session

-

- """ - else: - html_element = f""" -
-

Creating Dataproc Spark Session

-

Dataproc Session

-
- """ - self._output_element_or_message(plain_message, html_element) - - def _print_session_created_message(self): - plain_message = f"Dataproc Session was successfully created" - html_element = f"

{plain_message}

" - - self._output_element_or_message(plain_message, html_element) - - def _output_element_or_message(self, plain_message, html_element): - """ - Display / print the needed rich HTML element or plain text depending - on whether rich element is supported or not. - - :param plain_message: Message to print on non-IPython or - non-interactive shell - :param html_element: HTML element to display for interactive IPython - environment - """ - # Don't print any output (Rich or Plain) for non-interactive - if not environment.is_interactive(): - return - - if environment.is_interactive_terminal(): - print(plain_message) - return - - try: - from IPython.display import display, HTML - - display(HTML(html_element)) - except ImportError: - print(plain_message) - - def _get_exiting_active_session( - self, - ) -> Optional["DataprocSparkSession"]: - s8s_session_id = DataprocSparkSession._active_s8s_session_id - session_name = f"projects/{self._project_id}/locations/{self._region}/sessions/{s8s_session_id}" - session_response = None - session = None - if s8s_session_id is not None: - session_response = get_active_s8s_session_response( - session_name, self._client_options - ) - session = DataprocSparkSession.getActiveSession() - - if session is None: - session = DataprocSparkSession._default_session - - if session_response is not None: - print( - f"Using existing Dataproc Session (configuration changes may not be applied): {_DATAPROC_SESSIONS_BASE_URL}/{self._region}/{s8s_session_id}?project={self._project_id}" - ) - self._display_view_session_details_button(s8s_session_id) - if session is None: - session_response = self._wait_for_session_available( - session_name - ) - session = self.__create_spark_connect_session_from_s8s( - session_response, session_name - ) - return session - else: - if session is not None: - print( - f"{s8s_session_id} Dataproc Session is not active, stopping and creating a new one" - ) - session.stop() - - return None - - def getOrCreate(self) -> "DataprocSparkSession": - with DataprocSparkSession._lock: - if environment.is_dataproc_batch(): - # For Dataproc batch workloads, connect to the already initialized local SparkSession - from pyspark.sql import SparkSession as PySparkSQLSession - - session = PySparkSQLSession.builder.getOrCreate() - return session # type: ignore - - if self._project_id is None: - raise DataprocSparkConnectException( - f"Error while creating Dataproc Session: project ID is not set" - ) - - if self._region is None: - raise DataprocSparkConnectException( - f"Error while creating Dataproc Session: location is not set" - ) - - # Handle custom session ID by setting it early and letting existing logic handle it - if self._custom_session_id: - self._handle_custom_session_id() - - session = self._get_exiting_active_session() - if session is None: - session = self.__create() - - # Register this session as the instantiated SparkSession for compatibility - # with tools and libraries that expect SparkSession._instantiatedSession - from pyspark.sql import SparkSession as PySparkSQLSession - - PySparkSQLSession._instantiatedSession = session - - return session - - def _handle_custom_session_id(self): - """Handle custom session ID by checking if it exists and setting _active_s8s_session_id.""" - session_response = self._get_session_by_id(self._custom_session_id) - if session_response is not None: - # Found an active session with the custom ID, set it as the active session - DataprocSparkSession._active_s8s_session_id = ( - self._custom_session_id - ) - # Mark that this session uses a custom ID - DataprocSparkSession._active_session_uses_custom_id = True - else: - # No existing session found, clear any existing active session ID - # so we'll create a new one with the custom ID - DataprocSparkSession._active_s8s_session_id = None - - def _get_dataproc_config(self): - # Use the property to ensure we always have a config - dataproc_config = self.dataproc_config - for k, v in self._options.items(): - dataproc_config.runtime_config.properties[k] = v - dataproc_config.spark_connect_session = ( - sessions.SparkConnectConfig() - ) - if not dataproc_config.runtime_config.version: - dataproc_config.runtime_config.version = ( - DataprocSparkSession._DEFAULT_RUNTIME_VERSION - ) - - # Check for Python version mismatch with runtime for UDF compatibility - self._check_python_version_compatibility( - dataproc_config.runtime_config.version - ) - - # Use local variable to improve readability of deeply nested attribute access - exec_config = dataproc_config.environment_config.execution_config - - # Set service account from environment if not already set - if ( - not exec_config.service_account - and "DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT" in os.environ - ): - exec_config.service_account = os.getenv( - "DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT" - ) - - # Auto-set authentication type to SERVICE_ACCOUNT when service account is provided - if exec_config.service_account: - # When service account is provided, explicitly set auth type to SERVICE_ACCOUNT - exec_config.authentication_config.user_workload_authentication_type = ( - AuthenticationConfig.AuthenticationType.SERVICE_ACCOUNT - ) - elif ( - not exec_config.authentication_config.user_workload_authentication_type - and "DATAPROC_SPARK_CONNECT_AUTH_TYPE" in os.environ - ): - # Only set auth type from environment if no service account is present - exec_config.authentication_config.user_workload_authentication_type = AuthenticationConfig.AuthenticationType[ - os.getenv("DATAPROC_SPARK_CONNECT_AUTH_TYPE") - ] - if ( - not dataproc_config.environment_config.execution_config.subnetwork_uri - and "DATAPROC_SPARK_CONNECT_SUBNET" in os.environ - ): - dataproc_config.environment_config.execution_config.subnetwork_uri = os.getenv( - "DATAPROC_SPARK_CONNECT_SUBNET" - ) - if ( - not dataproc_config.environment_config.execution_config.ttl - and "DATAPROC_SPARK_CONNECT_TTL_SECONDS" in os.environ - ): - dataproc_config.environment_config.execution_config.ttl = { - "seconds": int( - os.getenv("DATAPROC_SPARK_CONNECT_TTL_SECONDS") - ) - } - if ( - not dataproc_config.environment_config.execution_config.idle_ttl - and "DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS" in os.environ - ): - dataproc_config.environment_config.execution_config.idle_ttl = { - "seconds": int( - os.getenv("DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS") - ) - } - client_environment = environment.get_client_environment_label() - dataproc_config.labels["dataproc-session-client"] = ( - client_environment - ) - if "COLAB_NOTEBOOK_ID" in os.environ: - colab_notebook_name = os.environ["COLAB_NOTEBOOK_ID"] - # Extract the last part of the path, which is the ID - notebook_id = os.path.basename(colab_notebook_name) - if _is_valid_label_value(notebook_id): - dataproc_config.labels["goog-colab-notebook-id"] = ( - notebook_id - ) - else: - logger.warning( - f"Warning while processing notebook ID: Notebook ID '{notebook_id}' is not compliant with label value format. " - f"Only lowercase letters, numbers, and dashes are allowed. " - f"The value must start with lowercase letter or number and end with a lowercase letter or number. " - f"Maximum length is 63 characters. " - f"Ignoring notebook ID label." - ) - default_datasource = os.getenv( - "DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE" - ) - match default_datasource: - case "bigquery": - # Merge default configs with existing properties, - # user configs take precedence - for k, v in { - "spark.sql.catalog.spark_catalog": "com.google.cloud.spark.bigquery.BigQuerySparkSessionCatalog", - "spark.sql.sources.default": "bigquery", - }.items(): - if k not in dataproc_config.runtime_config.properties: - dataproc_config.runtime_config.properties[k] = v - case _: - if default_datasource: - logger.warning( - f"DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE is set to an invalid value:" - f" {default_datasource}. Supported value is 'bigquery'." - ) - - return dataproc_config - - def _check_python_version_compatibility(self, runtime_version): - """Check if client Python version matches server Python version for UDF compatibility.""" - import sys - import warnings - - # Runtime version to server Python version mapping - RUNTIME_PYTHON_MAP = { - "3.0": (3, 12), - } - - client_python = sys.version_info[:2] # (major, minor) - - if runtime_version in RUNTIME_PYTHON_MAP: - server_python = RUNTIME_PYTHON_MAP[runtime_version] - - if client_python != server_python: - warnings.warn( - f"Python version mismatch detected: Client is using Python {client_python[0]}.{client_python[1]}, " - f"but Dataproc runtime {runtime_version} uses Python {server_python[0]}.{server_python[1]}. " - f"This mismatch may cause issues with Python UDF (User Defined Function) compatibility. " - f"Consider using Python {server_python[0]}.{server_python[1]} for optimal UDF execution.", - stacklevel=3, - ) - - def _check_runtime_compatibility(self, dataproc_config): - """Check if runtime version 3.0 client is compatible with older runtime versions. - - Runtime version 3.0 clients do not support older runtime versions (pre-3.0). - There is no backward or forward compatibility between different runtime versions. - - Args: - dataproc_config: The Session configuration containing runtime version - - Raises: - DataprocSparkConnectException: If server is using pre-3.0 runtime version - """ - runtime_version = dataproc_config.runtime_config.version - - if not runtime_version: - return - - logger.debug(f"Detected server runtime version: {runtime_version}") - - # Parse runtime version to check if it's below minimum supported version - try: - server_version = version.parse(runtime_version) - min_version = version.parse( - DataprocSparkSession._MIN_RUNTIME_VERSION - ) - - if server_version < min_version: - raise DataprocSparkConnectException( - f"Specified {runtime_version} Dataproc Runtime version is not supported, " - f"use {DataprocSparkSession._MIN_RUNTIME_VERSION} version or higher." - ) - except version.InvalidVersion: - # If we can't parse the version, log a warning but continue - logger.warning( - f"Could not parse runtime version: {runtime_version}" - ) - - def _display_view_session_details_button(self, session_id): - # Display button is only supported in colab enterprise - if not environment.is_colab_enterprise(): - return - - # Skip button display for colab enterprise IPython terminals - if environment.is_interactive_terminal(): - return - - try: - session_url = f"{_DATAPROC_SESSIONS_BASE_URL}/{self._region}/{session_id}?project={self._project_id}" - from IPython.core.interactiveshell import InteractiveShell - - if not InteractiveShell.initialized(): - return - - from google.cloud.aiplatform.utils import _ipython_utils - - _ipython_utils.display_link( - "View Session Details", f"{session_url}", "dashboard" - ) - except ImportError as e: - logger.debug(f"Import error: {e}") - - def _get_session_by_id(self, session_id: str) -> Optional[Session]: - """ - Get existing session by ID. - - Returns: - Session if ACTIVE/CREATING, None if not found or not usable - """ - session_name = f"projects/{self._project_id}/locations/{self._region}/sessions/{session_id}" - - try: - get_request = GetSessionRequest(name=session_name) - session = self.session_controller_client.get_session( - get_request - ) - - logger.debug( - f"Found existing session {session_id} in state: {session.state}" - ) - - if session.state in [ - Session.State.ACTIVE, - Session.State.CREATING, - ]: - # Reuse the active session - logger.info(f"Reusing existing session: {session_id}") - return session - else: - # Session exists but is not usable (terminated/failed/terminating) - logger.info( - f"Session {session_id} in {session.state.name} state, cannot reuse" - ) - return None - - except NotFound: - # Session doesn't exist, can create new one - logger.debug( - f"Session {session_id} not found, can create new one" - ) - return None - except Exception as e: - logger.error(f"Error checking session {session_id}: {e}") - return None - - def _delete_session(self, session_name: str): - """Delete a session to free up the session ID for reuse.""" - try: - delete_request = DeleteSessionRequest(name=session_name) - self.session_controller_client.delete_session(delete_request) - logger.debug(f"Deleted session: {session_name}") - except NotFound: - logger.debug(f"Session already deleted: {session_name}") - - def _wait_for_termination(self, session_name: str, timeout: int = 180): - """Wait for a session to finish terminating.""" - start_time = time.time() - - while time.time() - start_time < timeout: - try: - get_request = GetSessionRequest(name=session_name) - session = self.session_controller_client.get_session( - get_request - ) - - if session.state in [ - Session.State.TERMINATED, - Session.State.FAILED, - ]: - return - elif session.state != Session.State.TERMINATING: - # Session is in unexpected state - logger.warning( - f"Session {session_name} in unexpected state while waiting for termination: {session.state}" - ) - return - - time.sleep(2) - except NotFound: - # Session was deleted - return - - logger.warning( - f"Timeout waiting for session {session_name} to terminate" - ) - - @staticmethod - def generate_dataproc_session_id(): - timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") - suffix_length = 6 - random_suffix = "".join( - random.choices( - string.ascii_lowercase + string.digits, k=suffix_length - ) - ) - return f"sc-{timestamp}-{random_suffix}" - - def __init__( - self, - connection: Union[str, DataprocChannelBuilder], - user_id: Optional[str] = None, - ): - """ - Creates a new DataprocSparkSession for the Spark Connect interface. - - Parameters - ---------- - connection : str or :class:`DataprocChannelBuilder` - Connection string that is used to extract the connection parameters - and configure the GRPC connection. Or instance of ChannelBuilder / - DataprocChannelBuilder that creates GRPC connection. - user_id : str, optional - If not set, will default to the $USER environment. Defining the user - ID as part of the connection string takes precedence. - """ - - super().__init__(connection, user_id) - - execute_plan_request_base_method = ( - self.client._execute_plan_request_with_metadata - ) - execute_base_method = self.client._execute - execute_and_fetch_as_iterator_base_method = ( - self.client._execute_and_fetch_as_iterator - ) - - def execute_plan_request_wrapped_method(*args, **kwargs): - req = execute_plan_request_base_method(*args, **kwargs) - if not req.operation_id: - req.operation_id = str(uuid.uuid4()) - logger.debug( - f"No operation_id found. Setting operation_id: {req.operation_id}" - ) - return req - - self.client._execute_plan_request_with_metadata = ( - execute_plan_request_wrapped_method - ) - - 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) - - self.client._execute = MethodType(execute_wrapped_method, self.client) - - 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 - ) - - # Patching clearProgressHandlers method to not remove Dataproc Progress Handler - clearProgressHandlers_base_method = self.clearProgressHandlers - - def clearProgressHandlers_wrapper_method(_, *args, **kwargs): - clearProgressHandlers_base_method(*args, **kwargs) - - self._register_progress_execution_handler() - - self.clearProgressHandlers = MethodType( - clearProgressHandlers_wrapper_method, self - ) - - @staticmethod - @functools.lru_cache(maxsize=1) - def get_tqdm_bar(): - """ - Return a tqdm implementation that works in the current environment. - - - Uses CLI tqdm for interactive terminals. - - Uses the notebook tqdm if available, otherwise falls back to CLI tqdm. - """ - from tqdm import tqdm as cli_tqdm - - if environment.is_interactive_terminal(): - return cli_tqdm - - try: - import ipywidgets - from tqdm.notebook import tqdm as notebook_tqdm - - return notebook_tqdm - except ImportError: - return cli_tqdm - - def _register_progress_execution_handler(self): - from pyspark.sql.connect.shell.progress import StageInfo - - def handler( - stages: Optional[Iterable[StageInfo]], - inflight_tasks: int, - operation_id: Optional[str], - done: bool, - ): - if operation_id is None: - return - - # Don't build / render progress bar for non-interactive (despite - # Ipython or non-IPython) - if not environment.is_interactive(): - return - - total_tasks = 0 - completed_tasks = 0 - - for stage in stages or []: - total_tasks += stage.num_tasks - completed_tasks += stage.num_completed_tasks - - # Don't show progress bar till we receive some tasks - if total_tasks == 0: - return - - # Get correct tqdm (notebook or CLI) - tqdm_pbar = self.get_tqdm_bar() - - # Use a lock to ensure only one thread can access and modify - # the shared dictionaries at a time. - with self._lock: - if operation_id in self._execution_progress_bar: - pbar = self._execution_progress_bar[operation_id] - if pbar.total != total_tasks: - pbar.reset( - total=total_tasks - ) # This force resets the progress bar % too on next refresh - else: - pbar = tqdm_pbar( - total=total_tasks, - leave=True, - dynamic_ncols=True, - bar_format="{l_bar}{bar} {n_fmt}/{total_fmt} Tasks", - ) - self._execution_progress_bar[operation_id] = pbar - - # To handle skipped or failed tasks. - # StageInfo proto doesn't have skipped and failed tasks information to process. - if done and completed_tasks < total_tasks: - completed_tasks = total_tasks - - pbar.n = completed_tasks - pbar.refresh() - - if done: - pbar.close() - self._execution_progress_bar.pop(operation_id, None) - - self.registerProgressHandler(handler) - - @staticmethod - def _sql_lazy_transformation(req): - # Select SQL command - try: - query = req.plan.command.sql_command.input.sql.query - return "select" in query.strip().lower().split() - except AttributeError: - return False - - def _repr_html_(self) -> str: - if not self._active_s8s_session_id: - return """ -
No Active Dataproc Session
- """ - - s8s_session = f"{_DATAPROC_SESSIONS_BASE_URL}/{self._region}/{self._active_s8s_session_id}" - ui = f"{s8s_session}/sparkApplications/applications" - return f""" -
-

Spark Connect

- -

Dataproc Session

-

Spark UI

-
- """ - - 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(): - return - - assert 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, - ] - ) - - url = ( - f"{_DATAPROC_SESSIONS_BASE_URL}/{self._region}/" - f"{self._active_s8s_session_id}/sparkApplications/application;" - f"associatedSqlOperationId={operation_id}?project={self._project_id}" - ) - - if environment.is_interactive_terminal(): - print(f"Spark Query: {url}") - return - - try: - from IPython.display import display, HTML - - html_element = f""" -
-

Spark Query (Operation: {operation_id})

-
- """ - display(HTML(html_element)) - except ImportError: - return - - @staticmethod - def _remove_stopped_session_from_file(): - file_path = DataprocSparkSession._get_active_session_file_path() - if file_path is not None: - try: - with open(file_path, "w"): - pass - except Exception as e: - logger.error( - f"Exception while removing active session in file {file_path}, {e}" - ) - - def addArtifacts( - self, - *artifact: str, - pyfile: bool = False, - archive: bool = False, - file: bool = False, - pypi: bool = False, - ) -> None: - """ - Add artifact(s) to the client session. Currently only local files & pypi installations are supported. - - .. versionadded:: 3.5.0 - - Parameters - ---------- - *artifact : tuple of str - Artifact's URIs to add. - pyfile : bool - Whether to add them as Python dependencies such as .py, .egg, .zip or .jar files. - The pyfiles are directly inserted into the path when executing Python functions - in executors. - archive : bool - Whether to add them as archives such as .zip, .jar, .tar.gz, .tgz, or .tar files. - The archives are unpacked on the executor side automatically. - file : bool - Add a file to be downloaded with this Spark job on every node. - The ``path`` passed can only be a local file for now. - pypi : bool - This option is only available with DataprocSparkSession. e.g. `spark.addArtifacts("spacy==3.8.4", "torch", pypi=True)` - Installs PyPi package (with its dependencies) in the active Spark session on the driver and executors. - - Notes - ----- - This is an API dedicated to Spark Connect client only. With regular Spark Session, it throws - an exception. - Regarding pypi: Popular packages are already pre-installed in s8s runtime. - https://cloud.google.com/dataproc-serverless/docs/concepts/versions/spark-runtime-2.3#python_libraries - If there are conflicts/package doesn't exist, it throws an exception. - """ - if sum([pypi, file, pyfile, archive]) > 1: - raise ValueError( - "'pyfile', 'archive', 'file' and/or 'pypi' cannot be True together." - ) - if pypi: - artifacts = PyPiArtifacts(set(artifact)) - logger.debug("Making addArtifact call to install packages") - self.addArtifact( - artifacts.write_packages_config(self._active_s8s_session_uuid), - file=True, - ) - else: - super().addArtifacts( - *artifact, pyfile=pyfile, archive=archive, file=file - ) - - @staticmethod - def _get_active_session_file_path(): - return os.getenv("DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH") - - def stop(self, terminate: Optional[bool] = None) -> None: - """ - Stop the Spark session and optionally terminate the server-side session. - - Parameters - ---------- - terminate : bool, optional - Control server-side termination behavior. - - - None (default): Auto-detect based on session type - - - Managed sessions (auto-generated ID): terminate server - - Named sessions (custom ID): client-side cleanup only - - - True: Always terminate the server-side session - - False: Never terminate the server-side session (client cleanup only) - - Examples - -------- - Auto-detect termination behavior (existing behavior): - - >>> spark.stop() - - Force terminate a named session: - - >>> spark.stop(terminate=True) - - Prevent termination of a managed session: - - >>> spark.stop(terminate=False) - """ - with DataprocSparkSession._lock: - if DataprocSparkSession._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 DataprocSparkSession._active_session_uses_custom_id - ) - else: - should_terminate = terminate - - if should_terminate: - # Terminate the server-side session - logger.debug( - f"Terminating session {DataprocSparkSession._active_s8s_session_id}" - ) - terminate_s8s_session( - DataprocSparkSession._project_id, - DataprocSparkSession._region, - DataprocSparkSession._active_s8s_session_id, - self._client_options, - ) - else: - # Client-side cleanup only - logger.debug( - f"Stopping session {DataprocSparkSession._active_s8s_session_id} without termination" - ) - - self._remove_stopped_session_from_file() - - # Clean up SparkSession._instantiatedSession if it points to this session - try: - from pyspark.sql import SparkSession as PySparkSQLSession - - if PySparkSQLSession._instantiatedSession is self: - PySparkSQLSession._instantiatedSession = None - logger.debug( - "Cleared SparkSession._instantiatedSession reference" - ) - except (ImportError, AttributeError): - # PySpark not available or _instantiatedSession doesn't exist - pass - - DataprocSparkSession._active_s8s_session_uuid = None - DataprocSparkSession._active_s8s_session_id = None - DataprocSparkSession._active_session_uses_custom_id = False - DataprocSparkSession._project_id = None - DataprocSparkSession._region = None - DataprocSparkSession._client_options = None - - self.client.close() - if self is DataprocSparkSession._default_session: - DataprocSparkSession._default_session = None - if self is getattr( - DataprocSparkSession._active_session, "session", None - ): - DataprocSparkSession._active_session.session = None - - -def terminate_s8s_session( - project_id, region, active_s8s_session_id, client_options=None -): - from google.cloud.dataproc_v1 import SessionControllerClient - - logger.debug(f"Terminating Dataproc Session: {active_s8s_session_id}") - terminate_session_request = TerminateSessionRequest() - session_name = f"projects/{project_id}/locations/{region}/sessions/{active_s8s_session_id}" - terminate_session_request.name = session_name - state = None - try: - session_client = SessionControllerClient(client_options=client_options) - session_client.terminate_session(terminate_session_request) - get_session_request = GetSessionRequest() - get_session_request.name = session_name - state = Session.State.ACTIVE - while ( - state != Session.State.TERMINATING - and state != Session.State.TERMINATED - and state != Session.State.FAILED - ): - session = session_client.get_session(get_session_request) - state = session.state - time.sleep(1) - except NotFound: - logger.debug( - f"{active_s8s_session_id} Dataproc Session already deleted" - ) - # Client will get 'Aborted' error if session creation is still in progress and - # 'FailedPrecondition' if another termination is still in progress. - # Both are retryable, but we catch it and let TTL take care of cleanups. - except (FailedPrecondition, Aborted): - logger.debug( - f"{active_s8s_session_id} Dataproc Session already terminated manually or automatically due to TTL" - ) - if state is not None and state == Session.State.FAILED: - raise RuntimeError("Dataproc Session termination failed") - - -def get_active_s8s_session_response( - session_name, client_options -) -> Optional[sessions.Session]: - get_session_request = GetSessionRequest() - get_session_request.name = session_name - try: - get_session_response = SessionControllerClient( - client_options=client_options - ).get_session(get_session_request) - state = get_session_response.state - except Exception as e: - print(f"{session_name} Dataproc Session deleted: {e}") - return None - if state is not None and ( - state == Session.State.ACTIVE or state == Session.State.CREATING - ): - return get_session_response - return None - - -def is_s8s_session_active(session_name, client_options) -> bool: - if get_active_s8s_session_response(session_name, client_options) is None: - return False - return True +DataprocSparkSession = ManagedSparkSession diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py new file mode 100644 index 00000000..70ff97c2 --- /dev/null +++ b/google/cloud/managed_spark_connect/__init__.py @@ -0,0 +1,30 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import importlib.metadata +import warnings + +from .session import ManagedSparkSession + +old_package_names = ["google-spark-connect", "dataproc-spark-connect"] +current_package_name = "managed-spark-connect" +for old_package_name in old_package_names: + try: + importlib.metadata.distribution(old_package_name) + 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}'." + ) + except Exception: + pass diff --git a/google/cloud/managed_spark_connect/client/__init__.py b/google/cloud/managed_spark_connect/client/__init__.py new file mode 100644 index 00000000..4ddcf7d1 --- /dev/null +++ b/google/cloud/managed_spark_connect/client/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from .core import ManagedSparkChannelBuilder diff --git a/google/cloud/managed_spark_connect/client/core.py b/google/cloud/managed_spark_connect/client/core.py new file mode 100644 index 00000000..843e741c --- /dev/null +++ b/google/cloud/managed_spark_connect/client/core.py @@ -0,0 +1,141 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import logging + +import google +import grpc +from pyspark.sql.connect.client import DefaultChannelBuilder + +from . import proxy + +logger = logging.getLogger(__name__) + + +class ManagedSparkChannelBuilder(DefaultChannelBuilder): + """ + This is a helper class that is used to create a GRPC channel based on the given + connection string per the documentation of Spark Connect. + + This implementation of ChannelBuilder uses `secure_authorized_channel` from the + `google.auth.transport.grpc` package for authenticating secure channel. + + Examples + -------- + >>> cb = ChannelBuilder("sc://localhost") + ... cb.endpoint + + >>> cb = ChannelBuilder("sc://localhost/;use_ssl=true;token=aaa") + ... cb.secure + True + """ + + def __init__(self, url, is_active_callback=None): + self._is_active_callback = is_active_callback + super().__init__(url) + + def toChannel(self) -> grpc.Channel: + """ + Applies the parameters of the connection string and creates a new + GRPC channel according to the configuration. Passes optional channel options to + construct the channel. + + Returns + ------- + GRPC Channel instance. + """ + # TODO: Replace with a direct channel once all compatibility issues with + # grpc have been resolved. + return self._proxied_channel() + + def _proxied_channel(self) -> grpc.Channel: + return ProxiedChannel(self.host, self._is_active_callback) + + def _direct_channel(self) -> grpc.Channel: + destination = f"{self.host}:{self.port}" + + credentials, project = google.auth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + # Get an HTTP request function to refresh credentials. + request = google.auth.transport.requests.Request() + # Create a channel. + + return google.auth.transport.grpc.secure_authorized_channel( + credentials, + request, + destination, + None, + None, + options=self._channel_options, + ) + + +class ProxiedChannel(grpc.Channel): + + def __init__(self, target_host, is_active_callback): + self._is_active_callback = is_active_callback + self._proxy = proxy.ManagedSparkSessionProxy(0, target_host) + self._proxy.start() + self._proxied_connect_url = f"sc://localhost:{self._proxy.port}" + self._wrapped = DefaultChannelBuilder( + self._proxied_connect_url + ).toChannel() + + def __enter__(self): + return self + + def __exit__(self, *args): + ret = self._wrapped.__exit__(*args) + self._proxy.stop() + return ret + + def close(self): + ret = self._wrapped.close() + self._proxy.stop() + return ret + + def _wrap_method(self, wrapped_method): + if self._is_active_callback is None: + return wrapped_method + + def checked_method(*margs, **mkwargs): + if ( + self._is_active_callback is not None + and not self._is_active_callback() + ): + logger.warning(f"Session is no longer active") + raise RuntimeError( + "Session not active. Please create a new session" + ) + return wrapped_method(*margs, **mkwargs) + + return checked_method + + def stream_stream(self, *args, **kwargs): + return self._wrap_method(self._wrapped.stream_stream(*args, **kwargs)) + + def stream_unary(self, *args, **kwargs): + return self._wrap_method(self._wrapped.stream_unary(*args, **kwargs)) + + def subscribe(self, *args, **kwargs): + return self._wrap_method(self._wrapped.subscribe(*args, **kwargs)) + + def unary_stream(self, *args, **kwargs): + return self._wrap_method(self._wrapped.unary_stream(*args, **kwargs)) + + def unary_unary(self, *args, **kwargs): + return self._wrap_method(self._wrapped.unary_unary(*args, **kwargs)) + + def unsubscribe(self, *args, **kwargs): + return self._wrap_method(self._wrapped.unsubscribe(*args, **kwargs)) diff --git a/google/cloud/managed_spark_connect/client/proxy.py b/google/cloud/managed_spark_connect/client/proxy.py new file mode 100755 index 00000000..cf680439 --- /dev/null +++ b/google/cloud/managed_spark_connect/client/proxy.py @@ -0,0 +1,269 @@ +#!/bin/env python + +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import contextlib +import logging +import socket +import threading + +import websockets.sync.client as websocketclient + +from google import auth as googleauth +from google.auth.transport import requests as googleauthrequests + +parser = argparse.ArgumentParser() +parser.add_argument("port") +parser.add_argument("target_host") + + +logger = logging.getLogger(__name__) + + +class bridged_socket(object): + """Socket-like object that uses a websocket-over-TCP Bridge transport. + + See: https://github.com/google/inverting-proxy/tree/master/utils/tcpbridge + """ + + def __init__(self, websocket_conn): + self._conn = websocket_conn + + def recv(self, buff_size): + # N.B. The websockets [recv method](https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.recv) + # does not support the buff_size parameter, but it does add a `timeout` keyword parameter not supported by normal + # socket objects. + # + # We set that timeout to 60 seconds to prevent any scenarios where we wind up stuck waiting for a message from a websocket connection + # that never comes. + msg = self._conn.recv(timeout=60) + return bytes.fromhex(msg) + + def send(self, msg_bytes): + msg = bytes.hex(msg_bytes) + self._conn.send(msg) + + def close(self): + return self._conn.close() + + +def connect_tcp_bridge(hostname): + """Create a socket-like connection to the given hostname using websocket. + + The backend server connected to over the websocket connection must be + running the TCP-bridge backend corresponding to this frontend. + + Args: + hostname: The hostname of the server running the TCP-bridge backend. + + Returns: + A socket-like object with `recv` and `send` methods. + """ + path = "tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96" + creds, _ = googleauth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + creds.refresh(googleauthrequests.Request()) + + return websocketclient.connect( + f"wss://{hostname}/{path}", + additional_headers={"Authorization": f"Bearer {creds.token}"}, + open_timeout=30, + ) + + +def forward_bytes(name, from_sock, to_sock): + """Continuously stream bytes from the `from_sock` to the `to_sock`. + + This method terminates when either the `from_sock` is closed (causing + it to return a Falsy value from its `recv` method), or the first time + it hits an exception. + + This method is intended to be run in a separate thread of execution. + + Args: + name: forwarding thread name + from_sock: A socket-like object to stream bytes from. + to_sock: A socket-like object to stream bytes to. + """ + while True: + try: + bs = from_sock.recv(1024) + if not bs: + to_sock.close() + return + attempt = 0 + while bs and (attempt < 10): + attempt += 1 + try: + to_sock.send(bs) + bs = None + except TimeoutError: + # On timeouts during a send, we retry just the send + # to make sure we don't lose any bytes. + pass + if bs: + raise Exception(f"Failed to forward bytes for {name}") + except TimeoutError: + # On timeouts during a receive, we retry the entire flow. + pass + except Exception as ex: + logger.debug(f"[{name}] Exception forwarding bytes: {ex}") + to_sock.close() + return + + +def connect_sockets(conn_number, from_sock, to_sock): + """Create a connection between the two given ports. + + This method continuously streams bytes in both directions between the + given `from_sock` and `to_sock` socket-like objects. + + The caller is responsible for creating and closing the supplied sockets. + """ + forward_name = f"{conn_number}-forward" + t1 = threading.Thread( + name=forward_name, + target=forward_bytes, + args=[forward_name, from_sock, to_sock], + daemon=True, + ) + t1.start() + backward_name = f"{conn_number}-backward" + t2 = threading.Thread( + name=backward_name, + target=forward_bytes, + args=[backward_name, to_sock, from_sock], + daemon=True, + ) + t2.start() + t1.join() + t2.join() + + +def forward_connection(conn_number, conn, addr, target_host): + """Create a connection to the target and forward `conn` to it. + + This method creates a socket-like object holding a connection to the given + target host, and then continuously streams bytes in both directions between + `conn` and that newly created connection. + + Both the supplied incoming connection (`conn`) and the created outgoing + connection are automatically closed when this method terminates. + + This method should be run inside a daemon thread so that it will not + block program termination. + """ + with conn: + with connect_tcp_bridge(target_host) as websocket_conn: + backend_socket = bridged_socket(websocket_conn) + # 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(10) + connect_sockets(conn_number, conn, backend_socket) + + +class ManagedSparkSessionProxy(object): + """A TCP proxy for forwarding requests to Dataproc Serverless Sessions. + + Spark Connect clients connect to this proxy using the h2c (without-SSL) + protocol, and this proxy adds SSL by tunneling those connections over + an HTTPS/WebSocket connection to the backend server. + + The tunneled requests are authenticated using the Google Application + Default Credentials. + """ + + def __init__(self, port, target_host): + self._port = port + self._target_host = target_host + self._started = False + self._killed = False + self._conn_number = 0 + + @property + def port(self): + """The local port the proxy is listening on""" + return self._port + + def start(self, daemon=True): + """Start the proxy. + + By the time this method returns the proxy has already started listening + on its local port will accept incoming connections. + """ + if self._started: + raise Exception("Managed Spark session proxy already started") + self._started = True + s = threading.Semaphore(value=0) + t = threading.Thread(target=self._run, args=[s], daemon=daemon) + t.start() + s.acquire() + + def _run(self, s): + with socket.create_server(("127.0.0.1", self._port)) as frontend_socket: + if self._port == 0: + self._port = frontend_socket.getsockname()[1] + s.release() + while not self._killed: + conn, addr = frontend_socket.accept() + 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], + daemon=True, + ).start() + + def stop(self): + """Stop the proxy.""" + self._killed = True + + +@contextlib.contextmanager +def managed_spark_session_proxy(port, target_host): + """Context manager for creating a Managed Spark session proxy. + + Usage: + with managed_spark_session_proxy(0, backend_hostname) as p: + local_port = p.port + ... + + Args: + port: The local port to listen on. Use `0` to pick a free port. + target_host: The backend to proxy connections to. + + Returns: + A context manager wrapping a ManagedSparkSessionProxy instance. + """ + proxy = ManagedSparkSessionProxy(port, target_host) + try: + proxy.start(daemon=False) + yield proxy + finally: + proxy.stop() + + +if __name__ == "__main__": + args = parser.parse_args() + with managed_spark_session_proxy(int(args.port), args.target_host) as p: + print(f"Proxy listening on port {p.port}") + try: + while True: + pass + except KeyboardInterrupt: + pass diff --git a/google/cloud/managed_spark_connect/environment.py b/google/cloud/managed_spark_connect/environment.py new file mode 100644 index 00000000..e19dd973 --- /dev/null +++ b/google/cloud/managed_spark_connect/environment.py @@ -0,0 +1,190 @@ +# Copyright 2025 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 +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +from typing import Callable, Tuple, List + + +def is_antigravity() -> bool: + """True if running inside the Antigravity environment.""" + return "antigravity" in os.getenv("__CFBundleIdentifier", "").lower() + + +def is_vscode() -> bool: + """True if running inside VS Code at all.""" + return os.getenv("VSCODE_PID") is not None + + +def is_jupyter() -> bool: + """True if running in a Jupyter environment.""" + return os.getenv("JPY_PARENT_PID") is not None + + +def is_colab_enterprise() -> bool: + """True if running in Colab Enterprise (Vertex AI).""" + return os.getenv("VERTEX_PRODUCT") == "COLAB_ENTERPRISE" + + +def is_colab() -> bool: + """True if running in Google Colab.""" + return os.getenv("COLAB_RELEASE_TAG") is not None + + +def is_workbench() -> bool: + """True if running in Vertex Workbench Instance (managed Jupyter).""" + return os.getenv("VERTEX_PRODUCT") == "WORKBENCH_INSTANCE" + + +def is_kaggle() -> bool: + """True if running in Kaggle Notebooks.""" + return os.getenv("KAGGLE_KERNEL_RUN_TYPE") is not None + + +def is_databricks() -> bool: + """True if running in Databricks.""" + return os.getenv("DATABRICKS_RUNTIME_VERSION") is not None + + +def is_sagemaker() -> bool: + """True if running in AWS SageMaker.""" + return os.getenv("SAGEMAKER_INTERNAL_IMAGE_URI") is not None + + +def is_deepnote() -> bool: + """True if running in Deepnote.""" + return os.getenv("DEEPNOTE_PROJECT_ID") is not None + + +def is_datalore() -> bool: + """True if running in JetBrains Datalore.""" + return os.getenv("DATALORE_USER") is not None + + +def is_spyder() -> bool: + """True if running inside Spyder IDE.""" + return any(k.startswith("SPYDER") for k in os.environ) + + +def is_cloud_shell() -> bool: + """True if running in Google Cloud Shell.""" + return os.getenv("CLOUD_SHELL") is not None + + +def is_codespaces() -> bool: + """True if running in GitHub Codespaces.""" + return os.getenv("CODESPACES") is not None + + +def is_jetbrains_ide() -> bool: + """True if running inside JetBrains IDE.""" + return ( + "jetbrains" in os.getenv("TERMINAL_EMULATOR", "").lower() + or "PYCHARM_HOSTED" in os.environ + ) + + +def is_hex() -> bool: + """True if running in Hex.""" + return os.getenv("HEX_PROJECT_ID") is not None + + +def is_polynote() -> bool: + """True if running in Polynote.""" + return os.getenv("POLYNOTE_VERSION") is not None + + +def is_eclipse() -> bool: + """True if running inside Eclipse IDE.""" + return "ECLIPSE_HOME" in os.environ or any( + k.startswith("ECLIPSE") for k in os.environ + ) + + +def is_interactive() -> bool: + try: + from IPython import get_ipython + + if get_ipython() is not None: + return True + except ImportError: + pass + + return hasattr(sys, "ps1") or bool(sys.flags.interactive) + + +def is_terminal() -> bool: + return sys.stdin.isatty() + + +def is_interactive_terminal() -> bool: + return is_interactive() and is_terminal() + + +def is_dataproc_batch() -> bool: + return os.getenv("DATAPROC_WORKLOAD_TYPE") == "batch" + + +def get_client_environment_label() -> str: + """ + Map current environment to a standardized client label. + + Priority order: + 1. Colab Enterprise ("colab-enterprise") + 2. Colab ("colab") + 3. Vertex Workbench Instance ("workbench-jupyter") + 4. Kaggle ("kaggle") + 5. AWS SageMaker ("sagemaker") + 6. Databricks ("databricks") + 7. Deepnote ("deepnote") + 8. JetBrains Datalore ("datalore") + 9. GitHub Codespaces ("codespaces") + 10. Google Cloud Shell ("cloud-shell") + 11. Hex ("hex") + 12. Polynote ("polynote") + 13. Antigravity ("antigravity") + 14. VS Code ("vscode") + 15. JetBrains IDE ("jetbrains") + 16. Spyder ("spyder") + 17. Eclipse ("eclipse") + 18. Jupyter ("jupyter") + 19. Unknown ("unknown") + """ + checks: List[Tuple[Callable[[], bool], str]] = [ + (is_colab_enterprise, "colab-enterprise"), + (is_colab, "colab"), + (is_workbench, "workbench-jupyter"), + (is_kaggle, "kaggle"), + (is_sagemaker, "sagemaker"), + (is_databricks, "databricks"), + (is_deepnote, "deepnote"), + (is_datalore, "datalore"), + (is_codespaces, "codespaces"), + (is_cloud_shell, "cloud-shell"), + (is_hex, "hex"), + (is_polynote, "polynote"), + (is_antigravity, "antigravity"), + (is_vscode, "vscode"), + (is_jetbrains_ide, "jetbrains"), + (is_spyder, "spyder"), + (is_eclipse, "eclipse"), + (is_jupyter, "jupyter"), + ] + for detector, label in checks: + try: + if detector(): + return label + except Exception: + pass + return "unknown" diff --git a/google/cloud/managed_spark_connect/exceptions.py b/google/cloud/managed_spark_connect/exceptions.py new file mode 100644 index 00000000..afdea9e2 --- /dev/null +++ b/google/cloud/managed_spark_connect/exceptions.py @@ -0,0 +1,27 @@ +# Copyright 2025 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. + + +class ManagedSparkConnectException(Exception): + """A custom exception class to only print the error messages. + This would be used for exceptions where the stack trace + doesn't provide any additional information. + """ + + def __init__(self, message): + self.message = message + super().__init__(message) + + def _render_traceback_(self): + return [self.message] diff --git a/google/cloud/managed_spark_connect/pypi_artifacts.py b/google/cloud/managed_spark_connect/pypi_artifacts.py new file mode 100644 index 00000000..b31d3864 --- /dev/null +++ b/google/cloud/managed_spark_connect/pypi_artifacts.py @@ -0,0 +1,48 @@ +import json +import logging +import os +import tempfile + +from packaging.requirements import Requirement + +logger = logging.getLogger(__name__) + + +class PyPiArtifacts: + """ + This is a helper class to serialize the PYPI package installation request with a "magic" file name + that Spark Connect server understands + """ + + @staticmethod + 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") + return reqs + + def __init__(self, packages: set[str]): + self.requirements = PyPiArtifacts.__try_parsing_package(packages) + + def write_packages_config(self, s8s_session_uuid: str) -> str: + """ + Can't use the same file-name as Spark throws exception that file already exists + Keep the filename/format in sync with server + """ + dependencies = { + "version": "0.5", + "packageType": "PYPI", + "packages": [str(req) for req in self.requirements], + } + + file_path = os.path.join( + tempfile.gettempdir(), + s8s_session_uuid, + "add-artifacts-1729-" + self.__str__() + ".json", + ) + + os.makedirs(os.path.dirname(file_path), exist_ok=True) + with open(file_path, "w") as json_file: + json.dump(dependencies, json_file, indent=4) + logger.debug("Dumping dependencies request in file: " + file_path) + return file_path diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py new file mode 100644 index 00000000..12952351 --- /dev/null +++ b/google/cloud/managed_spark_connect/session.py @@ -0,0 +1,1440 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import atexit +import datetime +import functools +import json +import logging +import os +import random +import re +import string +import threading +import time +import uuid +import warnings +import tqdm +from packaging import version +from types import MethodType +from typing import Any, cast, ClassVar, Dict, Iterable, Optional, Union + +from google.api_core import retry +from google.api_core.client_options import ClientOptions +from google.api_core.exceptions import ( + Aborted, + FailedPrecondition, + InvalidArgument, + NotFound, + PermissionDenied, +) +from google.api_core.future.polling import POLLING_PREDICATE +from google.auth.exceptions import DefaultCredentialsError +from google.cloud.managed_spark_connect.client import ManagedSparkChannelBuilder +from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException +from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts +from google.cloud.dataproc_v1 import ( + AuthenticationConfig, + CreateSessionRequest, + DeleteSessionRequest, + GetSessionRequest, + Session, + SessionControllerClient, + TerminateSessionRequest, +) +from google.cloud.dataproc_v1.types import sessions +from google.cloud.managed_spark_connect import environment +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 +SYSTEM_LABELS = { + "dataproc-session-client", + "goog-colab-notebook-id", +} + +_MANAGED_SPARK_SESSIONS_BASE_URL = ( + "https://console.cloud.google.com/dataproc/interactive" +) + + +def _env_var_set(new_name: str, old_name: str) -> bool: + return new_name in os.environ or old_name in os.environ + + +def _getenv_with_deprecated_alias( + new_name: str, old_name: str, default: Optional[str] = None +) -> Optional[str]: + if new_name in os.environ: + return os.environ[new_name] + if old_name in os.environ: + warnings.warn( + f"Environment variable '{old_name}' is deprecated, use '{new_name}' instead.", + DeprecationWarning, + stacklevel=2, + ) + return os.environ[old_name] + return default + + +def _is_valid_label_value(value: str) -> bool: + """ + Validates if a string complies with Google Cloud label value format. + Only lowercase letters, numbers, and dashes are allowed. + The value must start with lowercase letter or number and end with a lowercase letter or number. + Maximum length is 63 characters. + """ + if not value: + return False + + # Check maximum length (63 characters for GCP label values) + if len(value) > 63: + return False + + # Check if the value matches the pattern: starts and ends with alphanumeric, + # contains only lowercase letters, numbers, and dashes + pattern = r"^[a-z0-9]([a-z0-9-]*[a-z0-9])?$" + return bool(re.match(pattern, value)) + + +def _is_valid_session_id(session_id: str) -> bool: + """ + Validates if a string complies with Google Cloud session ID format. + - Must be 4-63 characters + - Only lowercase letters, numbers, and dashes are allowed + - Must start with a lowercase letter + - Cannot end with a dash + """ + if not session_id: + return False + + # The pattern is sufficient for validation and already enforces length constraints. + pattern = r"^[a-z][a-z0-9-]{2,61}[a-z0-9]$" + return bool(re.match(pattern, session_id)) + + +class ManagedSparkSession(SparkSession): + """The entry point to programming Spark with the Dataset and DataFrame API. + + A ManagedSparkSession can be used to create :class:`DataFrame`, register :class:`DataFrame` as + tables, execute SQL over tables, cache tables, and read parquet files. + + Examples + -------- + + Create a Spark session with Managed Spark Connect. + + >>> spark = ( + ... ManagedSparkSession.builder + ... .appName("Word Count") + ... .dataprocSessionConfig(Session()) + ... .getOrCreate() + ... ) # doctest: +SKIP + """ + + _DEFAULT_RUNTIME_VERSION = "3.0" + _MIN_RUNTIME_VERSION = "3.0" + + _active_s8s_session_uuid: ClassVar[Optional[str]] = None + _project_id = None + _region = None + _client_options = None + _active_s8s_session_id: ClassVar[Optional[str]] = None + _active_session_uses_custom_id: ClassVar[bool] = False + _execution_progress_bar = dict() + + class Builder(SparkSession.Builder): + + def __init__(self): + self._options: Dict[str, Any] = {} + self._channel_builder: Optional[ManagedSparkChannelBuilder] = None + self._dataproc_config: Optional[Session] = None + self._custom_session_id: Optional[str] = None + self._project_id = os.getenv("GOOGLE_CLOUD_PROJECT") + self._region = os.getenv("GOOGLE_CLOUD_REGION") + self._client_options = ClientOptions( + api_endpoint=os.getenv( + "GOOGLE_CLOUD_DATAPROC_API_ENDPOINT", + f"{self._region}-dataproc.googleapis.com", + ) + ) + self._session_controller_client: Optional[ + SessionControllerClient + ] = None + + @property + def session_controller_client(self) -> SessionControllerClient: + """Get or create a SessionControllerClient instance.""" + if self._session_controller_client is None: + self._session_controller_client = SessionControllerClient( + client_options=self._client_options + ) + return self._session_controller_client + + def projectId(self, project_id): + self._project_id = project_id + return self + + def location(self, location): + self._region = location + self._client_options.api_endpoint = os.getenv( + "GOOGLE_CLOUD_DATAPROC_API_ENDPOINT", + f"{self._region}-dataproc.googleapis.com", + ) + return self + + def dataprocSessionId(self, session_id: str): + """ + Set a custom session ID for creating or reusing sessions. + + The session ID must: + - Be 4-63 characters long + - Start with a lowercase letter + - Contain only lowercase letters, numbers, and hyphens + - Not end with a hyphen + + Args: + session_id: The custom session ID to use + + Returns: + This Builder instance for method chaining + + Raises: + ValueError: If the session ID format is invalid + """ + if not _is_valid_session_id(session_id): + raise ValueError( + f"Invalid session ID: '{session_id}'. " + "Session ID must be 4-63 characters, start with a lowercase letter, " + "contain only lowercase letters, numbers, and hyphens, " + "and not end with a hyphen." + ) + self._custom_session_id = session_id + return self + + def dataprocSessionConfig(self, dataproc_config: Session): + self._dataproc_config = dataproc_config + for k, v in dataproc_config.runtime_config.properties.items(): + self._options[cast(str, k)] = to_str(v) + return self + + @property + def dataproc_config(self): + with self._lock: + self._dataproc_config = self._dataproc_config or Session() + return self._dataproc_config + + def runtimeVersion(self, version: str): + self.dataproc_config.runtime_config.version = version + return self + + def serviceAccount(self, account: str): + self.dataproc_config.environment_config.execution_config.service_account = ( + account + ) + return self + + def subnetwork(self, subnet: str): + self.dataproc_config.environment_config.execution_config.subnetwork_uri = ( + subnet + ) + return self + + def ttl(self, duration: datetime.timedelta): + """Set the time-to-live (TTL) for the session using a timedelta object.""" + return self.ttlSeconds(int(duration.total_seconds())) + + def ttlSeconds(self, seconds: int): + """Set the time-to-live (TTL) for the session in seconds.""" + self.dataproc_config.environment_config.execution_config.ttl = { + "seconds": seconds + } + return self + + def idleTtl(self, duration: datetime.timedelta): + """Set the idle time-to-live (idle TTL) for the session using a timedelta object.""" + return self.idleTtlSeconds(int(duration.total_seconds())) + + def idleTtlSeconds(self, seconds: int): + """Set the idle time-to-live (idle TTL) for the session in seconds.""" + self.dataproc_config.environment_config.execution_config.idle_ttl = { + "seconds": seconds + } + return self + + def runtimeProfile(self, profile: str): + """Set the Runtime Profile to use for the session.""" + self.dataproc_config.session_template = profile + return self + + def sessionTemplate(self, template: str): + """Deprecated: use :meth:`runtimeProfile` instead.""" + warnings.warn( + "sessionTemplate() is deprecated, use runtimeProfile() instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.runtimeProfile(template) + + def label(self, key: str, value: str): + """Add a single label to the session.""" + return self.labels({key: value}) + + def labels(self, labels: Dict[str, str]): + # Filter out system labels and warn user + filtered_labels = {} + for key, value in labels.items(): + if key in SYSTEM_LABELS: + logger.warning( + f"Label '{key}' is a system label and cannot be overridden by user. Ignoring." + ) + else: + filtered_labels[key] = value + + self.dataproc_config.labels.update(filtered_labels) + return self + + def remote(self, url: Optional[str] = None) -> "SparkSession.Builder": + if url: + raise NotImplemented( + "ManagedSparkSession does not support connecting to an existing remote server" + ) + else: + return self + + def create(self) -> "ManagedSparkSession": + raise NotImplemented( + "ManagedSparkSession allows session creation only through getOrCreate" + ) + + def __create_spark_connect_session_from_s8s( + self, session_response, session_name + ) -> "ManagedSparkSession": + ManagedSparkSession._active_s8s_session_uuid = session_response.uuid + ManagedSparkSession._project_id = self._project_id + ManagedSparkSession._region = self._region + ManagedSparkSession._client_options = self._client_options + spark_connect_url = session_response.runtime_info.endpoints.get( + "Spark Connect Server" + ) + url = f"{spark_connect_url}/;session_id={session_response.uuid};use_ssl=true" + logger.debug(f"Spark Connect URL: {url}") + self._channel_builder = ManagedSparkChannelBuilder( + url, + is_active_callback=lambda: is_s8s_session_active( + session_name, self._client_options + ), + ) + + assert self._channel_builder is not None + session = ManagedSparkSession(connection=self._channel_builder) + + # Register handler for Cell Execution Progress bar + session._register_progress_execution_handler() + + ManagedSparkSession._set_default_and_active_session(session) + + return session + + def __create(self) -> "ManagedSparkSession": + with self._lock: + + if self._options.get("spark.remote", False): + raise NotImplemented( + "ManagedSparkSession does not support connecting to an existing Spark Connect remote server" + ) + + from google.cloud.dataproc_v1 import SessionControllerClient + + dataproc_config: Session = self._get_dataproc_config() + + # Check runtime version compatibility before creating session + self._check_runtime_compatibility(dataproc_config) + + # Use custom session ID if provided, otherwise generate one + session_id = ( + self._custom_session_id + if self._custom_session_id + else self.generate_session_id() + ) + + dataproc_config.name = f"projects/{self._project_id}/locations/{self._region}/sessions/{session_id}" + logger.debug( + f"Managed Spark Session configuration:\n{dataproc_config}" + ) + + session_request = CreateSessionRequest() + session_request.session_id = session_id + session_request.session = dataproc_config + session_request.parent = ( + f"projects/{self._project_id}/locations/{self._region}" + ) + + logger.debug("Creating Managed Spark Session") + ManagedSparkSession._active_s8s_session_id = session_id + # Track whether this session uses a custom ID (unmanaged) or auto-generated ID (managed) + ManagedSparkSession._active_session_uses_custom_id = ( + self._custom_session_id is not None + ) + s8s_creation_start_time = time.time() + + stop_create_session_pbar_event = threading.Event() + + 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 + if i >= iterations - 1: + # Sleep until session created + while not stop_create_session_pbar_event.is_set(): + time.sleep(1) + else: + time.sleep(1) + + 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" + + try: + if ( + _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT", + "DATAPROC_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT", + "false", + ) + == "true" + ): + atexit.register( + lambda: terminate_s8s_session( + self._project_id, + self._region, + session_id, + self._client_options, + ) + ) + operation = SessionControllerClient( + client_options=self._client_options + ).create_session(session_request) + self._display_session_link_on_creation(session_id) + self._display_view_session_details_button(session_id) + create_session_pbar_thread.start() + session_response: Session = operation.result( + polling=retry.Retry( + predicate=POLLING_PREDICATE, + initial=5.0, # seconds + maximum=5.0, # seconds + multiplier=1.0, + timeout=600, # seconds + ) + ) + stop_create_session_pbar_event.set() + create_session_pbar_thread.join() + self._print_session_created_message() + file_path = ( + ManagedSparkSession._get_active_session_file_path() + ) + if file_path is not None: + try: + session_data = { + "session_name": session_response.name, + "session_uuid": session_response.uuid, + } + os.makedirs( + os.path.dirname(file_path), exist_ok=True + ) + with open(file_path, "w") as json_file: + json.dump(session_data, json_file, indent=4) + except Exception as e: + logger.error( + 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( + f"Error while creating Managed Spark Session" + ) from e + finally: + stop_create_session_pbar_event.set() + + logger.debug( + f"Managed Spark Session created: {session_id} in {int(time.time() - s8s_creation_start_time)} seconds" + ) + return self.__create_spark_connect_session_from_s8s( + session_response, dataproc_config.name + ) + + def _wait_for_session_available( + self, session_name: str, timeout: int = 300 + ) -> Session: + start_time = time.time() + while time.time() - start_time < timeout: + try: + session = self.session_controller_client.get_session( + name=session_name + ) + if "Spark Connect Server" in session.runtime_info.endpoints: + return session + time.sleep(5) + except Exception as e: + logger.warning( + f"Error while polling for Spark Connect endpoint: {e}" + ) + time.sleep(5) + raise RuntimeError( + f"Spark Connect endpoint not available for session {session_name} after {timeout} seconds." + ) + + def _display_session_link_on_creation(self, session_id): + session_url = f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/{session_id}?project={self._project_id}" + plain_message = ( + f"Creating Managed Spark Connect Session: {session_url}" + ) + if environment.is_colab_enterprise(): + html_element = f""" +
+

Creating Managed Spark Connect Session

+

+ """ + else: + html_element = f""" +
+

Creating Managed Spark Connect Session

+

Managed Spark Session

+
+ """ + self._output_element_or_message(plain_message, html_element) + + def _print_session_created_message(self): + plain_message = f"Managed Spark Session was successfully created" + html_element = f"

{plain_message}

" + + self._output_element_or_message(plain_message, html_element) + + def _output_element_or_message(self, plain_message, html_element): + """ + Display / print the needed rich HTML element or plain text depending + on whether rich element is supported or not. + + :param plain_message: Message to print on non-IPython or + non-interactive shell + :param html_element: HTML element to display for interactive IPython + environment + """ + # Don't print any output (Rich or Plain) for non-interactive + if not environment.is_interactive(): + return + + if environment.is_interactive_terminal(): + print(plain_message) + return + + try: + from IPython.display import display, HTML + + display(HTML(html_element)) + except ImportError: + print(plain_message) + + def _get_exiting_active_session( + self, + ) -> Optional["ManagedSparkSession"]: + s8s_session_id = ManagedSparkSession._active_s8s_session_id + session_name = f"projects/{self._project_id}/locations/{self._region}/sessions/{s8s_session_id}" + session_response = None + session = None + if s8s_session_id is not None: + session_response = get_active_s8s_session_response( + session_name, self._client_options + ) + session = ManagedSparkSession.getActiveSession() + + if session is None: + session = ManagedSparkSession._default_session + + if session_response is not None: + print( + f"Using existing Managed Spark Session (configuration changes may not be applied): {_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/{s8s_session_id}?project={self._project_id}" + ) + self._display_view_session_details_button(s8s_session_id) + if session is None: + session_response = self._wait_for_session_available( + session_name + ) + session = self.__create_spark_connect_session_from_s8s( + session_response, session_name + ) + 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" + ) + session.stop() + + return None + + def getOrCreate(self) -> "ManagedSparkSession": + with ManagedSparkSession._lock: + if environment.is_dataproc_batch(): + # For Dataproc batch workloads, connect to the already initialized local SparkSession + from pyspark.sql import SparkSession as PySparkSQLSession + + session = PySparkSQLSession.builder.getOrCreate() + return session # type: ignore + + if self._project_id is None: + raise ManagedSparkConnectException( + f"Error while creating Managed Spark Session: project ID is not set" + ) + + if self._region is None: + raise ManagedSparkConnectException( + f"Error while creating Managed Spark Session: location is not set" + ) + + # Handle custom session ID by setting it early and letting existing logic handle it + if self._custom_session_id: + self._handle_custom_session_id() + + session = self._get_exiting_active_session() + if session is None: + session = self.__create() + + # Register this session as the instantiated SparkSession for compatibility + # with tools and libraries that expect SparkSession._instantiatedSession + from pyspark.sql import SparkSession as PySparkSQLSession + + PySparkSQLSession._instantiatedSession = session + + return session + + def _handle_custom_session_id(self): + """Handle custom session ID by checking if it exists and setting _active_s8s_session_id.""" + session_response = self._get_session_by_id(self._custom_session_id) + if session_response is not None: + # Found an active session with the custom ID, set it as the active session + ManagedSparkSession._active_s8s_session_id = ( + self._custom_session_id + ) + # Mark that this session uses a custom ID + ManagedSparkSession._active_session_uses_custom_id = True + else: + # No existing session found, clear any existing active session ID + # so we'll create a new one with the custom ID + ManagedSparkSession._active_s8s_session_id = None + + def _get_dataproc_config(self): + # Use the property to ensure we always have a config + dataproc_config = self.dataproc_config + for k, v in self._options.items(): + dataproc_config.runtime_config.properties[k] = v + dataproc_config.spark_connect_session = ( + sessions.SparkConnectConfig() + ) + if not dataproc_config.runtime_config.version: + dataproc_config.runtime_config.version = ( + ManagedSparkSession._DEFAULT_RUNTIME_VERSION + ) + + # Check for Python version mismatch with runtime for UDF compatibility + self._check_python_version_compatibility( + dataproc_config.runtime_config.version + ) + + # Use local variable to improve readability of deeply nested attribute access + exec_config = dataproc_config.environment_config.execution_config + + # Set service account from environment if not already set + if not exec_config.service_account and _env_var_set( + "MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT", + "DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT", + ): + exec_config.service_account = _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT", + "DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT", + ) + + # Auto-set authentication type to SERVICE_ACCOUNT when service account is provided + if exec_config.service_account: + # When service account is provided, explicitly set auth type to SERVICE_ACCOUNT + exec_config.authentication_config.user_workload_authentication_type = ( + AuthenticationConfig.AuthenticationType.SERVICE_ACCOUNT + ) + elif ( + not exec_config.authentication_config.user_workload_authentication_type + and _env_var_set( + "MANAGED_SPARK_CONNECT_AUTH_TYPE", + "DATAPROC_SPARK_CONNECT_AUTH_TYPE", + ) + ): + # Only set auth type from environment if no service account is present + exec_config.authentication_config.user_workload_authentication_type = AuthenticationConfig.AuthenticationType[ + _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_AUTH_TYPE", + "DATAPROC_SPARK_CONNECT_AUTH_TYPE", + ) + ] + if ( + not dataproc_config.environment_config.execution_config.subnetwork_uri + and _env_var_set( + "MANAGED_SPARK_CONNECT_SUBNET", + "DATAPROC_SPARK_CONNECT_SUBNET", + ) + ): + dataproc_config.environment_config.execution_config.subnetwork_uri = _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_SUBNET", + "DATAPROC_SPARK_CONNECT_SUBNET", + ) + if ( + not dataproc_config.environment_config.execution_config.ttl + and _env_var_set( + "MANAGED_SPARK_CONNECT_TTL_SECONDS", + "DATAPROC_SPARK_CONNECT_TTL_SECONDS", + ) + ): + dataproc_config.environment_config.execution_config.ttl = { + "seconds": int( + _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_TTL_SECONDS", + "DATAPROC_SPARK_CONNECT_TTL_SECONDS", + ) + ) + } + if ( + not dataproc_config.environment_config.execution_config.idle_ttl + and _env_var_set( + "MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS", + "DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS", + ) + ): + dataproc_config.environment_config.execution_config.idle_ttl = { + "seconds": int( + _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS", + "DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS", + ) + ) + } + client_environment = environment.get_client_environment_label() + dataproc_config.labels["dataproc-session-client"] = ( + client_environment + ) + if "COLAB_NOTEBOOK_ID" in os.environ: + colab_notebook_name = os.environ["COLAB_NOTEBOOK_ID"] + # Extract the last part of the path, which is the ID + notebook_id = os.path.basename(colab_notebook_name) + if _is_valid_label_value(notebook_id): + dataproc_config.labels["goog-colab-notebook-id"] = ( + notebook_id + ) + else: + logger.warning( + f"Warning while processing notebook ID: Notebook ID '{notebook_id}' is not compliant with label value format. " + f"Only lowercase letters, numbers, and dashes are allowed. " + f"The value must start with lowercase letter or number and end with a lowercase letter or number. " + f"Maximum length is 63 characters. " + f"Ignoring notebook ID label." + ) + default_datasource = _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE", + "DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE", + ) + match default_datasource: + case "bigquery": + # Merge default configs with existing properties, + # user configs take precedence + for k, v in { + "spark.sql.catalog.spark_catalog": "com.google.cloud.spark.bigquery.BigQuerySparkSessionCatalog", + "spark.sql.sources.default": "bigquery", + }.items(): + if k not in dataproc_config.runtime_config.properties: + dataproc_config.runtime_config.properties[k] = v + case _: + if default_datasource: + logger.warning( + f"MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE is set to an invalid value:" + f" {default_datasource}. Supported value is 'bigquery'." + ) + + return dataproc_config + + def _check_python_version_compatibility(self, runtime_version): + """Check if client Python version matches server Python version for UDF compatibility.""" + import sys + import warnings + + # Runtime version to server Python version mapping + RUNTIME_PYTHON_MAP = { + "3.0": (3, 12), + } + + client_python = sys.version_info[:2] # (major, minor) + + if runtime_version in RUNTIME_PYTHON_MAP: + server_python = RUNTIME_PYTHON_MAP[runtime_version] + + if client_python != server_python: + warnings.warn( + f"Python version mismatch detected: Client is using Python {client_python[0]}.{client_python[1]}, " + f"but Managed Spark runtime {runtime_version} uses Python {server_python[0]}.{server_python[1]}. " + f"This mismatch may cause issues with Python UDF (User Defined Function) compatibility. " + f"Consider using Python {server_python[0]}.{server_python[1]} for optimal UDF execution.", + stacklevel=3, + ) + + def _check_runtime_compatibility(self, dataproc_config): + """Check if runtime version 3.0 client is compatible with older runtime versions. + + Runtime version 3.0 clients do not support older runtime versions (pre-3.0). + There is no backward or forward compatibility between different runtime versions. + + Args: + dataproc_config: The Session configuration containing runtime version + + Raises: + ManagedSparkConnectException: If server is using pre-3.0 runtime version + """ + runtime_version = dataproc_config.runtime_config.version + + if not runtime_version: + return + + logger.debug(f"Detected server runtime version: {runtime_version}") + + # Parse runtime version to check if it's below minimum supported version + try: + server_version = version.parse(runtime_version) + min_version = version.parse( + ManagedSparkSession._MIN_RUNTIME_VERSION + ) + + if server_version < min_version: + raise ManagedSparkConnectException( + f"Specified {runtime_version} Managed Spark Runtime version is not supported, " + f"use {ManagedSparkSession._MIN_RUNTIME_VERSION} version or higher." + ) + except version.InvalidVersion: + # If we can't parse the version, log a warning but continue + logger.warning( + f"Could not parse runtime version: {runtime_version}" + ) + + def _display_view_session_details_button(self, session_id): + # Display button is only supported in colab enterprise + if not environment.is_colab_enterprise(): + return + + # Skip button display for colab enterprise IPython terminals + if environment.is_interactive_terminal(): + return + + try: + session_url = f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/{session_id}?project={self._project_id}" + from IPython.core.interactiveshell import InteractiveShell + + if not InteractiveShell.initialized(): + return + + from google.cloud.aiplatform.utils import _ipython_utils + + _ipython_utils.display_link( + "View Session Details", f"{session_url}", "dashboard" + ) + except ImportError as e: + logger.debug(f"Import error: {e}") + + def _get_session_by_id(self, session_id: str) -> Optional[Session]: + """ + Get existing session by ID. + + Returns: + Session if ACTIVE/CREATING, None if not found or not usable + """ + session_name = f"projects/{self._project_id}/locations/{self._region}/sessions/{session_id}" + + try: + get_request = GetSessionRequest(name=session_name) + session = self.session_controller_client.get_session( + get_request + ) + + logger.debug( + f"Found existing session {session_id} in state: {session.state}" + ) + + if session.state in [ + Session.State.ACTIVE, + Session.State.CREATING, + ]: + # Reuse the active session + logger.info(f"Reusing existing session: {session_id}") + return session + else: + # Session exists but is not usable (terminated/failed/terminating) + logger.info( + f"Session {session_id} in {session.state.name} state, cannot reuse" + ) + return None + + except NotFound: + # Session doesn't exist, can create new one + logger.debug( + f"Session {session_id} not found, can create new one" + ) + return None + except Exception as e: + logger.error(f"Error checking session {session_id}: {e}") + return None + + def _delete_session(self, session_name: str): + """Delete a session to free up the session ID for reuse.""" + try: + delete_request = DeleteSessionRequest(name=session_name) + self.session_controller_client.delete_session(delete_request) + logger.debug(f"Deleted session: {session_name}") + except NotFound: + logger.debug(f"Session already deleted: {session_name}") + + def _wait_for_termination(self, session_name: str, timeout: int = 180): + """Wait for a session to finish terminating.""" + start_time = time.time() + + while time.time() - start_time < timeout: + try: + get_request = GetSessionRequest(name=session_name) + session = self.session_controller_client.get_session( + get_request + ) + + if session.state in [ + Session.State.TERMINATED, + Session.State.FAILED, + ]: + return + elif session.state != Session.State.TERMINATING: + # Session is in unexpected state + logger.warning( + f"Session {session_name} in unexpected state while waiting for termination: {session.state}" + ) + return + + time.sleep(2) + except NotFound: + # Session was deleted + return + + logger.warning( + f"Timeout waiting for session {session_name} to terminate" + ) + + @staticmethod + def generate_session_id(): + timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + suffix_length = 6 + random_suffix = "".join( + random.choices( + string.ascii_lowercase + string.digits, k=suffix_length + ) + ) + return f"sc-{timestamp}-{random_suffix}" + + def __init__( + self, + connection: Union[str, ManagedSparkChannelBuilder], + user_id: Optional[str] = None, + ): + """ + Creates a new ManagedSparkSession for the Spark Connect interface. + + Parameters + ---------- + connection : str or :class:`ManagedSparkChannelBuilder` + Connection string that is used to extract the connection parameters + and configure the GRPC connection. Or instance of ChannelBuilder / + ManagedSparkChannelBuilder that creates GRPC connection. + user_id : str, optional + If not set, will default to the $USER environment. Defining the user + ID as part of the connection string takes precedence. + """ + + super().__init__(connection, user_id) + + execute_plan_request_base_method = ( + self.client._execute_plan_request_with_metadata + ) + execute_base_method = self.client._execute + execute_and_fetch_as_iterator_base_method = ( + self.client._execute_and_fetch_as_iterator + ) + + def execute_plan_request_wrapped_method(*args, **kwargs): + req = execute_plan_request_base_method(*args, **kwargs) + if not req.operation_id: + req.operation_id = str(uuid.uuid4()) + logger.debug( + f"No operation_id found. Setting operation_id: {req.operation_id}" + ) + return req + + self.client._execute_plan_request_with_metadata = ( + execute_plan_request_wrapped_method + ) + + 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) + + self.client._execute = MethodType(execute_wrapped_method, self.client) + + 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 + ) + + # Patching clearProgressHandlers method to not remove Managed Spark Progress Handler + clearProgressHandlers_base_method = self.clearProgressHandlers + + def clearProgressHandlers_wrapper_method(_, *args, **kwargs): + clearProgressHandlers_base_method(*args, **kwargs) + + self._register_progress_execution_handler() + + self.clearProgressHandlers = MethodType( + clearProgressHandlers_wrapper_method, self + ) + + @staticmethod + @functools.lru_cache(maxsize=1) + def get_tqdm_bar(): + """ + Return a tqdm implementation that works in the current environment. + + - Uses CLI tqdm for interactive terminals. + - Uses the notebook tqdm if available, otherwise falls back to CLI tqdm. + """ + from tqdm import tqdm as cli_tqdm + + if environment.is_interactive_terminal(): + return cli_tqdm + + try: + import ipywidgets + from tqdm.notebook import tqdm as notebook_tqdm + + return notebook_tqdm + except ImportError: + return cli_tqdm + + def _register_progress_execution_handler(self): + from pyspark.sql.connect.shell.progress import StageInfo + + def handler( + stages: Optional[Iterable[StageInfo]], + inflight_tasks: int, + operation_id: Optional[str], + done: bool, + ): + if operation_id is None: + return + + # Don't build / render progress bar for non-interactive (despite + # Ipython or non-IPython) + if not environment.is_interactive(): + return + + total_tasks = 0 + completed_tasks = 0 + + for stage in stages or []: + total_tasks += stage.num_tasks + completed_tasks += stage.num_completed_tasks + + # Don't show progress bar till we receive some tasks + if total_tasks == 0: + return + + # Get correct tqdm (notebook or CLI) + tqdm_pbar = self.get_tqdm_bar() + + # Use a lock to ensure only one thread can access and modify + # the shared dictionaries at a time. + with self._lock: + if operation_id in self._execution_progress_bar: + pbar = self._execution_progress_bar[operation_id] + if pbar.total != total_tasks: + pbar.reset( + total=total_tasks + ) # This force resets the progress bar % too on next refresh + else: + pbar = tqdm_pbar( + total=total_tasks, + leave=True, + dynamic_ncols=True, + bar_format="{l_bar}{bar} {n_fmt}/{total_fmt} Tasks", + ) + self._execution_progress_bar[operation_id] = pbar + + # To handle skipped or failed tasks. + # StageInfo proto doesn't have skipped and failed tasks information to process. + if done and completed_tasks < total_tasks: + completed_tasks = total_tasks + + pbar.n = completed_tasks + pbar.refresh() + + if done: + pbar.close() + self._execution_progress_bar.pop(operation_id, None) + + self.registerProgressHandler(handler) + + @staticmethod + def _sql_lazy_transformation(req): + # Select SQL command + try: + query = req.plan.command.sql_command.input.sql.query + return "select" in query.strip().lower().split() + except AttributeError: + return False + + def _repr_html_(self) -> str: + if not self._active_s8s_session_id: + return """ +
No Active Managed Spark Session
+ """ + + s8s_session = f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/{self._active_s8s_session_id}" + ui = f"{s8s_session}/sparkApplications/applications" + return f""" +
+

Spark Connect

+ +

Managed Spark Session

+

Spark UI

+
+ """ + + 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(): + return + + assert 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, + ] + ) + + url = ( + f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/{self._region}/" + f"{self._active_s8s_session_id}/sparkApplications/application;" + f"associatedSqlOperationId={operation_id}?project={self._project_id}" + ) + + if environment.is_interactive_terminal(): + print(f"Spark Query: {url}") + return + + try: + from IPython.display import display, HTML + + html_element = f""" +
+

Spark Query (Operation: {operation_id})

+
+ """ + display(HTML(html_element)) + except ImportError: + return + + @staticmethod + def _remove_stopped_session_from_file(): + file_path = ManagedSparkSession._get_active_session_file_path() + if file_path is not None: + try: + with open(file_path, "w"): + pass + except Exception as e: + logger.error( + f"Exception while removing active session in file {file_path}, {e}" + ) + + def addArtifacts( + self, + *artifact: str, + pyfile: bool = False, + archive: bool = False, + file: bool = False, + pypi: bool = False, + ) -> None: + """ + Add artifact(s) to the client session. Currently only local files & pypi installations are supported. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + *artifact : tuple of str + Artifact's URIs to add. + pyfile : bool + Whether to add them as Python dependencies such as .py, .egg, .zip or .jar files. + The pyfiles are directly inserted into the path when executing Python functions + in executors. + archive : bool + Whether to add them as archives such as .zip, .jar, .tar.gz, .tgz, or .tar files. + The archives are unpacked on the executor side automatically. + file : bool + Add a file to be downloaded with this Spark job on every node. + The ``path`` passed can only be a local file for now. + pypi : bool + This option is only available with ManagedSparkSession. e.g. `spark.addArtifacts("spacy==3.8.4", "torch", pypi=True)` + Installs PyPi package (with its dependencies) in the active Spark session on the driver and executors. + + Notes + ----- + This is an API dedicated to Spark Connect client only. With regular Spark Session, it throws + an exception. + Regarding pypi: Popular packages are already pre-installed in s8s runtime. + https://cloud.google.com/dataproc-serverless/docs/concepts/versions/spark-runtime-2.3#python_libraries + If there are conflicts/package doesn't exist, it throws an exception. + """ + if sum([pypi, file, pyfile, archive]) > 1: + raise ValueError( + "'pyfile', 'archive', 'file' and/or 'pypi' cannot be True together." + ) + if pypi: + artifacts = PyPiArtifacts(set(artifact)) + logger.debug("Making addArtifact call to install packages") + self.addArtifact( + artifacts.write_packages_config(self._active_s8s_session_uuid), + file=True, + ) + else: + super().addArtifacts( + *artifact, pyfile=pyfile, archive=archive, file=file + ) + + @staticmethod + def _get_active_session_file_path(): + return _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH", + "DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH", + ) + + def stop(self, terminate: Optional[bool] = None) -> None: + """ + Stop the Spark session and optionally terminate the server-side session. + + Parameters + ---------- + terminate : bool, optional + Control server-side termination behavior. + + - None (default): Auto-detect based on session type + + - Managed sessions (auto-generated ID): terminate server + - Named sessions (custom ID): client-side cleanup only + + - True: Always terminate the server-side session + - False: Never terminate the server-side session (client cleanup only) + + Examples + -------- + Auto-detect termination behavior (existing behavior): + + >>> spark.stop() + + Force terminate a named session: + + >>> spark.stop(terminate=True) + + Prevent termination of a managed session: + + >>> spark.stop(terminate=False) + """ + 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 + + 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, + ) + else: + # Client-side cleanup only + logger.debug( + f"Stopping session {ManagedSparkSession._active_s8s_session_id} without termination" + ) + + self._remove_stopped_session_from_file() + + # Clean up SparkSession._instantiatedSession if it points to this session + try: + from pyspark.sql import SparkSession as PySparkSQLSession + + if PySparkSQLSession._instantiatedSession is self: + PySparkSQLSession._instantiatedSession = None + logger.debug( + "Cleared SparkSession._instantiatedSession reference" + ) + except (ImportError, AttributeError): + # PySpark not available or _instantiatedSession doesn't exist + pass + + 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 + + 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 + + +def terminate_s8s_session( + project_id, region, active_s8s_session_id, client_options=None +): + from google.cloud.dataproc_v1 import SessionControllerClient + + logger.debug(f"Terminating Managed Spark Session: {active_s8s_session_id}") + terminate_session_request = TerminateSessionRequest() + session_name = f"projects/{project_id}/locations/{region}/sessions/{active_s8s_session_id}" + terminate_session_request.name = session_name + state = None + try: + session_client = SessionControllerClient(client_options=client_options) + session_client.terminate_session(terminate_session_request) + get_session_request = GetSessionRequest() + get_session_request.name = session_name + state = Session.State.ACTIVE + while ( + state != Session.State.TERMINATING + and state != Session.State.TERMINATED + and state != Session.State.FAILED + ): + session = session_client.get_session(get_session_request) + state = session.state + time.sleep(1) + except NotFound: + logger.debug( + f"{active_s8s_session_id} Managed Spark Session already deleted" + ) + # Client will get 'Aborted' error if session creation is still in progress and + # 'FailedPrecondition' if another termination is still in progress. + # Both are retryable, but we catch it and let TTL take care of cleanups. + except (FailedPrecondition, Aborted): + logger.debug( + f"{active_s8s_session_id} Managed Spark Session already terminated manually or automatically due to TTL" + ) + if state is not None and state == Session.State.FAILED: + raise RuntimeError("Managed Spark Session termination failed") + + +def get_active_s8s_session_response( + session_name, client_options +) -> Optional[sessions.Session]: + get_session_request = GetSessionRequest() + get_session_request.name = session_name + try: + get_session_response = SessionControllerClient( + client_options=client_options + ).get_session(get_session_request) + state = get_session_response.state + except Exception as e: + print(f"{session_name} Managed Spark Session deleted: {e}") + return None + if state is not None and ( + state == Session.State.ACTIVE or state == Session.State.CREATING + ): + return get_session_response + return None + + +def is_s8s_session_active(session_name, client_options) -> bool: + if get_active_s8s_session_response(session_name, client_options) is None: + return False + return True diff --git a/google/cloud/managed_spark_magics/__init__.py b/google/cloud/managed_spark_magics/__init__.py new file mode 100644 index 00000000..79632f57 --- /dev/null +++ b/google/cloud/managed_spark_magics/__init__.py @@ -0,0 +1,19 @@ +# 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 .magics import ManagedSparkMagics + + +def load_ipython_extension(ipython): + ipython.register_magics(ManagedSparkMagics) diff --git a/google/cloud/managed_spark_magics/magics.py b/google/cloud/managed_spark_magics/magics.py new file mode 100644 index 00000000..54363ae1 --- /dev/null +++ b/google/cloud/managed_spark_magics/magics.py @@ -0,0 +1,76 @@ +# 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. + +"""Managed Spark magic implementations.""" + +import shlex +from IPython.core.magic import (Magics, magics_class, line_magic) +from google.cloud.managed_spark_connect import ManagedSparkSession + + +@magics_class +class ManagedSparkMagics(Magics): + + def __init__( + self, + shell, + **kwargs, + ): + super().__init__(shell, **kwargs) + + @line_magic + def dpip(self, line): + """ + Custom magic to install pip packages as Spark Connect artifacts. + Usage: %dpip install pandas numpy + """ + try: + args = shlex.split(line) + + if not args or args[0] != "install": + raise RuntimeError( + "Usage: %dpip install ..." + ) + + packages = args[1:] # remove `install` + + if not packages: + raise RuntimeError("Error: No packages specified.") + + if any(pkg.startswith("-") for pkg in packages): + raise RuntimeError("Error: Flags are not currently supported.") + + sessions = [ + (key, value) + for key, value in self.shell.user_ns.items() + if isinstance(value, ManagedSparkSession) + ] + + if not sessions: + raise RuntimeError( + "Error: No active Managed Spark Session found. Please create one first." + ) + if len(sessions) > 1: + raise RuntimeError( + "Error: Found more than one active Managed Spark Sessions." + ) + + ((name, session),) = sessions + print(f"Active session found: {name}") + print(f"Installing packages: {packages}") + session.addArtifacts(*packages, pypi=True) + + print("Finished installing packages.") + except Exception as e: + raise RuntimeError(f"Failed to install packages: {e}") from e diff --git a/setup.py b/setup.py index ed906106..ca9d32d8 100644 --- a/setup.py +++ b/setup.py @@ -19,13 +19,13 @@ setup( - name="dataproc-spark-connect", + name="managed-spark-connect", version="1.1.0", - description="Dataproc client library for Spark Connect", + description="Managed Spark client library for Spark Connect", long_description=long_description, long_description_content_type="text/markdown", author="Google LLC", - url="https://github.com/GoogleCloudDataproc/dataproc-spark-connect-python", + url="https://github.com/GoogleCloudDataproc/managed-spark-connect-python", license="Apache 2.0", packages=find_namespace_packages(include=["google.*"]), install_requires=[ diff --git a/tests/integration/dataproc_magics/__init__.py b/tests/integration/managed_spark_magics/__init__.py similarity index 100% rename from tests/integration/dataproc_magics/__init__.py rename to tests/integration/managed_spark_magics/__init__.py diff --git a/tests/integration/dataproc_magics/test_magics.py b/tests/integration/managed_spark_magics/test_magics.py similarity index 89% rename from tests/integration/dataproc_magics/test_magics.py rename to tests/integration/managed_spark_magics/test_magics.py index 67a09764..2b067178 100644 --- a/tests/integration/dataproc_magics/test_magics.py +++ b/tests/integration/managed_spark_magics/test_magics.py @@ -16,7 +16,7 @@ import certifi from unittest import mock -from google.cloud.dataproc_spark_connect import DataprocSparkSession +from google.cloud.managed_spark_connect import ManagedSparkSession _SERVICE_ACCOUNT_KEY_FILE_ = "service_account_key.json" @@ -63,12 +63,12 @@ def auth_type(request): @pytest.fixture def test_subnet(): - return os.getenv("DATAPROC_SPARK_CONNECT_SUBNET") + return os.getenv("MANAGED_SPARK_CONNECT_SUBNET") @pytest.fixture def test_subnetwork_uri(test_subnet): - # Make DATAPROC_SPARK_CONNECT_SUBNET the full URI + # Make MANAGED_SPARK_CONNECT_SUBNET the full URI # to align with how user would specify it in the project return test_subnet @@ -80,9 +80,9 @@ def os_environment(auth_type, image_version, test_project, test_region): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = ( _SERVICE_ACCOUNT_KEY_FILE_ ) - os.environ["DATAPROC_SPARK_CONNECT_AUTH_TYPE"] = auth_type + os.environ["MANAGED_SPARK_CONNECT_AUTH_TYPE"] = auth_type if auth_type == "END_USER_CREDENTIALS": - os.environ.pop("DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT", None) + os.environ.pop("MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT", None) # Add SSL certificate fix os.environ["SSL_CERT_FILE"] = certifi.where() os.environ["REQUESTS_CA_BUNDLE"] = certifi.where() @@ -94,7 +94,7 @@ def os_environment(auth_type, image_version, test_project, test_region): @pytest.fixture def connect_session(test_project, test_region, os_environment): session = ( - DataprocSparkSession.builder.projectId(test_project) + ManagedSparkSession.builder.projectId(test_project) .location(test_region) .getOrCreate() ) @@ -109,16 +109,16 @@ def connect_session(test_project, test_region, os_environment): @pytest.fixture def ipython_shell(connect_session): - """Provides an IPython shell with a DataprocSparkSession in user_ns.""" + """Provides an IPython shell with a ManagedSparkSession in user_ns.""" try: from IPython.terminal.interactiveshell import TerminalInteractiveShell - from google.cloud import dataproc_magics + from google.cloud import managed_spark_magics shell = TerminalInteractiveShell.instance() shell.user_ns = {"spark": connect_session} # Load magics - dataproc_magics.load_ipython_extension(shell) + managed_spark_magics.load_ipython_extension(shell) yield shell finally: @@ -186,7 +186,7 @@ def test_dpip_no_session(ipython_shell): """Test message when no Spark session is active.""" ipython_shell.user_ns = {} # Remove spark session from namespace with pytest.raises( - RuntimeError, match="No active Dataproc Spark Session found." + RuntimeError, match="No active Managed Spark Session found." ): ipython_shell.run_line_magic("dpip", "install pandas") @@ -206,6 +206,6 @@ def test_dpip_multiple_sessions(ipython_shell, connect_session): ipython_shell.user_ns["sparkanother"] = connect_session with pytest.raises( RuntimeError, - match="Error: Found more than one active Dataproc Spark Sessions.", + match="Error: Found more than one active Managed Spark Sessions.", ): ipython_shell.run_line_magic("dpip", "install pandas") diff --git a/tests/integration/test_session.py b/tests/integration/test_session.py index d39292a0..61d4a2f6 100644 --- a/tests/integration/test_session.py +++ b/tests/integration/test_session.py @@ -18,7 +18,7 @@ import certifi from google.api_core import client_options -from google.cloud.dataproc_spark_connect import DataprocSparkSession +from google.cloud.managed_spark_connect import ManagedSparkSession from google.cloud.dataproc_v1 import ( CreateSessionTemplateRequest, DeleteSessionRequest, @@ -79,12 +79,12 @@ def test_region(): @pytest.fixture def test_subnet(): - return os.getenv("DATAPROC_SPARK_CONNECT_SUBNET") + return os.getenv("MANAGED_SPARK_CONNECT_SUBNET") @pytest.fixture def test_subnetwork_uri(test_subnet): - # Make DATAPROC_SPARK_CONNECT_SUBNET the full URI to align with how user would specify it in the project + # Make MANAGED_SPARK_CONNECT_SUBNET the full URI to align with how user would specify it in the project return test_subnet @@ -95,9 +95,9 @@ def os_environment(auth_type, image_version, test_project, test_region): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = ( _SERVICE_ACCOUNT_KEY_FILE_ ) - os.environ["DATAPROC_SPARK_CONNECT_AUTH_TYPE"] = auth_type + os.environ["MANAGED_SPARK_CONNECT_AUTH_TYPE"] = auth_type if auth_type == "END_USER_CREDENTIALS": - os.environ.pop("DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT", None) + os.environ.pop("MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT", None) # Add SSL certificate fix os.environ["SSL_CERT_FILE"] = certifi.where() os.environ["REQUESTS_CA_BUNDLE"] = certifi.where() @@ -132,7 +132,7 @@ def session_template_controller_client(test_client_options): @pytest.fixture def connect_session(test_project, test_region, os_environment): session = ( - DataprocSparkSession.builder.projectId(test_project) + ManagedSparkSession.builder.projectId(test_project) .location(test_region) .getOrCreate() ) @@ -147,7 +147,7 @@ def connect_session(test_project, test_region, os_environment): @pytest.fixture def session_name(test_project, test_region, connect_session): - return f"projects/{test_project}/locations/{test_region}/sessions/{DataprocSparkSession._active_s8s_session_id}" + return f"projects/{test_project}/locations/{test_region}/sessions/{ManagedSparkSession._active_s8s_session_id}" def test_create_spark_session_with_default_notebook_behavior( @@ -172,7 +172,7 @@ def test_create_spark_session_with_default_notebook_behavior( assert "[TABLE_OR_VIEW_ALREADY_EXISTS]" in str(ex) - assert DataprocSparkSession._active_s8s_session_uuid is not None + assert ManagedSparkSession._active_s8s_session_uuid is not None connect_session.sql("DROP TABLE IF EXISTS FOO") connect_session.stop() session = session_controller_client.get_session(get_session_request) @@ -181,26 +181,26 @@ def test_create_spark_session_with_default_notebook_behavior( Session.State.TERMINATING, Session.State.TERMINATED, ] - assert DataprocSparkSession._active_s8s_session_uuid is None + assert ManagedSparkSession._active_s8s_session_uuid is None def test_reuse_s8s_spark_session( connect_session, session_name, session_controller_client ): """Test that Spark sessions can be reused within the same process.""" - assert DataprocSparkSession._active_s8s_session_uuid is not None + assert ManagedSparkSession._active_s8s_session_uuid is not None - first_session_id = DataprocSparkSession._active_s8s_session_id - first_session_uuid = DataprocSparkSession._active_s8s_session_uuid + first_session_id = ManagedSparkSession._active_s8s_session_id + first_session_uuid = ManagedSparkSession._active_s8s_session_uuid - connect_session = DataprocSparkSession.builder.getOrCreate() - second_session_id = DataprocSparkSession._active_s8s_session_id - second_session_uuid = DataprocSparkSession._active_s8s_session_uuid + connect_session = ManagedSparkSession.builder.getOrCreate() + second_session_id = ManagedSparkSession._active_s8s_session_id + second_session_uuid = ManagedSparkSession._active_s8s_session_uuid assert first_session_id == second_session_id assert first_session_uuid == second_session_uuid - assert DataprocSparkSession._active_s8s_session_uuid is not None - assert DataprocSparkSession._active_s8s_session_id is not None + assert ManagedSparkSession._active_s8s_session_uuid is not None + assert ManagedSparkSession._active_s8s_session_id is not None connect_session.stop() @@ -209,7 +209,7 @@ def test_stop_spark_session_with_deleted_serverless_session( connect_session, session_name, session_controller_client ): """Test stopping a Spark session when the serverless session has been deleted.""" - assert DataprocSparkSession._active_s8s_session_uuid is not None + assert ManagedSparkSession._active_s8s_session_uuid is not None delete_session_request = DeleteSessionRequest() delete_session_request.name = session_name @@ -217,15 +217,15 @@ def test_stop_spark_session_with_deleted_serverless_session( operation.result() connect_session.stop() - assert DataprocSparkSession._active_s8s_session_uuid is None - assert DataprocSparkSession._active_s8s_session_id is None + assert ManagedSparkSession._active_s8s_session_uuid is None + assert ManagedSparkSession._active_s8s_session_id is None def test_stop_spark_session_with_terminated_serverless_session( connect_session, session_name, session_controller_client ): """Test stopping a Spark session when the serverless session has been terminated.""" - assert DataprocSparkSession._active_s8s_session_uuid is not None + assert ManagedSparkSession._active_s8s_session_uuid is not None terminate_session_request = TerminateSessionRequest() terminate_session_request.name = session_name @@ -235,8 +235,8 @@ def test_stop_spark_session_with_terminated_serverless_session( operation.result() connect_session.stop() - assert DataprocSparkSession._active_s8s_session_uuid is None - assert DataprocSparkSession._active_s8s_session_id is None + assert ManagedSparkSession._active_s8s_session_uuid is None + assert ManagedSparkSession._active_s8s_session_id is None def test_get_or_create_spark_session_with_terminated_serverless_session( @@ -249,22 +249,22 @@ def test_get_or_create_spark_session_with_terminated_serverless_session( """Test creating a new Spark session when the previous serverless session has been terminated.""" first_session_name = session_name - assert DataprocSparkSession._active_s8s_session_uuid is not None + assert ManagedSparkSession._active_s8s_session_uuid is not None - first_session = DataprocSparkSession._active_s8s_session_uuid + first_session = ManagedSparkSession._active_s8s_session_uuid terminate_session_request = TerminateSessionRequest() terminate_session_request.name = first_session_name operation = session_controller_client.terminate_session( terminate_session_request ) operation.result() - connect_session = DataprocSparkSession.builder.getOrCreate() - second_session = DataprocSparkSession._active_s8s_session_uuid - second_session_name = f"projects/{test_project}/locations/{test_region}/sessions/{DataprocSparkSession._active_s8s_session_id}" + connect_session = ManagedSparkSession.builder.getOrCreate() + second_session = ManagedSparkSession._active_s8s_session_uuid + second_session_name = f"projects/{test_project}/locations/{test_region}/sessions/{ManagedSparkSession._active_s8s_session_id}" assert first_session != second_session - assert DataprocSparkSession._active_s8s_session_uuid is not None - assert DataprocSparkSession._active_s8s_session_id is not None + assert ManagedSparkSession._active_s8s_session_uuid is not None + assert ManagedSparkSession._active_s8s_session_id is not None get_session_request = GetSessionRequest() get_session_request.name = first_session_name @@ -315,7 +315,7 @@ def session_template_name( assert ( session_template.runtime_config.version == image_version if image_version - else DataprocSparkSession._DEFAULT_RUNTIME_VERSION + else ManagedSparkSession._DEFAULT_RUNTIME_VERSION ) yield session_template.name @@ -333,17 +333,17 @@ def test_create_spark_session_with_session_template_and_user_provided_dataproc_c session_template_name, session_controller_client, ): - """Test creating a Spark session with a session template and user-provided Dataproc configuration.""" + """Test creating a Spark session with a Runtime Profile and user-provided Dataproc configuration.""" dataproc_config = Session() dataproc_config.environment_config.execution_config.ttl = {"seconds": 64800} dataproc_config.session_template = session_template_name connect_session = ( - DataprocSparkSession.builder.config("spark.executor.cores", "7") + ManagedSparkSession.builder.config("spark.executor.cores", "7") .dataprocSessionConfig(dataproc_config) .config("spark.executor.cores", "16") .getOrCreate() ) - session_name = f"projects/{test_project}/locations/{test_region}/sessions/{DataprocSparkSession._active_s8s_session_id}" + session_name = f"projects/{test_project}/locations/{test_region}/sessions/{ManagedSparkSession._active_s8s_session_id}" get_session_request = GetSessionRequest() get_session_request.name = session_name @@ -358,7 +358,7 @@ def test_create_spark_session_with_session_template_and_user_provided_dataproc_c assert ( session.runtime_config.properties["spark:spark.executor.cores"] == "16" ) - assert DataprocSparkSession._active_s8s_session_uuid is not None + assert ManagedSparkSession._active_s8s_session_uuid is not None connect_session.stop() get_session_request = GetSessionRequest() @@ -369,7 +369,7 @@ def test_create_spark_session_with_session_template_and_user_provided_dataproc_c Session.State.TERMINATING, Session.State.TERMINATED, ] - assert DataprocSparkSession._active_s8s_session_uuid is None + assert ManagedSparkSession._active_s8s_session_uuid is None @pytest.mark.skip( @@ -380,7 +380,7 @@ def test_add_artifacts_pypi_package(): Note: Skipped in CI due to infrastructure issues with PyPI package installation. """ - connect_session = DataprocSparkSession.builder.getOrCreate() + connect_session = ManagedSparkSession.builder.getOrCreate() from pyspark.sql.connect.functions import udf, sum from pyspark.sql.types import IntegerType @@ -489,9 +489,9 @@ def test_session_reuse_with_custom_id( custom_session_id = f"ml-pipeline-session-{uuid.uuid4().hex[:8]}" # Stop any existing session first to ensure clean state - if DataprocSparkSession._active_s8s_session_id: + if ManagedSparkSession._active_s8s_session_id: try: - existing_session = DataprocSparkSession.getActiveSession() + existing_session = ManagedSparkSession.getActiveSession() if existing_session: existing_session.stop() except Exception: @@ -499,14 +499,14 @@ def test_session_reuse_with_custom_id( # PHASE 1: Create initial session with custom ID spark1 = ( - DataprocSparkSession.builder.dataprocSessionId(custom_session_id) + ManagedSparkSession.builder.dataprocSessionId(custom_session_id) .projectId(test_project) .location(test_region) .getOrCreate() ) # Verify session is created with custom ID - assert DataprocSparkSession._active_s8s_session_id == custom_session_id + assert ManagedSparkSession._active_s8s_session_id == custom_session_id first_session_uuid = spark1._active_s8s_session_uuid # Test basic functionality @@ -516,17 +516,17 @@ def test_session_reuse_with_custom_id( # PHASE 2: Test session reuse while active # Clear cache to force session lookup - DataprocSparkSession._default_session = None + ManagedSparkSession._default_session = None spark2 = ( - DataprocSparkSession.builder.dataprocSessionId(custom_session_id) + ManagedSparkSession.builder.dataprocSessionId(custom_session_id) .projectId(test_project) .location(test_region) .getOrCreate() ) # Should reuse the same active session - assert DataprocSparkSession._active_s8s_session_id == custom_session_id + assert ManagedSparkSession._active_s8s_session_id == custom_session_id assert spark2._active_s8s_session_uuid == first_session_uuid # Test functionality on reused session @@ -539,19 +539,19 @@ def test_session_reuse_with_custom_id( # PHASE 4: Recreate with same ID - this tests the cleanup and recreation logic # Clear all session state to ensure fresh lookup - DataprocSparkSession._default_session = None - DataprocSparkSession._active_s8s_session_id = None - DataprocSparkSession._active_s8s_session_uuid = None + ManagedSparkSession._default_session = None + ManagedSparkSession._active_s8s_session_id = None + ManagedSparkSession._active_s8s_session_uuid = None spark3 = ( - DataprocSparkSession.builder.dataprocSessionId(custom_session_id) + ManagedSparkSession.builder.dataprocSessionId(custom_session_id) .projectId(test_project) .location(test_region) .getOrCreate() ) # Should be a same session and same ID - assert DataprocSparkSession._active_s8s_session_id == custom_session_id + assert ManagedSparkSession._active_s8s_session_id == custom_session_id third_session_uuid = spark3._active_s8s_session_uuid # Should be same UUID @@ -573,13 +573,13 @@ def test_session_id_validation_in_integration( # Test invalid session ID raises ValueError with pytest.raises(ValueError) as exc_info: - DataprocSparkSession.builder.dataprocSessionId("123-invalid-id") + ManagedSparkSession.builder.dataprocSessionId("123-invalid-id") assert "Invalid session ID" in str(exc_info.value) # Test that valid session ID works valid_id = "valid-session-id-123" builder = ( - DataprocSparkSession.builder.dataprocSessionId(valid_id) + ManagedSparkSession.builder.dataprocSessionId(valid_id) .projectId(test_project) .location(test_region) ) @@ -614,15 +614,15 @@ def test_sparksql_magic_library_available(connect_session): assert magic_loaded, "sparksql_magic should be available as a dependency" - # Test that DataprocSparkSession can execute SQL (ensuring basic compatibility) + # Test that ManagedSparkSession can execute SQL (ensuring basic compatibility) result = connect_session.sql("SELECT 'integration_test' as test_column") data = result.collect() assert len(data) == 1 assert data[0]["test_column"] == "integration_test" -def test_sparksql_magic_with_dataproc_session(connect_session): - """Test that sparksql-magic works with registered DataprocSparkSession.""" +def test_sparksql_magic_with_managed_spark_session(connect_session): + """Test that sparksql-magic works with registered ManagedSparkSession.""" pytest.importorskip( "IPython", reason="IPython not available (install with magic extra)" ) @@ -633,7 +633,7 @@ def test_sparksql_magic_with_dataproc_session(connect_session): from IPython.terminal.interactiveshell import TerminalInteractiveShell - # Create real IPython shell (DataprocSparkSession is already registered globally) + # Create real IPython shell (ManagedSparkSession is already registered globally) shell = TerminalInteractiveShell.instance() # Load the sparksql_magic extension @@ -679,14 +679,14 @@ def test_stop_named_session_with_terminate_true( # Create a session with custom ID spark = ( - DataprocSparkSession.builder.dataprocSessionId(custom_session_id) + ManagedSparkSession.builder.dataprocSessionId(custom_session_id) .projectId(test_project) .location(test_region) .getOrCreate() ) # Verify session is created - assert DataprocSparkSession._active_s8s_session_id == custom_session_id + assert ManagedSparkSession._active_s8s_session_id == custom_session_id session_name = f"projects/{test_project}/locations/{test_region}/sessions/{custom_session_id}" # Test basic functionality @@ -697,7 +697,7 @@ def test_stop_named_session_with_terminate_true( spark.stop(terminate=True) # Verify client-side cleanup - assert DataprocSparkSession._active_s8s_session_id is None + assert ManagedSparkSession._active_s8s_session_id is None # Verify server-side session is terminating or terminated get_session_request = GetSessionRequest() @@ -720,15 +720,15 @@ def test_stop_managed_session_with_terminate_false( """Test that stop(terminate=False) does NOT terminate a managed session on the server.""" # Create a managed session (auto-generated ID) spark = ( - DataprocSparkSession.builder.projectId(test_project) + ManagedSparkSession.builder.projectId(test_project) .location(test_region) .getOrCreate() ) # Verify it's a managed session (auto-generated ID) - assert DataprocSparkSession._active_s8s_session_id is not None - assert DataprocSparkSession._active_session_uses_custom_id is False - session_id = DataprocSparkSession._active_s8s_session_id + assert ManagedSparkSession._active_s8s_session_id is not None + assert ManagedSparkSession._active_session_uses_custom_id is False + session_id = ManagedSparkSession._active_s8s_session_id session_name = ( f"projects/{test_project}/locations/{test_region}/sessions/{session_id}" ) @@ -741,7 +741,7 @@ def test_stop_managed_session_with_terminate_false( spark.stop(terminate=False) # Verify client-side cleanup - assert DataprocSparkSession._active_s8s_session_id is None + assert ManagedSparkSession._active_s8s_session_id is None # Verify server-side session is still ACTIVE (not terminated) get_session_request = GetSessionRequest() @@ -768,10 +768,10 @@ def local_spark_session(): from pyspark.sql import SparkSession as PySparkSession # Stop any existing session to ensure a clean environment for creating a local session. - # This prevents test isolation failures where a Dataproc session from a previous + # This prevents test isolation failures where a Managed Spark session from a previous # test might be picked up by getOrCreate(). - if DataprocSparkSession.getActiveSession(): - DataprocSparkSession.getActiveSession().stop() + if ManagedSparkSession.getActiveSession(): + ManagedSparkSession.getActiveSession().stop() session = PySparkSession.builder.master("local").getOrCreate() yield session @@ -782,12 +782,12 @@ def test_create_local_spark_session(batch_workload_env, local_spark_session): """Test creating a local Spark session.""" from pyspark.sql import SparkSession as PySparkSession - dataproc_spark_session = DataprocSparkSession.builder.getOrCreate() + managed_spark_session = ManagedSparkSession.builder.getOrCreate() try: - assert isinstance(dataproc_spark_session, PySparkSession) - assert not isinstance(dataproc_spark_session, DataprocSparkSession) + assert isinstance(managed_spark_session, PySparkSession) + assert not isinstance(managed_spark_session, ManagedSparkSession) # Compare configurations to ensure they are both local sessions - assert dataproc_spark_session == local_spark_session + assert managed_spark_session == local_spark_session finally: - dataproc_spark_session.stop() + managed_spark_session.stop() diff --git a/tests/unit/dataproc_magics/__init__.py b/tests/unit/managed_spark_magics/__init__.py similarity index 100% rename from tests/unit/dataproc_magics/__init__.py rename to tests/unit/managed_spark_magics/__init__.py diff --git a/tests/unit/dataproc_magics/test_magics.py b/tests/unit/managed_spark_magics/test_magics.py similarity index 84% rename from tests/unit/dataproc_magics/test_magics.py rename to tests/unit/managed_spark_magics/test_magics.py index 83d0b3ed..f546057c 100644 --- a/tests/unit/dataproc_magics/test_magics.py +++ b/tests/unit/managed_spark_magics/test_magics.py @@ -17,19 +17,19 @@ from contextlib import redirect_stdout from unittest import mock -from google.cloud.dataproc_spark_connect import DataprocSparkSession -from google.cloud.dataproc_magics import DataprocMagics +from google.cloud.managed_spark_connect import ManagedSparkSession +from google.cloud.managed_spark_magics import ManagedSparkMagics from IPython.core.interactiveshell import InteractiveShell from traitlets.config import Config -class DataprocMagicsTest(unittest.TestCase): +class ManagedSparkMagicsTest(unittest.TestCase): def setUp(self): self.shell = mock.create_autospec(InteractiveShell, instance=True) self.shell.user_ns = {} self.shell.config = Config() - self.magics = DataprocMagics(shell=self.shell) + self.magics = ManagedSparkMagics(shell=self.shell) def test_dpip_with_flags(self): with self.assertRaisesRegex( @@ -51,18 +51,18 @@ def test_dpip_invalid_command(self): def test_dpip_no_session(self): with self.assertRaisesRegex( - RuntimeError, "Error: No active Dataproc Spark Session found" + RuntimeError, "Error: No active Managed Spark Session found" ): self.magics.dpip("install pandas") def test_dpip_multiple_sessions(self): - mock_session = mock.Mock(spec=DataprocSparkSession) + mock_session = mock.Mock(spec=ManagedSparkSession) self.shell.user_ns["spark1"] = mock_session self.shell.user_ns["spark2"] = mock_session with self.assertRaisesRegex( RuntimeError, - "Error: Found more than one active Dataproc Spark Sessions", + "Error: Found more than one active Managed Spark Sessions", ): self.magics.dpip("install pandas") @@ -73,7 +73,7 @@ def test_dpip_no_packages_specified(self): self.magics.dpip("install") def test_dpip_install_packages_success(self): - mock_session = mock.Mock(spec=DataprocSparkSession) + mock_session = mock.Mock(spec=ManagedSparkSession) self.shell.user_ns["spark"] = mock_session f = io.StringIO() @@ -87,7 +87,7 @@ def test_dpip_install_packages_success(self): self.assertIn("Finished installing packages.", f.getvalue()) def test_dpip_add_artifacts_fails(self): - mock_session = mock.Mock(spec=DataprocSparkSession) + mock_session = mock.Mock(spec=ManagedSparkSession) mock_session.addArtifacts.side_effect = Exception("Failed") self.shell.user_ns["spark"] = mock_session diff --git a/tests/unit/test_deprecated_shims.py b/tests/unit/test_deprecated_shims.py new file mode 100644 index 00000000..aca687cd --- /dev/null +++ b/tests/unit/test_deprecated_shims.py @@ -0,0 +1,75 @@ +# 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. +"""Tests that the pre-rename `dataproc_*` import paths still work as deprecated aliases.""" +import importlib +import sys +import unittest + + +def _fresh_import(module_name): + """Import (or re-import) a module, forcing its top-level code to run. + + Needed because a module already cached in sys.modules from an earlier + test or import elsewhere would otherwise not re-emit its deprecation + warning, making assertWarns order-dependent. + """ + for name in list(sys.modules): + if name == module_name or name.startswith(module_name + "."): + del sys.modules[name] + return importlib.import_module(module_name) + + +class DeprecatedPackageShimTests(unittest.TestCase): + + def test_dataproc_spark_connect_package_warns_and_aliases_session(self): + from google.cloud.managed_spark_connect import ManagedSparkSession + + with self.assertWarns(DeprecationWarning): + module = _fresh_import("google.cloud.dataproc_spark_connect") + + self.assertIs(module.DataprocSparkSession, ManagedSparkSession) + + def test_dataproc_spark_connect_exceptions_alias(self): + from google.cloud.managed_spark_connect.exceptions import ( + ManagedSparkConnectException, + ) + from google.cloud.dataproc_spark_connect.exceptions import ( + DataprocSparkConnectException, + ) + + self.assertIs( + DataprocSparkConnectException, ManagedSparkConnectException + ) + + def test_dataproc_spark_connect_client_alias(self): + from google.cloud.managed_spark_connect.client import ( + ManagedSparkChannelBuilder, + ) + from google.cloud.dataproc_spark_connect.client import ( + DataprocChannelBuilder, + ) + + self.assertIs(DataprocChannelBuilder, ManagedSparkChannelBuilder) + + def test_dataproc_magics_package_warns_and_aliases_magics(self): + from google.cloud.managed_spark_magics import ManagedSparkMagics + + with self.assertWarns(DeprecationWarning): + module = _fresh_import("google.cloud.dataproc_magics") + + self.assertIs(module.DataprocMagics, ManagedSparkMagics) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_environment.py b/tests/unit/test_environment.py index a64387af..d8ef2d6f 100644 --- a/tests/unit/test_environment.py +++ b/tests/unit/test_environment.py @@ -17,7 +17,7 @@ import unittest from unittest import mock -from google.cloud.dataproc_spark_connect import environment +from google.cloud.managed_spark_connect import environment class TestEnvironment(unittest.TestCase): @@ -200,75 +200,75 @@ def test_is_jetbrains_ide_false_env_var_not_jetbrains(self): # ---- get_client_environment_label tests ---- @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + "google.cloud.managed_spark_connect.environment.is_colab_enterprise", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab", + "google.cloud.managed_spark_connect.environment.is_colab", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_workbench", + "google.cloud.managed_spark_connect.environment.is_workbench", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_kaggle", + "google.cloud.managed_spark_connect.environment.is_kaggle", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_sagemaker", + "google.cloud.managed_spark_connect.environment.is_sagemaker", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_databricks", + "google.cloud.managed_spark_connect.environment.is_databricks", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_deepnote", + "google.cloud.managed_spark_connect.environment.is_deepnote", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_datalore", + "google.cloud.managed_spark_connect.environment.is_datalore", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_codespaces", + "google.cloud.managed_spark_connect.environment.is_codespaces", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_cloud_shell", + "google.cloud.managed_spark_connect.environment.is_cloud_shell", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_hex", + "google.cloud.managed_spark_connect.environment.is_hex", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_polynote", + "google.cloud.managed_spark_connect.environment.is_polynote", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_antigravity", + "google.cloud.managed_spark_connect.environment.is_antigravity", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_vscode", + "google.cloud.managed_spark_connect.environment.is_vscode", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_jetbrains_ide", + "google.cloud.managed_spark_connect.environment.is_jetbrains_ide", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_spyder", + "google.cloud.managed_spark_connect.environment.is_spyder", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_eclipse", + "google.cloud.managed_spark_connect.environment.is_eclipse", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_jupyter", + "google.cloud.managed_spark_connect.environment.is_jupyter", return_value=False, ) def test_get_client_environment_label_unknown(self, *mocks): @@ -278,11 +278,11 @@ def test_get_client_environment_label_unknown(self, *mocks): ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + "google.cloud.managed_spark_connect.environment.is_colab_enterprise", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab", + "google.cloud.managed_spark_connect.environment.is_colab", return_value=True, ) def test_get_client_environment_label_colab(self, *mocks): @@ -292,7 +292,7 @@ def test_get_client_environment_label_colab(self, *mocks): ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + "google.cloud.managed_spark_connect.environment.is_colab_enterprise", return_value=True, ) def test_get_client_environment_label_colab_enterprise( @@ -304,15 +304,15 @@ def test_get_client_environment_label_colab_enterprise( ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + "google.cloud.managed_spark_connect.environment.is_colab_enterprise", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab", + "google.cloud.managed_spark_connect.environment.is_colab", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_workbench", + "google.cloud.managed_spark_connect.environment.is_workbench", return_value=True, ) def test_get_client_environment_label_workbench(self, *mocks): @@ -322,19 +322,19 @@ def test_get_client_environment_label_workbench(self, *mocks): ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + "google.cloud.managed_spark_connect.environment.is_colab_enterprise", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab", + "google.cloud.managed_spark_connect.environment.is_colab", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_workbench", + "google.cloud.managed_spark_connect.environment.is_workbench", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_kaggle", + "google.cloud.managed_spark_connect.environment.is_kaggle", return_value=True, ) def test_get_client_environment_label_kaggle(self, *mocks): @@ -344,55 +344,55 @@ def test_get_client_environment_label_kaggle(self, *mocks): ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + "google.cloud.managed_spark_connect.environment.is_colab_enterprise", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab", + "google.cloud.managed_spark_connect.environment.is_colab", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_workbench", + "google.cloud.managed_spark_connect.environment.is_workbench", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_kaggle", + "google.cloud.managed_spark_connect.environment.is_kaggle", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_sagemaker", + "google.cloud.managed_spark_connect.environment.is_sagemaker", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_databricks", + "google.cloud.managed_spark_connect.environment.is_databricks", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_deepnote", + "google.cloud.managed_spark_connect.environment.is_deepnote", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_datalore", + "google.cloud.managed_spark_connect.environment.is_datalore", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_codespaces", + "google.cloud.managed_spark_connect.environment.is_codespaces", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_cloud_shell", + "google.cloud.managed_spark_connect.environment.is_cloud_shell", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_hex", + "google.cloud.managed_spark_connect.environment.is_hex", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_polynote", + "google.cloud.managed_spark_connect.environment.is_polynote", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_antigravity", + "google.cloud.managed_spark_connect.environment.is_antigravity", return_value=True, ) def test_get_client_environment_label_antigravity(self, *mocks): @@ -402,59 +402,59 @@ def test_get_client_environment_label_antigravity(self, *mocks): ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + "google.cloud.managed_spark_connect.environment.is_colab_enterprise", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab", + "google.cloud.managed_spark_connect.environment.is_colab", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_workbench", + "google.cloud.managed_spark_connect.environment.is_workbench", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_kaggle", + "google.cloud.managed_spark_connect.environment.is_kaggle", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_sagemaker", + "google.cloud.managed_spark_connect.environment.is_sagemaker", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_databricks", + "google.cloud.managed_spark_connect.environment.is_databricks", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_deepnote", + "google.cloud.managed_spark_connect.environment.is_deepnote", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_datalore", + "google.cloud.managed_spark_connect.environment.is_datalore", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_codespaces", + "google.cloud.managed_spark_connect.environment.is_codespaces", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_cloud_shell", + "google.cloud.managed_spark_connect.environment.is_cloud_shell", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_hex", + "google.cloud.managed_spark_connect.environment.is_hex", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_polynote", + "google.cloud.managed_spark_connect.environment.is_polynote", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_antigravity", + "google.cloud.managed_spark_connect.environment.is_antigravity", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_vscode", + "google.cloud.managed_spark_connect.environment.is_vscode", return_value=True, ) def test_get_client_environment_label_vscode(self, *mocks): @@ -464,63 +464,63 @@ def test_get_client_environment_label_vscode(self, *mocks): ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + "google.cloud.managed_spark_connect.environment.is_colab_enterprise", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab", + "google.cloud.managed_spark_connect.environment.is_colab", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_workbench", + "google.cloud.managed_spark_connect.environment.is_workbench", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_kaggle", + "google.cloud.managed_spark_connect.environment.is_kaggle", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_sagemaker", + "google.cloud.managed_spark_connect.environment.is_sagemaker", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_databricks", + "google.cloud.managed_spark_connect.environment.is_databricks", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_deepnote", + "google.cloud.managed_spark_connect.environment.is_deepnote", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_datalore", + "google.cloud.managed_spark_connect.environment.is_datalore", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_codespaces", + "google.cloud.managed_spark_connect.environment.is_codespaces", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_cloud_shell", + "google.cloud.managed_spark_connect.environment.is_cloud_shell", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_hex", + "google.cloud.managed_spark_connect.environment.is_hex", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_polynote", + "google.cloud.managed_spark_connect.environment.is_polynote", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_antigravity", + "google.cloud.managed_spark_connect.environment.is_antigravity", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_vscode", + "google.cloud.managed_spark_connect.environment.is_vscode", return_value=False, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_jetbrains_ide", + "google.cloud.managed_spark_connect.environment.is_jetbrains_ide", return_value=True, ) def test_get_client_environment_label_jetbrains_ide(self, *mocks): @@ -530,11 +530,11 @@ def test_get_client_environment_label_jetbrains_ide(self, *mocks): ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + "google.cloud.managed_spark_connect.environment.is_colab_enterprise", return_value=True, ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.is_colab", + "google.cloud.managed_spark_connect.environment.is_colab", return_value=True, ) def test_get_client_environment_label_precedence( @@ -550,7 +550,7 @@ def test_is_interactive_ipython_true(self, mock_get_ipython): self.assertTrue(environment.is_interactive()) @mock.patch("IPython.get_ipython", return_value=None) - @mock.patch("google.cloud.dataproc_spark_connect.environment.sys") + @mock.patch("google.cloud.managed_spark_connect.environment.sys") def test_is_interactive_ipython_false(self, mock_sys, mock_get_ipython): if hasattr(mock_sys, "ps1"): del mock_sys.ps1 @@ -558,7 +558,7 @@ def test_is_interactive_ipython_false(self, mock_sys, mock_get_ipython): self.assertFalse(environment.is_interactive()) @mock.patch("IPython.get_ipython", side_effect=ImportError) - @mock.patch("google.cloud.dataproc_spark_connect.environment.sys") + @mock.patch("google.cloud.managed_spark_connect.environment.sys") def test_is_interactive_true_via_ps1(self, mock_sys, mock_get_ipython): # Simulate interactive environment by setting ps1 mock_sys.ps1 = ">>>" @@ -566,7 +566,7 @@ def test_is_interactive_true_via_ps1(self, mock_sys, mock_get_ipython): self.assertTrue(environment.is_interactive()) @mock.patch("IPython.get_ipython", side_effect=ImportError) - @mock.patch("google.cloud.dataproc_spark_connect.environment.sys") + @mock.patch("google.cloud.managed_spark_connect.environment.sys") def test_is_interactive_true_via_flags(self, mock_sys, mock_get_ipython): # Simulate interactive environment via sys.flags.interactive if hasattr(mock_sys, "ps1"): @@ -575,7 +575,7 @@ def test_is_interactive_true_via_flags(self, mock_sys, mock_get_ipython): self.assertTrue(environment.is_interactive()) @mock.patch("IPython.get_ipython", side_effect=ImportError) - @mock.patch("google.cloud.dataproc_spark_connect.environment.sys") + @mock.patch("google.cloud.managed_spark_connect.environment.sys") def test_is_interactive_false(self, mock_sys, mock_get_ipython): # Simulate non-interactive environment if hasattr(mock_sys, "ps1"): @@ -594,14 +594,14 @@ def test_is_terminal_false(self, mock_stdin): self.assertFalse(environment.is_terminal()) @mock.patch("sys.stdin") - @mock.patch("google.cloud.dataproc_spark_connect.environment.sys") + @mock.patch("google.cloud.managed_spark_connect.environment.sys") def test_is_interactive_terminal_true(self, mock_sys, mock_stdin): mock_sys.ps1 = ">>>" mock_stdin.isatty.return_value = True self.assertTrue(environment.is_interactive_terminal()) @mock.patch("sys.stdin") - @mock.patch("google.cloud.dataproc_spark_connect.environment.sys") + @mock.patch("google.cloud.managed_spark_connect.environment.sys") @mock.patch("IPython.get_ipython", side_effect=ImportError) def test_is_interactive_terminal_false( self, mock_get_ipython, mock_sys, mock_stdin diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index 38e3e440..0edd19a9 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -14,8 +14,8 @@ import unittest from unittest import mock -from google.cloud.dataproc_spark_connect.session import DataprocSparkSession -from google.cloud.dataproc_spark_connect.exceptions import DataprocSparkConnectException +from google.cloud.managed_spark_connect.session import ManagedSparkSession +from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException class TestPythonVersionCheck(unittest.TestCase): @@ -30,14 +30,14 @@ def test_python_version_mismatch_warning_for_runtime_30(self): "sys.version_info", (client_py_major, client_py_minor, 0) ): with mock.patch("warnings.warn") as mock_warn: - session_builder = DataprocSparkSession.Builder() + session_builder = ManagedSparkSession.Builder() session_builder._check_python_version_compatibility( runtime_version ) expected_warning = ( f"Python version mismatch detected: Client is using Python {client_py_major}.{client_py_minor}, " - f"but Dataproc runtime {runtime_version} uses Python {server_py_major}.{server_py_minor}. " + f"but Managed Spark runtime {runtime_version} uses Python {server_py_major}.{server_py_minor}. " "This mismatch may cause issues with Python UDF (User Defined Function) compatibility. " f"Consider using Python {server_py_major}.{server_py_minor} for optimal UDF execution." ) @@ -53,7 +53,7 @@ def test_no_warning_when_python_versions_match_runtime_30(self): "sys.version_info", (client_py_major, client_py_minor, 0) ): with mock.patch("warnings.warn") as mock_warn: - session_builder = DataprocSparkSession.Builder() + session_builder = ManagedSparkSession.Builder() session_builder._check_python_version_compatibility( runtime_version ) @@ -64,7 +64,7 @@ def test_no_warning_for_unknown_runtime_version(self): """Test that no warning is shown for unknown runtime versions""" with mock.patch("sys.version_info", (3, 10, 0)): with mock.patch("warnings.warn") as mock_warn: - session_builder = DataprocSparkSession.Builder() + session_builder = ManagedSparkSession.Builder() session_builder._check_python_version_compatibility("unknown") mock_warn.assert_not_called() @@ -73,8 +73,8 @@ def test_no_warning_for_unknown_runtime_version(self): class TestRuntimeVersionCompatibility(unittest.TestCase): def test_older_runtimes_raise_exception(self): - """Test that runtime versions < MIN_SUPPORTED_RUNTIME_VERSION raise DataprocSparkConnectException""" - session_builder = DataprocSparkSession.Builder() + """Test that runtime versions < MIN_SUPPORTED_RUNTIME_VERSION raise ManagedSparkConnectException""" + session_builder = ManagedSparkSession.Builder() old_versions = ["2.4", "2.2", "1.0"] for version in old_versions: @@ -82,23 +82,21 @@ def test_older_runtimes_raise_exception(self): mock_dataproc_config = mock.Mock() mock_dataproc_config.runtime_config.version = version - with self.assertRaises( - DataprocSparkConnectException - ) as context: + with self.assertRaises(ManagedSparkConnectException) as context: session_builder._check_runtime_compatibility( mock_dataproc_config ) - min_version = DataprocSparkSession._MIN_RUNTIME_VERSION + min_version = ManagedSparkSession._MIN_RUNTIME_VERSION expected_message = ( - f"Specified {version} Dataproc Runtime version is not supported, " + f"Specified {version} Managed Spark Runtime version is not supported, " f"use {min_version} version or higher." ) self.assertEqual(str(context.exception), expected_message) def test_newer_runtimes_succeed(self): """Test that runtime versions >= MIN_RUNTIME_VERSION succeed""" - session_builder = DataprocSparkSession.Builder() + session_builder = ManagedSparkSession.Builder() new_versions = ["3.0", "3.1", "4.0"] for version in new_versions: @@ -110,15 +108,15 @@ def test_newer_runtimes_succeed(self): session_builder._check_runtime_compatibility( mock_dataproc_config ) - except DataprocSparkConnectException: + except ManagedSparkConnectException: self.fail( - f"_check_runtime_compatibility raised DataprocSparkConnectException unexpectedly for version {version}" + f"_check_runtime_compatibility raised ManagedSparkConnectException unexpectedly for version {version}" ) - @mock.patch("google.cloud.dataproc_spark_connect.session.logger") + @mock.patch("google.cloud.managed_spark_connect.session.logger") def test_invalid_runtime_version_logs_warning(self, mock_logger): """Test that invalid runtime versions are logged as warnings but don't fail""" - session_builder = DataprocSparkSession.Builder() + session_builder = ManagedSparkSession.Builder() # Mock dataproc config with invalid runtime version mock_dataproc_config = mock.Mock() diff --git a/tests/unit/test_proxy.py b/tests/unit/test_proxy.py index fece7d84..23339391 100644 --- a/tests/unit/test_proxy.py +++ b/tests/unit/test_proxy.py @@ -17,7 +17,7 @@ import pytest -from google.cloud.dataproc_spark_connect.client.proxy import connect_sockets +from google.cloud.managed_spark_connect.client.proxy import connect_sockets @pytest.fixture diff --git a/tests/unit/test_pypi_artifacts.py b/tests/unit/test_pypi_artifacts.py index 22ef3600..da073ee2 100644 --- a/tests/unit/test_pypi_artifacts.py +++ b/tests/unit/test_pypi_artifacts.py @@ -4,7 +4,7 @@ from packaging.requirements import InvalidRequirement -from google.cloud.dataproc_spark_connect.pypi_artifacts import PyPiArtifacts +from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts class PyPiArtifactsTest(unittest.TestCase): diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 2b1a6245..624a8ca4 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -22,9 +22,14 @@ InvalidArgument, NotFound, ) -from google.cloud.dataproc_spark_connect import DataprocSparkSession -from google.cloud.dataproc_spark_connect.exceptions import DataprocSparkConnectException -from google.cloud.dataproc_spark_connect.session import _is_valid_label_value, _is_valid_session_id +from google.cloud.managed_spark_connect import ManagedSparkSession +from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException +from google.cloud.managed_spark_connect.session import ( + _env_var_set, + _getenv_with_deprecated_alias, + _is_valid_label_value, + _is_valid_session_id, +) from google.cloud.dataproc_v1 import ( AuthenticationConfig, CreateSessionRequest, @@ -38,16 +43,16 @@ from pyspark.sql.connect.proto import Command, ConfigResponse, ExecutePlanRequest, Plan, Relation, SQL, SqlCommand, UserContext from unittest import mock -_DATAPROC_SESSIONS_BASE_URL = ( +_MANAGED_SPARK_SESSIONS_BASE_URL = ( "https://console.cloud.google.com/dataproc/interactive" ) -class DataprocRemoteSparkSessionBuilderTests(unittest.TestCase): +class ManagedSparkSessionBuilderTests(unittest.TestCase): def setUp(self): self._default_runtime_version = ( - DataprocSparkSession._DEFAULT_RUNTIME_VERSION + ManagedSparkSession._DEFAULT_RUNTIME_VERSION ) self.original_environment = dict(os.environ) os.environ.clear() @@ -71,7 +76,7 @@ def stopSession(mock_session_controller_client_instance, session): @staticmethod def _setup_session_creation_mocks( mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -83,7 +88,7 @@ def _setup_session_creation_mocks( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = session_id + mock_session_id.return_value = session_id mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) @@ -109,13 +114,13 @@ def _setup_session_creation_mocks( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.get_client_environment_label" + "google.cloud.managed_spark_connect.environment.get_client_environment_label" ) @mock.patch( "IPython.core.interactiveshell.InteractiveShell.initialized", @@ -134,7 +139,7 @@ def test_create_spark_session_with_default_notebook_behavior( mock_interactive_shell, mock_get_client_environment_label, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -146,7 +151,7 @@ def test_create_spark_session_with_default_notebook_behavior( ) session_id = "sc-20240702-103952-abcdef" - mock_dataproc_session_id.return_value = session_id + mock_session_id.return_value = session_id mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) @@ -167,7 +172,7 @@ def test_create_spark_session_with_default_notebook_behavior( mock_ipython_utils = mock.sys.modules[ "google.cloud.aiplatform.utils" ]._ipython_utils - test_session_url = f"{_DATAPROC_SESSIONS_BASE_URL}/test-region/{session_id}?project=test-project" + test_session_url = f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/test-region/{session_id}?project=test-project" mock_display_link = mock_ipython_utils.display_link mock.patch.dict( os.environ, @@ -193,7 +198,7 @@ def test_create_spark_session_with_default_notebook_behavior( ) try: session = ( - DataprocSparkSession.builder.projectId("test-project") + ManagedSparkSession.builder.projectId("test-project") .location("test-region") .getOrCreate() ) @@ -239,8 +244,8 @@ def test_pypi_add_artifacts( mock_session_controller_client_instance.create_session.return_value = ( mock_operation ) - session = DataprocSparkSession.builder.getOrCreate() - self.assertTrue(isinstance(session, DataprocSparkSession)) + session = ManagedSparkSession.builder.getOrCreate() + self.assertTrue(isinstance(session, ManagedSparkSession)) session.addArtifact = mock.MagicMock() # Setting two flags together @@ -273,15 +278,15 @@ def test_pypi_add_artifacts( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_create_session_with_user_provided_dataproc_config( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -294,7 +299,7 @@ def test_create_session_with_user_provided_dataproc_config( mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" cred = mock.MagicMock() cred.token = "token" mock_credentials.return_value = (cred, "") @@ -346,7 +351,7 @@ def test_create_session_with_user_provided_dataproc_config( "spark.executor.cores": "8" } session = ( - DataprocSparkSession.builder.config("spark.executor.cores", "6") + ManagedSparkSession.builder.config("spark.executor.cores", "6") .dataprocSessionConfig(dataproc_config) .config("spark.executor.cores", "16") .getOrCreate() @@ -373,15 +378,15 @@ def test_create_session_with_user_provided_dataproc_config( @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_create_session_with_env_vars_config( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_session_controller_client, mock_credentials, ): @@ -390,7 +395,7 @@ def test_create_session_with_env_vars_config( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" cred = mock.MagicMock() cred.token = "token" mock_credentials.return_value = (cred, "") @@ -408,11 +413,11 @@ def test_create_session_with_env_vars_config( mock.patch.dict( os.environ, { - "DATAPROC_SPARK_CONNECT_AUTH_TYPE": "SERVICE_ACCOUNT", - "DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT": "test-acc@example.com", - "DATAPROC_SPARK_CONNECT_SUBNET": "test-subnet-from-env", - "DATAPROC_SPARK_CONNECT_TTL_SECONDS": "12", - "DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS": "89", + "MANAGED_SPARK_CONNECT_AUTH_TYPE": "SERVICE_ACCOUNT", + "MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT": "test-acc@example.com", + "MANAGED_SPARK_CONNECT_SUBNET": "test-subnet-from-env", + "MANAGED_SPARK_CONNECT_TTL_SECONDS": "12", + "MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS": "89", "COLAB_NOTEBOOK_ID": "/embedded/projects/company.com%3Aproject1/locations/us-central1/repositories/test-notebook-id", }, ).start() @@ -452,7 +457,7 @@ def test_create_session_with_env_vars_config( ) try: - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() mock_session_controller_client_instance.create_session.assert_called_once_with( create_session_request ) @@ -476,15 +481,15 @@ def test_create_session_with_env_vars_config( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_create_session_with_session_template( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -494,7 +499,7 @@ def test_create_session_with_session_template( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) @@ -532,7 +537,7 @@ def test_create_session_with_session_template( try: dataproc_config = Session() dataproc_config.session_template = "projects/test-project/locations/test-region/sessionTemplates/test_template" - session = DataprocSparkSession.builder.dataprocSessionConfig( + session = ManagedSparkSession.builder.dataprocSessionConfig( dataproc_config ).getOrCreate() mock_session_controller_client_instance.create_session.assert_called_once_with( @@ -558,15 +563,15 @@ def test_create_session_with_session_template( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_create_session_with_user_provided_dataproc_config_and_session_template( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -576,7 +581,7 @@ def test_create_session_with_user_provided_dataproc_config_and_session_template( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) @@ -620,7 +625,7 @@ def test_create_session_with_user_provided_dataproc_config_and_session_template( "seconds": 10 } dataproc_config.session_template = "projects/test-project/locations/test-region/sessionTemplates/test_template" - session = DataprocSparkSession.builder.dataprocSessionConfig( + session = ManagedSparkSession.builder.dataprocSessionConfig( dataproc_config ).getOrCreate() mock_session_controller_client_instance.create_session.assert_called_once_with( @@ -645,15 +650,15 @@ def test_create_session_with_user_provided_dataproc_config_and_session_template( @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) def test_create_spark_session_with_create_session_failed( self, - mock_dataproc_session_id, + mock_session_id, mock_session_controller_client, mock_credentials, ): - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) @@ -668,11 +673,11 @@ def test_create_spark_session_with_create_session_failed( cred.token = "token" mock_credentials.return_value = (cred, "") with self.assertRaises(RuntimeError) as e: - DataprocSparkSession.builder.dataprocSessionConfig( + ManagedSparkSession.builder.dataprocSessionConfig( Session() ).getOrCreate() self.assertEqual( - "Error while creating Dataproc Session", e.exception.args[0] + "Error while creating Managed Spark Session", e.exception.args[0] ) @mock.patch("google.auth.default") @@ -695,28 +700,28 @@ def test_create_spark_session_with_invalid_argument( cred = mock.MagicMock() cred.token = "token" mock_credentials.return_value = (cred, "") - with self.assertRaises(DataprocSparkConnectException) as e: - DataprocSparkSession.builder.dataprocSessionConfig( + with self.assertRaises(ManagedSparkConnectException) as e: + ManagedSparkSession.builder.dataprocSessionConfig( Session() ).getOrCreate() self.assertEqual( e.exception.error_message, - "Error while creating Dataproc Session: " + "Error while creating Managed Spark Session: " "400 Network does not have permissions", ) @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_spark_session_with_inactive_s8s_session( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_session_controller_client, mock_credentials, ): @@ -726,7 +731,7 @@ def test_spark_session_with_inactive_s8s_session( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" cred = mock.MagicMock() cred.token = "token" @@ -742,7 +747,7 @@ def test_spark_session_with_inactive_s8s_session( mock_operation ) with self.assertRaises(RuntimeError) as e: - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() session.createDataFrame([(1, "Sarah"), (2, "Maria")]).toDF( "id", "name" ).show() @@ -756,7 +761,7 @@ def test_spark_session_with_inactive_s8s_session( @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_stop_spark_session_with_terminated_s8s_session( self, @@ -787,7 +792,7 @@ def test_stop_spark_session_with_terminated_s8s_session( mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() finally: mock_session_controller_client_instance.terminate_session.side_effect = FailedPrecondition( @@ -795,13 +800,13 @@ def test_stop_spark_session_with_terminated_s8s_session( ) if session is not None: session.stop() - self.assertIsNone(DataprocSparkSession._active_s8s_session_uuid) + self.assertIsNone(ManagedSparkSession._active_s8s_session_uuid) @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_stop_spark_session_with_creating_s8s_session( self, @@ -832,7 +837,7 @@ def test_stop_spark_session_with_creating_s8s_session( mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() finally: mock_session_controller_client_instance.terminate_session.side_effect = Aborted( @@ -840,13 +845,13 @@ def test_stop_spark_session_with_creating_s8s_session( ) if session is not None: session.stop() - self.assertIsNone(DataprocSparkSession._active_s8s_session_uuid) + self.assertIsNone(ManagedSparkSession._active_s8s_session_uuid) @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_stop_spark_session_with_deleted_s8s_session( self, @@ -877,7 +882,7 @@ def test_stop_spark_session_with_deleted_s8s_session( mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() finally: mock_session_controller_client_instance.terminate_session.side_effect = NotFound( @@ -885,28 +890,28 @@ def test_stop_spark_session_with_deleted_s8s_session( ) if session is not None: session.stop() - self.assertIsNone(DataprocSparkSession._active_s8s_session_uuid) + self.assertIsNone(ManagedSparkSession._active_s8s_session_uuid) @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_stop_spark_session_wait_for_terminating_state( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_session_controller_client, mock_credentials, mock_client_config, ): session = None mock_is_s8s_session_active.return_value = True - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) @@ -927,7 +932,7 @@ def test_stop_spark_session_wait_for_terminating_state( mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() finally: mock_session_controller_client_instance.terminate_session.return_value = ( @@ -949,19 +954,19 @@ def test_stop_spark_session_wait_for_terminating_state( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.logger" + "google.cloud.managed_spark_connect.session.logger" ) # Mock the logger def test_create_session_with_default_datasource_env_var( self, mock_logger, # Add mock logger parameter mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -971,7 +976,7 @@ def test_create_session_with_default_datasource_env_var( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = ( + mock_session_id.return_value = ( "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" # Use a valid UUID ) mock_client_config.return_value = ConfigResult.fromProto( @@ -1000,11 +1005,11 @@ def test_create_session_with_default_datasource_env_var( mock_operation ) - # Scenario 1: DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE is not set + # Scenario 1: MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE is not set with mock.patch.dict(os.environ, {}, clear=True): os.environ["GOOGLE_CLOUD_PROJECT"] = "test-project" os.environ["GOOGLE_CLOUD_REGION"] = "test-region" - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() create_session_request = mock_session_controller_client_instance.create_session.call_args[ 0 ][ @@ -1019,15 +1024,15 @@ def test_create_session_with_default_datasource_env_var( mock_session_controller_client_instance.create_session.reset_mock() mock_logger.warning.reset_mock() - # Scenario 2: DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE is set to "bigquery" + # Scenario 2: MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE is set to "bigquery" with mock.patch.dict( os.environ, - {"DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE": "bigquery"}, + {"MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE": "bigquery"}, clear=True, ): os.environ["GOOGLE_CLOUD_PROJECT"] = "test-project" os.environ["GOOGLE_CLOUD_REGION"] = "test-region" - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() create_session_request = mock_session_controller_client_instance.create_session.call_args[ 0 ][ @@ -1051,15 +1056,15 @@ def test_create_session_with_default_datasource_env_var( mock_session_controller_client_instance.create_session.reset_mock() mock_logger.warning.reset_mock() - # Scenario 3: DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE is set to an invalid value + # Scenario 3: MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE is set to an invalid value with mock.patch.dict( os.environ, - {"DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE": "invalid_datasource"}, + {"MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE": "invalid_datasource"}, clear=True, ): os.environ["GOOGLE_CLOUD_PROJECT"] = "test-project" os.environ["GOOGLE_CLOUD_REGION"] = "test-region" - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() create_session_request = mock_session_controller_client_instance.create_session.call_args[ 0 ][ @@ -1070,16 +1075,16 @@ def test_create_session_with_default_datasource_env_var( create_session_request.session.runtime_config.properties, ) mock_logger.warning.assert_called_once_with( - "DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE is set to an invalid value: invalid_datasource. Supported value is 'bigquery'." + "MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE is set to an invalid value: invalid_datasource. Supported value is 'bigquery'." ) self.stopSession(mock_session_controller_client_instance, session) mock_session_controller_client_instance.create_session.reset_mock() mock_logger.warning.reset_mock() - # Scenario 4: DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE is set to "bigquery" with pre-existing properties + # Scenario 4: MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE is set to "bigquery" with pre-existing properties with mock.patch.dict( os.environ, - {"DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE": "bigquery"}, + {"MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE": "bigquery"}, clear=True, ): os.environ["GOOGLE_CLOUD_PROJECT"] = "test-project" @@ -1090,7 +1095,7 @@ def test_create_session_with_default_datasource_env_var( "spark.sql.sources.default": "override_source", "spark.some.other.property": "some_value", } - session = DataprocSparkSession.builder.dataprocSessionConfig( + session = ManagedSparkSession.builder.dataprocSessionConfig( dataproc_config ).getOrCreate() create_session_request = mock_session_controller_client_instance.create_session.call_args[ @@ -1128,7 +1133,7 @@ def test_create_session_with_default_datasource_env_var( "IPython.core.interactiveshell.InteractiveShell.initialized", return_value=True, ) - @mock.patch("google.cloud.dataproc_spark_connect.session.logger") + @mock.patch("google.cloud.managed_spark_connect.session.logger") def test_display_button_with_aiplatform_not_installed( self, mock_logger, _mock_ipy ): @@ -1138,7 +1143,7 @@ def test_display_button_with_aiplatform_not_installed( "VERTEX_PRODUCT": "COLAB_ENTERPRISE", }, ).start() - DataprocSparkSession.builder._display_view_session_details_button( + ManagedSparkSession.builder._display_view_session_details_button( "test_session" ) mock_logger.debug.assert_called_once_with( @@ -1169,10 +1174,10 @@ def test_display_button_with_aiplatform_installed_ipython_interactive( mock_ipython_utils = mock.sys.modules[ "google.cloud.aiplatform.utils" ]._ipython_utils - test_session_url = f"{_DATAPROC_SESSIONS_BASE_URL}/test-region/test_session?project=test-project" + test_session_url = f"{_MANAGED_SPARK_SESSIONS_BASE_URL}/test-region/test_session?project=test-project" mock_display_link = mock_ipython_utils.display_link - DataprocSparkSession.builder._display_view_session_details_button( + ManagedSparkSession.builder._display_view_session_details_button( "test_session" ) mock_display_link.assert_called_once_with( @@ -1205,7 +1210,7 @@ def test_display_button_with_aiplatform_installed_ipython_non_interactive( ]._ipython_utils mock_display_link = mock_ipython_utils.display_link - DataprocSparkSession.builder._display_view_session_details_button( + ManagedSparkSession.builder._display_view_session_details_button( "test_session" ) mock_display_link.assert_not_called() @@ -1226,15 +1231,15 @@ def test_display_session_link_on_creation_colab_enterprise( "VERTEX_PRODUCT": "COLAB_ENTERPRISE", }, ).start() - DataprocSparkSession.builder._display_session_link_on_creation( + ManagedSparkSession.builder._display_session_link_on_creation( "test_session" ) mock_display.assert_called_once() args, _ = mock_display.call_args html_output = args[0].data - self.assertIn("Creating Dataproc Spark Session", html_output) - self.assertNotIn("Dataproc Session", html_output) + self.assertIn("Creating Managed Spark Connect Session", html_output) + self.assertNotIn("Managed Spark Session", html_output) @mock.patch( "IPython.core.interactiveshell.InteractiveShell.initialized", @@ -1250,15 +1255,15 @@ def test_display_session_link_on_creation_not_colab_enterprise( os.environ, {}, ).start() - DataprocSparkSession.builder._display_session_link_on_creation( + ManagedSparkSession.builder._display_session_link_on_creation( "test_session" ) mock_display.assert_called_once() args, _ = mock_display.call_args html_output = args[0].data - self.assertIn("Creating Dataproc Spark Session", html_output) - self.assertIn("Dataproc Session", html_output) + self.assertIn("Creating Managed Spark Connect Session", html_output) + self.assertIn("Managed Spark Session", html_output) def test_is_valid_label_value(self): # Valid label values @@ -1303,17 +1308,17 @@ def test_is_valid_label_value(self): @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) - @mock.patch("google.cloud.dataproc_spark_connect.session.logger") + @mock.patch("google.cloud.managed_spark_connect.session.logger") def test_create_session_with_invalid_notebook_id( self, mock_logger, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_session_controller_client, mock_credentials, ): @@ -1322,7 +1327,7 @@ def test_create_session_with_invalid_notebook_id( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" cred = mock.MagicMock() cred.token = "token" mock_credentials.return_value = (cred, "") @@ -1363,7 +1368,7 @@ def test_create_session_with_invalid_notebook_id( # Note: No notebook label should be set due to invalid format try: - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() mock_session_controller_client_instance.create_session.assert_called_once_with( create_session_request ) @@ -1395,17 +1400,17 @@ def test_create_session_with_invalid_notebook_id( @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) - @mock.patch("google.cloud.dataproc_spark_connect.session.logger") + @mock.patch("google.cloud.managed_spark_connect.session.logger") def test_create_session_with_valid_notebook_id( self, mock_logger, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_session_controller_client, mock_credentials, ): @@ -1414,7 +1419,7 @@ def test_create_session_with_valid_notebook_id( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" cred = mock.MagicMock() cred.token = "token" mock_credentials.return_value = (cred, "") @@ -1458,7 +1463,7 @@ def test_create_session_with_valid_notebook_id( ) try: - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() mock_session_controller_client_instance.create_session.assert_called_once_with( create_session_request ) @@ -1475,32 +1480,32 @@ def test_create_session_without_project_id(self): """Tests that an exception is raised when project ID is not provided.""" os.environ.clear() try: - DataprocSparkSession.builder.location("test-region").getOrCreate() - except DataprocSparkConnectException as e: + ManagedSparkSession.builder.location("test-region").getOrCreate() + except ManagedSparkConnectException as e: self.assertIn("project ID is not set", str(e)) def test_create_session_without_location(self): """Tests that an exception is raised when location is not provided.""" os.environ.clear() try: - DataprocSparkSession.builder.projectId("test-project").getOrCreate() - except DataprocSparkConnectException as e: + ManagedSparkSession.builder.projectId("test-project").getOrCreate() + except ManagedSparkConnectException as e: self.assertIn("location is not set", str(e)) def test_create_session_without_application_default_credentials(self): """Tests that an exception is raised when application default credentials is not provided.""" os.environ.clear() try: - DataprocSparkSession.builder.location("test-region").projectId( + ManagedSparkSession.builder.location("test-region").projectId( "test-project" ).getOrCreate() - except DataprocSparkConnectException as e: + except ManagedSparkConnectException as e: self.assertIn( - "Credentials error while creating Dataproc Session", str(e) + "Credentials error while creating Managed Spark Session", str(e) ) -class DataprocSparkConnectClientTest(unittest.TestCase): +class ManagedSparkConnectClientTest(unittest.TestCase): def setUp(self): self.original_environment = dict(os.environ) @@ -1521,7 +1526,7 @@ def stopSession(mock_session_controller_client_instance, session): @staticmethod def _setup_session_creation_mocks( mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -1533,7 +1538,7 @@ def _setup_session_creation_mocks( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = session_id + mock_session_id.return_value = session_id mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) @@ -1559,10 +1564,10 @@ def _setup_session_creation_mocks( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) @mock.patch("uuid.uuid4") @mock.patch( @@ -1573,7 +1578,7 @@ def test_execute_plan_request_default_behaviour( mock_super_execute_plan_request, mock_uuid4, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -1597,7 +1602,7 @@ def test_execute_plan_request_default_behaviour( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) @@ -1616,7 +1621,7 @@ def test_execute_plan_request_default_behaviour( ) try: - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() client = session.client result_request = client._execute_plan_request_with_metadata() @@ -1651,10 +1656,10 @@ def test_execute_plan_request_default_behaviour( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) @mock.patch("uuid.uuid4") @mock.patch( @@ -1665,7 +1670,7 @@ def test_execute_plan_request_with_operation_id_provided( mock_super_execute_plan_request, mock_uuid4, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -1690,7 +1695,7 @@ def test_execute_plan_request_with_operation_id_provided( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) @@ -1709,7 +1714,7 @@ def test_execute_plan_request_with_operation_id_provided( ) try: - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() client = session.client result_request = client._execute_plan_request_with_metadata() @@ -1789,17 +1794,17 @@ def test_sql_lazy_transformation(self): ) self.assertTrue( - DataprocSparkSession._sql_lazy_transformation( + ManagedSparkSession._sql_lazy_transformation( test_execute_plan_request_1 ) ) self.assertFalse( - DataprocSparkSession._sql_lazy_transformation( + ManagedSparkSession._sql_lazy_transformation( test_execute_plan_request_2 ) ) self.assertFalse( - DataprocSparkSession._sql_lazy_transformation( + ManagedSparkSession._sql_lazy_transformation( test_execute_plan_request_3 ) ) @@ -1808,15 +1813,15 @@ def test_sql_lazy_transformation(self): @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_builder_pattern_runtime_config( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -1825,7 +1830,7 @@ def test_builder_pattern_runtime_config( mock_session_controller_client_instance = ( self._setup_session_creation_mocks( mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -1834,7 +1839,7 @@ def test_builder_pattern_runtime_config( try: session = ( - DataprocSparkSession.builder.runtimeVersion("3.0") + ManagedSparkSession.builder.runtimeVersion("3.0") .config( "spark.executor.cores", "8" ) # Use existing Spark config method @@ -1865,15 +1870,15 @@ def test_builder_pattern_runtime_config( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_builder_pattern_environment_config( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -1882,7 +1887,7 @@ def test_builder_pattern_environment_config( mock_session_controller_client_instance = ( self._setup_session_creation_mocks( mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -1891,7 +1896,7 @@ def test_builder_pattern_environment_config( try: session = ( - DataprocSparkSession.builder.serviceAccount( + ManagedSparkSession.builder.serviceAccount( "test-service@project.iam.gserviceaccount.com" ) .subnetwork( @@ -1943,15 +1948,15 @@ def test_builder_pattern_environment_config( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_service_account_sets_auth_type_automatically( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -1961,7 +1966,7 @@ def test_service_account_sets_auth_type_automatically( mock_session_controller_client_instance = ( self._setup_session_creation_mocks( mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -1969,7 +1974,7 @@ def test_service_account_sets_auth_type_automatically( ) try: - session = DataprocSparkSession.builder.serviceAccount( + session = ManagedSparkSession.builder.serviceAccount( "test-service@project.iam.gserviceaccount.com" ).getOrCreate() @@ -2002,15 +2007,15 @@ def test_service_account_sets_auth_type_automatically( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_builder_pattern_ttl_with_timedelta( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2019,7 +2024,7 @@ def test_builder_pattern_ttl_with_timedelta( mock_session_controller_client_instance = ( self._setup_session_creation_mocks( mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2029,7 +2034,7 @@ def test_builder_pattern_ttl_with_timedelta( try: # Test using timedelta objects session = ( - DataprocSparkSession.builder.ttl(datetime.timedelta(hours=1)) + ManagedSparkSession.builder.ttl(datetime.timedelta(hours=1)) .idleTtl(datetime.timedelta(minutes=30)) .getOrCreate() ) @@ -2069,15 +2074,15 @@ def test_builder_pattern_ttl_with_timedelta( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) - def test_builder_pattern_session_template_and_labels( + def test_builder_pattern_runtime_profile_and_labels( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2086,7 +2091,7 @@ def test_builder_pattern_session_template_and_labels( mock_session_controller_client_instance = ( self._setup_session_creation_mocks( mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2095,7 +2100,7 @@ def test_builder_pattern_session_template_and_labels( try: session = ( - DataprocSparkSession.builder.sessionTemplate( + ManagedSparkSession.builder.runtimeProfile( "projects/test-project/locations/us-central1/sessionTemplates/test-template" ) .label("environment", "production") @@ -2104,7 +2109,7 @@ def test_builder_pattern_session_template_and_labels( .getOrCreate() ) - # Verify the session was created with the correct session template and labels + # Verify the session was created with the correct Runtime Profile and labels create_session_request = mock_session_controller_client_instance.create_session.call_args[ 0 ][ @@ -2135,19 +2140,31 @@ def test_builder_pattern_session_template_and_labels( ) self.stopSession(mock_session_controller_client_instance, session) + def test_session_template_is_deprecated_alias_for_runtime_profile(self): + """sessionTemplate() should still work but warn in favor of runtimeProfile().""" + builder = ManagedSparkSession.Builder() + with self.assertWarns(DeprecationWarning): + builder.sessionTemplate( + "projects/test-project/locations/us-central1/sessionTemplates/test-template" + ) + self.assertEqual( + builder.dataproc_config.session_template, + "projects/test-project/locations/us-central1/sessionTemplates/test-template", + ) + @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_builder_pattern_combined_with_dataprocSessionConfig( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2156,7 +2173,7 @@ def test_builder_pattern_combined_with_dataprocSessionConfig( mock_session_controller_client_instance = ( self._setup_session_creation_mocks( mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2171,7 +2188,7 @@ def test_builder_pattern_combined_with_dataprocSessionConfig( base_config.labels["base-label"] = "base-value" session = ( - DataprocSparkSession.builder.dataprocSessionConfig(base_config) + ManagedSparkSession.builder.dataprocSessionConfig(base_config) .config( "spark.executor.cores", "8" ) # Override using existing Spark method @@ -2210,17 +2227,17 @@ def test_builder_pattern_combined_with_dataprocSessionConfig( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) - @mock.patch("google.cloud.dataproc_spark_connect.session.logger") + @mock.patch("google.cloud.managed_spark_connect.session.logger") def test_builder_pattern_system_label_protection( self, mock_logger, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2229,7 +2246,7 @@ def test_builder_pattern_system_label_protection( mock_session_controller_client_instance = ( self._setup_session_creation_mocks( mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2238,7 +2255,7 @@ def test_builder_pattern_system_label_protection( try: session = ( - DataprocSparkSession.builder.label( + ManagedSparkSession.builder.label( "dataproc-session-client", "malicious-override" ) # Try to override system label .label( @@ -2310,19 +2327,19 @@ def test_builder_pattern_system_label_protection( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) @mock.patch( - "google.cloud.dataproc_spark_connect.environment.get_client_environment_label" + "google.cloud.managed_spark_connect.environment.get_client_environment_label" ) def test_create_session_with_client_environment_label( self, mock_get_client_environment_label, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2333,9 +2350,7 @@ def test_create_session_with_client_environment_label( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = ( - "6fa459ea-ee8a-3ca4-894e-db77e160355e" - ) + mock_session_id.return_value = "6fa459ea-ee8a-3ca4-894e-db77e160355e" mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) @@ -2392,12 +2407,12 @@ def test_create_session_with_client_environment_label( try: # Reset singleton state before each subtest run - DataprocSparkSession._active_s8s_session_id = None - DataprocSparkSession._default_session = None + ManagedSparkSession._active_s8s_session_id = None + ManagedSparkSession._default_session = None # Set up project and region for the builder session = ( - DataprocSparkSession.builder.projectId("test-project") + ManagedSparkSession.builder.projectId("test-project") .location("test-region") .getOrCreate() ) @@ -2416,15 +2431,15 @@ def test_create_session_with_client_environment_label( @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @mock.patch( - "google.cloud.dataproc_spark_connect.DataprocSparkSession.Builder.generate_dataproc_session_id" + "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id" ) @mock.patch( - "google.cloud.dataproc_spark_connect.session.is_s8s_session_active" + "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) def test_execution_progress_handler( self, mock_is_s8s_session_active, - mock_dataproc_session_id, + mock_session_id, mock_client_config, mock_session_controller_client, mock_credentials, @@ -2434,7 +2449,7 @@ def test_execution_progress_handler( mock_session_controller_client_instance = ( mock_session_controller_client.return_value ) - mock_dataproc_session_id.return_value = "sc-20240702-103952-abcdef" + mock_session_id.return_value = "sc-20240702-103952-abcdef" mock_client_config.return_value = ConfigResult.fromProto( ConfigResponse() ) @@ -2453,13 +2468,13 @@ def test_execution_progress_handler( ) try: - session = DataprocSparkSession.builder.getOrCreate() + session = ManagedSparkSession.builder.getOrCreate() client = session.client - # By default Dataproc handler is registered + # By default Managed Spark handler is registered self.assertEqual(len(client._progress_handlers), 1) - # Dataproc handler isn't cleared with clearProgressHandlers() method + # Managed Spark handler isn't cleared with clearProgressHandlers() method session.clearProgressHandlers() self.assertEqual(len(client._progress_handlers), 1) @@ -2498,7 +2513,7 @@ def test_wait_for_session_available_success( session_ready, ] - builder = DataprocSparkSession.Builder() + builder = ManagedSparkSession.Builder() builder._session_controller_client = ( mock_client # Inject the mock client ) @@ -2526,7 +2541,7 @@ def test_wait_for_session_available_timeout( mock_client.get_session.return_value = session_pending - builder = DataprocSparkSession.Builder() + builder = ManagedSparkSession.Builder() builder._session_controller_client = ( mock_client # Inject the mock client ) @@ -2583,7 +2598,7 @@ def test_invalid_session_ids(self): def test_dataproc_session_id_builder_method(self): """Test the dataprocSessionId() builder method.""" - builder = DataprocSparkSession.builder + builder = ManagedSparkSession.builder # Test valid session ID result = builder.dataprocSessionId("test-session") @@ -2596,7 +2611,7 @@ def test_dataproc_session_id_builder_method(self): self.assertIn("Invalid session ID", str(context.exception)) @mock.patch( - "google.cloud.dataproc_spark_connect.session.SessionControllerClient" + "google.cloud.managed_spark_connect.session.SessionControllerClient" ) def test_session_reuse_with_custom_id(self, mock_session_controller_client): """Test that sessions are reused when custom ID is provided.""" @@ -2611,7 +2626,7 @@ def test_session_reuse_with_custom_id(self, mock_session_controller_client): } mock_client.get_session.return_value = active_session - builder = DataprocSparkSession.Builder() + builder = ManagedSparkSession.Builder() builder._project_id = "test-project" builder._region = "test-region" builder._custom_session_id = "my-session" @@ -2622,7 +2637,7 @@ def test_session_reuse_with_custom_id(self, mock_session_controller_client): mock_client.get_session.assert_called_once() @mock.patch( - "google.cloud.dataproc_spark_connect.session.SessionControllerClient" + "google.cloud.managed_spark_connect.session.SessionControllerClient" ) def test_session_skip_terminated(self, mock_session_controller_client): """Test that terminated sessions are skipped, not cleaned up.""" @@ -2633,7 +2648,7 @@ def test_session_skip_terminated(self, mock_session_controller_client): terminated_session.state = Session.State.TERMINATED mock_client.get_session.return_value = terminated_session - builder = DataprocSparkSession.Builder() + builder = ManagedSparkSession.Builder() builder._project_id = "test-project" builder._region = "test-region" builder._custom_session_id = "my-session" @@ -2644,5 +2659,69 @@ def test_session_skip_terminated(self, mock_session_controller_client): mock_client.get_session.assert_called_once() +class DeprecatedEnvVarAliasTests(unittest.TestCase): + """Test cases for the MANAGED_SPARK_CONNECT_* / DATAPROC_SPARK_CONNECT_* env var fallback.""" + + def setUp(self): + for name in ( + "MANAGED_SPARK_CONNECT_TEST_VAR", + "DATAPROC_SPARK_CONNECT_TEST_VAR", + ): + os.environ.pop(name, None) + + tearDown = setUp + + def test_new_name_takes_precedence(self): + os.environ["MANAGED_SPARK_CONNECT_TEST_VAR"] = "new" + os.environ["DATAPROC_SPARK_CONNECT_TEST_VAR"] = "old" + self.assertEqual( + _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_TEST_VAR", + "DATAPROC_SPARK_CONNECT_TEST_VAR", + ), + "new", + ) + + def test_old_name_used_with_deprecation_warning(self): + os.environ["DATAPROC_SPARK_CONNECT_TEST_VAR"] = "old" + with self.assertWarns(DeprecationWarning): + value = _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_TEST_VAR", + "DATAPROC_SPARK_CONNECT_TEST_VAR", + ) + self.assertEqual(value, "old") + + def test_default_when_neither_set(self): + self.assertIsNone( + _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_TEST_VAR", + "DATAPROC_SPARK_CONNECT_TEST_VAR", + ) + ) + self.assertEqual( + _getenv_with_deprecated_alias( + "MANAGED_SPARK_CONNECT_TEST_VAR", + "DATAPROC_SPARK_CONNECT_TEST_VAR", + "fallback", + ), + "fallback", + ) + + def test_env_var_set_checks_both_names(self): + self.assertFalse( + _env_var_set( + "MANAGED_SPARK_CONNECT_TEST_VAR", + "DATAPROC_SPARK_CONNECT_TEST_VAR", + ) + ) + os.environ["DATAPROC_SPARK_CONNECT_TEST_VAR"] = "old" + self.assertTrue( + _env_var_set( + "MANAGED_SPARK_CONNECT_TEST_VAR", + "DATAPROC_SPARK_CONNECT_TEST_VAR", + ) + ) + + if __name__ == "__main__": unittest.main() From 719168c13c99b6d90e44a7809ca5cfaa99e0da7f Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:57:04 -0700 Subject: [PATCH 2/8] style: reformat with pyink to match latest version used in CI requirements-dev.txt pins pyink~=24.0, but the pyink CI workflow installs the latest unpinned version (26.5.1), which enforces a blank line after module docstrings. --- google/cloud/dataproc_magics/__init__.py | 1 + google/cloud/dataproc_magics/magics.py | 1 + google/cloud/dataproc_spark_connect/__init__.py | 1 + google/cloud/dataproc_spark_connect/client/__init__.py | 1 + google/cloud/dataproc_spark_connect/client/core.py | 1 + google/cloud/dataproc_spark_connect/environment.py | 1 + google/cloud/dataproc_spark_connect/exceptions.py | 1 + google/cloud/dataproc_spark_connect/session.py | 1 + tests/integration/managed_spark_magics/test_magics.py | 1 - tests/integration/test_session.py | 1 - tests/unit/test_deprecated_shims.py | 1 + 11 files changed, 9 insertions(+), 2 deletions(-) diff --git a/google/cloud/dataproc_magics/__init__.py b/google/cloud/dataproc_magics/__init__.py index a7001f09..6cce7b3d 100644 --- a/google/cloud/dataproc_magics/__init__.py +++ b/google/cloud/dataproc_magics/__init__.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Deprecated: this package has been renamed to ``google.cloud.managed_spark_magics``.""" + import warnings from .magics import DataprocMagics diff --git a/google/cloud/dataproc_magics/magics.py b/google/cloud/dataproc_magics/magics.py index 014d519f..b243be84 100644 --- a/google/cloud/dataproc_magics/magics.py +++ b/google/cloud/dataproc_magics/magics.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Deprecated: use ``google.cloud.managed_spark_magics.magics`` instead.""" + from google.cloud.managed_spark_magics.magics import ManagedSparkMagics DataprocMagics = ManagedSparkMagics diff --git a/google/cloud/dataproc_spark_connect/__init__.py b/google/cloud/dataproc_spark_connect/__init__.py index 862d30dc..afb03187 100644 --- a/google/cloud/dataproc_spark_connect/__init__.py +++ b/google/cloud/dataproc_spark_connect/__init__.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Deprecated: this package has been renamed to ``google.cloud.managed_spark_connect``.""" + import warnings from google.cloud.managed_spark_connect import ManagedSparkSession diff --git a/google/cloud/dataproc_spark_connect/client/__init__.py b/google/cloud/dataproc_spark_connect/client/__init__.py index da634080..40fabb0c 100644 --- a/google/cloud/dataproc_spark_connect/client/__init__.py +++ b/google/cloud/dataproc_spark_connect/client/__init__.py @@ -12,4 +12,5 @@ # See the License for the specific language governing permissions and # limitations under the License. """Deprecated: use ``google.cloud.managed_spark_connect.client`` instead.""" + from .core import DataprocChannelBuilder diff --git a/google/cloud/dataproc_spark_connect/client/core.py b/google/cloud/dataproc_spark_connect/client/core.py index 7fedc5fc..fa81a159 100644 --- a/google/cloud/dataproc_spark_connect/client/core.py +++ b/google/cloud/dataproc_spark_connect/client/core.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Deprecated: use ``google.cloud.managed_spark_connect.client.core`` instead.""" + from google.cloud.managed_spark_connect.client.core import ( ManagedSparkChannelBuilder, ProxiedChannel, diff --git a/google/cloud/dataproc_spark_connect/environment.py b/google/cloud/dataproc_spark_connect/environment.py index e04708cc..4bfc38b3 100644 --- a/google/cloud/dataproc_spark_connect/environment.py +++ b/google/cloud/dataproc_spark_connect/environment.py @@ -12,4 +12,5 @@ # See the License for the specific language governing permissions and # limitations under the License. """Deprecated: use ``google.cloud.managed_spark_connect.environment`` instead.""" + from google.cloud.managed_spark_connect.environment import * # noqa: F401,F403 diff --git a/google/cloud/dataproc_spark_connect/exceptions.py b/google/cloud/dataproc_spark_connect/exceptions.py index d958a7cb..a24b2ea7 100644 --- a/google/cloud/dataproc_spark_connect/exceptions.py +++ b/google/cloud/dataproc_spark_connect/exceptions.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Deprecated: use ``google.cloud.managed_spark_connect.exceptions`` instead.""" + from google.cloud.managed_spark_connect.exceptions import ( ManagedSparkConnectException, ) diff --git a/google/cloud/dataproc_spark_connect/session.py b/google/cloud/dataproc_spark_connect/session.py index 4e838d96..647a496a 100644 --- a/google/cloud/dataproc_spark_connect/session.py +++ b/google/cloud/dataproc_spark_connect/session.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Deprecated: use ``google.cloud.managed_spark_connect.session`` instead.""" + from google.cloud.managed_spark_connect.session import ( ManagedSparkSession, _is_valid_label_value, diff --git a/tests/integration/managed_spark_magics/test_magics.py b/tests/integration/managed_spark_magics/test_magics.py index 2b067178..b0c8f115 100644 --- a/tests/integration/managed_spark_magics/test_magics.py +++ b/tests/integration/managed_spark_magics/test_magics.py @@ -18,7 +18,6 @@ from google.cloud.managed_spark_connect import ManagedSparkSession - _SERVICE_ACCOUNT_KEY_FILE_ = "service_account_key.json" diff --git a/tests/integration/test_session.py b/tests/integration/test_session.py index 61d4a2f6..efea4c48 100644 --- a/tests/integration/test_session.py +++ b/tests/integration/test_session.py @@ -34,7 +34,6 @@ from pyspark.errors.exceptions import connect as connect_exceptions from pyspark.sql.types import StringType - _SERVICE_ACCOUNT_KEY_FILE_ = "service_account_key.json" diff --git a/tests/unit/test_deprecated_shims.py b/tests/unit/test_deprecated_shims.py index aca687cd..62be9611 100644 --- a/tests/unit/test_deprecated_shims.py +++ b/tests/unit/test_deprecated_shims.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Tests that the pre-rename `dataproc_*` import paths still work as deprecated aliases.""" + import importlib import sys import unittest From 4cb538027778c528f8fc7d7f4fa399d8513c2100 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:32:17 -0700 Subject: [PATCH 3/8] fix: address zizmor security findings in integration-tests.yaml Pins all action references to commit SHAs, adds an explicit contents: read permission, sets persist-credentials: false on checkout, and stops extracting secret fields via fromJson(secrets.*) in expressions (which bypasses GitHub's log redaction) in favor of a dedicated step that masks the derived values explicitly. --- .github/workflows/integration-tests.yaml | 32 ++++++++++++++++++------ 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/.github/workflows/integration-tests.yaml b/.github/workflows/integration-tests.yaml index 8f5b7627..e2573e7b 100644 --- a/.github/workflows/integration-tests.yaml +++ b/.github/workflows/integration-tests.yaml @@ -27,6 +27,9 @@ on: branches: [ main ] workflow_dispatch: +permissions: + contents: read + jobs: integration-test: name: Run integration tests @@ -37,15 +40,17 @@ jobs: steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 + with: + persist-credentials: false - name: Setup Python - uses: actions/setup-python@v5 + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5 with: python-version: "3.12" - name: Cache pip dependencies - uses: actions/cache@v4 + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4 with: path: ~/.cache/pip key: ${{ runner.os }}-pip-integration-${{ hashFiles('requirements-dev.txt', 'requirements-test.txt') }} @@ -59,19 +64,30 @@ jobs: pip install -r requirements-test.txt - name: Authenticate to Google Cloud - uses: google-github-actions/auth@v2 + uses: google-github-actions/auth@c200f3691d83b41bf9bbd8638997a462592937ed # v2 with: credentials_json: ${{ secrets.GCP_SA_KEY }} - name: Set up Cloud SDK - uses: google-github-actions/setup-gcloud@v2 + uses: google-github-actions/setup-gcloud@e427ad8a34f8676edf47cf7d7925499adf3eb74f # v2 + + - name: Extract service account details + env: + GCP_SA_KEY_JSON: ${{ secrets.GCP_SA_KEY }} + run: | + SA_EMAIL=$(echo "$GCP_SA_KEY_JSON" | jq -r '.client_email') + PROJECT_ID=$(echo "$GCP_SA_KEY_JSON" | jq -r '.project_id') + echo "::add-mask::$SA_EMAIL" + echo "::add-mask::$PROJECT_ID" + echo "SA_EMAIL=$SA_EMAIL" >> "$GITHUB_ENV" + echo "PROJECT_ID=$PROJECT_ID" >> "$GITHUB_ENV" - name: Run integration tests env: CI: "true" - # Extract from service account JSON automatically - GOOGLE_CLOUD_PROJECT: ${{ fromJson(secrets.GCP_SA_KEY).project_id }} - MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT: ${{ fromJson(secrets.GCP_SA_KEY).client_email }} + # Extracted from service account JSON in the previous step + GOOGLE_CLOUD_PROJECT: ${{ env.PROJECT_ID }} + MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT: ${{ env.SA_EMAIL }} # Infrastructure-specific secrets GOOGLE_CLOUD_REGION: ${{ secrets.GCP_REGION || 'us-central1' }} MANAGED_SPARK_CONNECT_SUBNET: ${{ secrets.GCP_SUBNET }} From 0d005abafc13a6b9e323648855b7d28a92419b0d Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Thu, 27 Aug 2026 10:10:28 -0700 Subject: [PATCH 4/8] fix: rename PyPI distribution to google-cloud-managed-spark-connect Aligns with Google's naming convention for other google-cloud-* client libraries rather than the shorter managed-spark-connect. --- README.md | 12 ++++++------ google/cloud/managed_spark_connect/__init__.py | 2 +- setup.py | 2 +- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 7d81c692..a7ec34fb 100644 --- a/README.md +++ b/README.md @@ -8,13 +8,13 @@ requiring additional steps. ## Install ```sh -pip install managed_spark_connect +pip install google-cloud-managed-spark-connect ``` ## Uninstall ```sh -pip uninstall managed_spark_connect +pip uninstall google-cloud-managed-spark-connect ``` ## Setup @@ -39,7 +39,7 @@ in your code using the builder API: 1. Install the latest version of Managed Spark Connect: ```sh - pip install -U managed-spark-connect + pip install -U google-cloud-managed-spark-connect ``` 2. Add the required imports into your PySpark application or notebook and start @@ -127,7 +127,7 @@ The package supports the [sparksql-magic](https://github.com/cryeo/sparksql-magi **Installation**: To use magic commands, install the required dependencies manually: ```bash -pip install managed-spark-connect +pip install google-cloud-managed-spark-connect pip install IPython sparksql-magic ``` @@ -165,12 +165,12 @@ See [sparksql-magic](https://github.com/cryeo/sparksql-magic) for more examples. **Note**: Magic commands are optional. If you only need basic ManagedSparkSession functionality without Jupyter magic support, install only the base package: ```bash -pip install managed-spark-connect +pip install google-cloud-managed-spark-connect ``` ## Migrating from dataproc-spark-connect -The `dataproc-spark-connect` package and the `google.cloud.dataproc_spark_connect` module have been renamed to `managed-spark-connect` / `google.cloud.managed_spark_connect`, and `DataprocSparkSession` has been renamed to `ManagedSparkSession`. The old import path and class name still work but emit a `DeprecationWarning` — update your imports when convenient: +The `dataproc-spark-connect` package and the `google.cloud.dataproc_spark_connect` module have been renamed to `google-cloud-managed-spark-connect` / `google.cloud.managed_spark_connect`, and `DataprocSparkSession` has been renamed to `ManagedSparkSession`. The old import path and class name still work but emit a `DeprecationWarning` — update your imports when convenient: ```python # Before diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index 70ff97c2..099d4f17 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -17,7 +17,7 @@ from .session import ManagedSparkSession old_package_names = ["google-spark-connect", "dataproc-spark-connect"] -current_package_name = "managed-spark-connect" +current_package_name = "google-cloud-managed-spark-connect" for old_package_name in old_package_names: try: importlib.metadata.distribution(old_package_name) diff --git a/setup.py b/setup.py index ca9d32d8..8978cf37 100644 --- a/setup.py +++ b/setup.py @@ -19,7 +19,7 @@ setup( - name="managed-spark-connect", + name="google-cloud-managed-spark-connect", version="1.1.0", description="Managed Spark client library for Spark Connect", long_description=long_description, From 33a3dda14c57e855a18268df50b64c64cac4fff2 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Thu, 27 Aug 2026 10:27:48 -0700 Subject: [PATCH 5/8] refactor: drop backward-compatible shims for the Dataproc -> Managed Spark rename Removes the google.cloud.dataproc_spark_connect / google.cloud.dataproc_magics shim packages, the DataprocSparkSession/DataprocMagics aliases, the sessionTemplate() deprecated alias for runtimeProfile(), and the DATAPROC_SPARK_CONNECT_* env var fallback. Anyone moving to the new google-cloud-managed-spark-connect package needs to update their imports, builder calls, and env vars directly rather than relying on a deprecation period. --- README.md | 9 +- google/cloud/dataproc_magics/__init__.py | 28 ------ google/cloud/dataproc_magics/magics.py | 18 ---- .../cloud/dataproc_spark_connect/__init__.py | 27 ------ .../dataproc_spark_connect/client/__init__.py | 16 ---- .../dataproc_spark_connect/client/core.py | 21 ----- .../dataproc_spark_connect/client/proxy.py | 13 --- .../dataproc_spark_connect/environment.py | 16 ---- .../dataproc_spark_connect/exceptions.py | 20 ---- .../dataproc_spark_connect/pypi_artifacts.py | 3 - .../cloud/dataproc_spark_connect/session.py | 22 ----- google/cloud/managed_spark_connect/session.py | 93 ++++--------------- tests/unit/test_deprecated_shims.py | 76 --------------- tests/unit/test_session.py | 78 ---------------- 14 files changed, 26 insertions(+), 414 deletions(-) delete mode 100644 google/cloud/dataproc_magics/__init__.py delete mode 100644 google/cloud/dataproc_magics/magics.py delete mode 100644 google/cloud/dataproc_spark_connect/__init__.py delete mode 100644 google/cloud/dataproc_spark_connect/client/__init__.py delete mode 100644 google/cloud/dataproc_spark_connect/client/core.py delete mode 100644 google/cloud/dataproc_spark_connect/client/proxy.py delete mode 100644 google/cloud/dataproc_spark_connect/environment.py delete mode 100644 google/cloud/dataproc_spark_connect/exceptions.py delete mode 100644 google/cloud/dataproc_spark_connect/pypi_artifacts.py delete mode 100644 google/cloud/dataproc_spark_connect/session.py delete mode 100644 tests/unit/test_deprecated_shims.py diff --git a/README.md b/README.md index a7ec34fb..c6ba9063 100644 --- a/README.md +++ b/README.md @@ -170,7 +170,14 @@ pip install google-cloud-managed-spark-connect ## Migrating from dataproc-spark-connect -The `dataproc-spark-connect` package and the `google.cloud.dataproc_spark_connect` module have been renamed to `google-cloud-managed-spark-connect` / `google.cloud.managed_spark_connect`, and `DataprocSparkSession` has been renamed to `ManagedSparkSession`. The old import path and class name still work but emit a `DeprecationWarning` — update your imports when convenient: +The `dataproc-spark-connect` package has been renamed to `google-cloud-managed-spark-connect`. This is a breaking change — update your code when you switch to the new package: + +* `pip install dataproc-spark-connect` → `pip install google-cloud-managed-spark-connect` +* `google.cloud.dataproc_spark_connect` → `google.cloud.managed_spark_connect` +* `DataprocSparkSession` → `ManagedSparkSession` +* `DataprocMagics` / `google.cloud.dataproc_magics` → `ManagedSparkMagics` / `google.cloud.managed_spark_magics` +* `.sessionTemplate(...)` builder method → `.runtimeProfile(...)` +* `DATAPROC_SPARK_CONNECT_*` environment variables → `MANAGED_SPARK_CONNECT_*` ```python # Before diff --git a/google/cloud/dataproc_magics/__init__.py b/google/cloud/dataproc_magics/__init__.py deleted file mode 100644 index 6cce7b3d..00000000 --- a/google/cloud/dataproc_magics/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# 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. -"""Deprecated: this package has been renamed to ``google.cloud.managed_spark_magics``.""" - -import warnings - -from .magics import DataprocMagics - -warnings.warn( - "google.cloud.dataproc_magics is deprecated, use google.cloud.managed_spark_magics instead.", - DeprecationWarning, - stacklevel=2, -) - - -def load_ipython_extension(ipython): - ipython.register_magics(DataprocMagics) diff --git a/google/cloud/dataproc_magics/magics.py b/google/cloud/dataproc_magics/magics.py deleted file mode 100644 index b243be84..00000000 --- a/google/cloud/dataproc_magics/magics.py +++ /dev/null @@ -1,18 +0,0 @@ -# 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. -"""Deprecated: use ``google.cloud.managed_spark_magics.magics`` instead.""" - -from google.cloud.managed_spark_magics.magics import ManagedSparkMagics - -DataprocMagics = ManagedSparkMagics diff --git a/google/cloud/dataproc_spark_connect/__init__.py b/google/cloud/dataproc_spark_connect/__init__.py deleted file mode 100644 index afb03187..00000000 --- a/google/cloud/dataproc_spark_connect/__init__.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Deprecated: this package has been renamed to ``google.cloud.managed_spark_connect``.""" - -import warnings - -from google.cloud.managed_spark_connect import ManagedSparkSession - -DataprocSparkSession = ManagedSparkSession - -warnings.warn( - "google.cloud.dataproc_spark_connect is deprecated, use google.cloud.managed_spark_connect instead. " - "DataprocSparkSession has been renamed to ManagedSparkSession.", - DeprecationWarning, - stacklevel=2, -) diff --git a/google/cloud/dataproc_spark_connect/client/__init__.py b/google/cloud/dataproc_spark_connect/client/__init__.py deleted file mode 100644 index 40fabb0c..00000000 --- a/google/cloud/dataproc_spark_connect/client/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Deprecated: use ``google.cloud.managed_spark_connect.client`` instead.""" - -from .core import DataprocChannelBuilder diff --git a/google/cloud/dataproc_spark_connect/client/core.py b/google/cloud/dataproc_spark_connect/client/core.py deleted file mode 100644 index fa81a159..00000000 --- a/google/cloud/dataproc_spark_connect/client/core.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Deprecated: use ``google.cloud.managed_spark_connect.client.core`` instead.""" - -from google.cloud.managed_spark_connect.client.core import ( - ManagedSparkChannelBuilder, - ProxiedChannel, -) - -DataprocChannelBuilder = ManagedSparkChannelBuilder diff --git a/google/cloud/dataproc_spark_connect/client/proxy.py b/google/cloud/dataproc_spark_connect/client/proxy.py deleted file mode 100644 index 5449e384..00000000 --- a/google/cloud/dataproc_spark_connect/client/proxy.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Deprecated: use ``google.cloud.managed_spark_connect.client.proxy`` instead.""" - -from google.cloud.managed_spark_connect.client.proxy import ( - ManagedSparkSessionProxy, - connect_sockets, - connect_tcp_bridge, - forward_bytes, - forward_connection, - managed_spark_session_proxy, -) - -DataprocSessionProxy = ManagedSparkSessionProxy -dataproc_session_proxy = managed_spark_session_proxy diff --git a/google/cloud/dataproc_spark_connect/environment.py b/google/cloud/dataproc_spark_connect/environment.py deleted file mode 100644 index 4bfc38b3..00000000 --- a/google/cloud/dataproc_spark_connect/environment.py +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright 2025 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 -# -# https://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. -"""Deprecated: use ``google.cloud.managed_spark_connect.environment`` instead.""" - -from google.cloud.managed_spark_connect.environment import * # noqa: F401,F403 diff --git a/google/cloud/dataproc_spark_connect/exceptions.py b/google/cloud/dataproc_spark_connect/exceptions.py deleted file mode 100644 index a24b2ea7..00000000 --- a/google/cloud/dataproc_spark_connect/exceptions.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright 2025 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. -"""Deprecated: use ``google.cloud.managed_spark_connect.exceptions`` instead.""" - -from google.cloud.managed_spark_connect.exceptions import ( - ManagedSparkConnectException, -) - -DataprocSparkConnectException = ManagedSparkConnectException diff --git a/google/cloud/dataproc_spark_connect/pypi_artifacts.py b/google/cloud/dataproc_spark_connect/pypi_artifacts.py deleted file mode 100644 index f6f27812..00000000 --- a/google/cloud/dataproc_spark_connect/pypi_artifacts.py +++ /dev/null @@ -1,3 +0,0 @@ -"""Deprecated: use ``google.cloud.managed_spark_connect.pypi_artifacts`` instead.""" - -from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts diff --git a/google/cloud/dataproc_spark_connect/session.py b/google/cloud/dataproc_spark_connect/session.py deleted file mode 100644 index 647a496a..00000000 --- a/google/cloud/dataproc_spark_connect/session.py +++ /dev/null @@ -1,22 +0,0 @@ -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Deprecated: use ``google.cloud.managed_spark_connect.session`` instead.""" - -from google.cloud.managed_spark_connect.session import ( - ManagedSparkSession, - _is_valid_label_value, - _is_valid_session_id, -) - -DataprocSparkSession = ManagedSparkSession diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index 12952351..4425b240 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -24,7 +24,6 @@ import threading import time import uuid -import warnings import tqdm from packaging import version from types import MethodType @@ -73,25 +72,6 @@ ) -def _env_var_set(new_name: str, old_name: str) -> bool: - return new_name in os.environ or old_name in os.environ - - -def _getenv_with_deprecated_alias( - new_name: str, old_name: str, default: Optional[str] = None -) -> Optional[str]: - if new_name in os.environ: - return os.environ[new_name] - if old_name in os.environ: - warnings.warn( - f"Environment variable '{old_name}' is deprecated, use '{new_name}' instead.", - DeprecationWarning, - stacklevel=2, - ) - return os.environ[old_name] - return default - - def _is_valid_label_value(value: str) -> bool: """ Validates if a string complies with Google Cloud label value format. @@ -282,15 +262,6 @@ def runtimeProfile(self, profile: str): self.dataproc_config.session_template = profile return self - def sessionTemplate(self, template: str): - """Deprecated: use :meth:`runtimeProfile` instead.""" - warnings.warn( - "sessionTemplate() is deprecated, use runtimeProfile() instead.", - DeprecationWarning, - stacklevel=2, - ) - return self.runtimeProfile(template) - def label(self, key: str, value: str): """Add a single label to the session.""" return self.labels({key: value}) @@ -426,9 +397,8 @@ def create_session_pbar(): try: if ( - _getenv_with_deprecated_alias( + os.getenv( "MANAGED_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT", - "DATAPROC_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT", "false", ) == "true" @@ -696,13 +666,12 @@ def _get_dataproc_config(self): exec_config = dataproc_config.environment_config.execution_config # Set service account from environment if not already set - if not exec_config.service_account and _env_var_set( - "MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT", - "DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT", + if ( + not exec_config.service_account + and "MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT" in os.environ ): - exec_config.service_account = _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT", - "DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT", + exec_config.service_account = os.getenv( + "MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT" ) # Auto-set authentication type to SERVICE_ACCOUNT when service account is provided @@ -713,57 +682,35 @@ def _get_dataproc_config(self): ) elif ( not exec_config.authentication_config.user_workload_authentication_type - and _env_var_set( - "MANAGED_SPARK_CONNECT_AUTH_TYPE", - "DATAPROC_SPARK_CONNECT_AUTH_TYPE", - ) + and "MANAGED_SPARK_CONNECT_AUTH_TYPE" in os.environ ): # Only set auth type from environment if no service account is present exec_config.authentication_config.user_workload_authentication_type = AuthenticationConfig.AuthenticationType[ - _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_AUTH_TYPE", - "DATAPROC_SPARK_CONNECT_AUTH_TYPE", - ) + os.getenv("MANAGED_SPARK_CONNECT_AUTH_TYPE") ] if ( not dataproc_config.environment_config.execution_config.subnetwork_uri - and _env_var_set( - "MANAGED_SPARK_CONNECT_SUBNET", - "DATAPROC_SPARK_CONNECT_SUBNET", - ) + and "MANAGED_SPARK_CONNECT_SUBNET" in os.environ ): - dataproc_config.environment_config.execution_config.subnetwork_uri = _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_SUBNET", - "DATAPROC_SPARK_CONNECT_SUBNET", + dataproc_config.environment_config.execution_config.subnetwork_uri = os.getenv( + "MANAGED_SPARK_CONNECT_SUBNET" ) if ( not dataproc_config.environment_config.execution_config.ttl - and _env_var_set( - "MANAGED_SPARK_CONNECT_TTL_SECONDS", - "DATAPROC_SPARK_CONNECT_TTL_SECONDS", - ) + and "MANAGED_SPARK_CONNECT_TTL_SECONDS" in os.environ ): dataproc_config.environment_config.execution_config.ttl = { "seconds": int( - _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_TTL_SECONDS", - "DATAPROC_SPARK_CONNECT_TTL_SECONDS", - ) + os.getenv("MANAGED_SPARK_CONNECT_TTL_SECONDS") ) } if ( not dataproc_config.environment_config.execution_config.idle_ttl - and _env_var_set( - "MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS", - "DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS", - ) + and "MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS" in os.environ ): dataproc_config.environment_config.execution_config.idle_ttl = { "seconds": int( - _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS", - "DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS", - ) + os.getenv("MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS") ) } client_environment = environment.get_client_environment_label() @@ -786,9 +733,8 @@ def _get_dataproc_config(self): f"Maximum length is 63 characters. " f"Ignoring notebook ID label." ) - default_datasource = _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE", - "DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE", + default_datasource = os.getenv( + "MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE" ) match default_datasource: case "bigquery": @@ -1280,10 +1226,7 @@ def addArtifacts( @staticmethod def _get_active_session_file_path(): - return _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH", - "DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH", - ) + return os.getenv("MANAGED_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH") def stop(self, terminate: Optional[bool] = None) -> None: """ diff --git a/tests/unit/test_deprecated_shims.py b/tests/unit/test_deprecated_shims.py deleted file mode 100644 index 62be9611..00000000 --- a/tests/unit/test_deprecated_shims.py +++ /dev/null @@ -1,76 +0,0 @@ -# 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. -"""Tests that the pre-rename `dataproc_*` import paths still work as deprecated aliases.""" - -import importlib -import sys -import unittest - - -def _fresh_import(module_name): - """Import (or re-import) a module, forcing its top-level code to run. - - Needed because a module already cached in sys.modules from an earlier - test or import elsewhere would otherwise not re-emit its deprecation - warning, making assertWarns order-dependent. - """ - for name in list(sys.modules): - if name == module_name or name.startswith(module_name + "."): - del sys.modules[name] - return importlib.import_module(module_name) - - -class DeprecatedPackageShimTests(unittest.TestCase): - - def test_dataproc_spark_connect_package_warns_and_aliases_session(self): - from google.cloud.managed_spark_connect import ManagedSparkSession - - with self.assertWarns(DeprecationWarning): - module = _fresh_import("google.cloud.dataproc_spark_connect") - - self.assertIs(module.DataprocSparkSession, ManagedSparkSession) - - def test_dataproc_spark_connect_exceptions_alias(self): - from google.cloud.managed_spark_connect.exceptions import ( - ManagedSparkConnectException, - ) - from google.cloud.dataproc_spark_connect.exceptions import ( - DataprocSparkConnectException, - ) - - self.assertIs( - DataprocSparkConnectException, ManagedSparkConnectException - ) - - def test_dataproc_spark_connect_client_alias(self): - from google.cloud.managed_spark_connect.client import ( - ManagedSparkChannelBuilder, - ) - from google.cloud.dataproc_spark_connect.client import ( - DataprocChannelBuilder, - ) - - self.assertIs(DataprocChannelBuilder, ManagedSparkChannelBuilder) - - def test_dataproc_magics_package_warns_and_aliases_magics(self): - from google.cloud.managed_spark_magics import ManagedSparkMagics - - with self.assertWarns(DeprecationWarning): - module = _fresh_import("google.cloud.dataproc_magics") - - self.assertIs(module.DataprocMagics, ManagedSparkMagics) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 624a8ca4..02b05b22 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -25,8 +25,6 @@ from google.cloud.managed_spark_connect import ManagedSparkSession from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException from google.cloud.managed_spark_connect.session import ( - _env_var_set, - _getenv_with_deprecated_alias, _is_valid_label_value, _is_valid_session_id, ) @@ -2140,18 +2138,6 @@ def test_builder_pattern_runtime_profile_and_labels( ) self.stopSession(mock_session_controller_client_instance, session) - def test_session_template_is_deprecated_alias_for_runtime_profile(self): - """sessionTemplate() should still work but warn in favor of runtimeProfile().""" - builder = ManagedSparkSession.Builder() - with self.assertWarns(DeprecationWarning): - builder.sessionTemplate( - "projects/test-project/locations/us-central1/sessionTemplates/test-template" - ) - self.assertEqual( - builder.dataproc_config.session_template, - "projects/test-project/locations/us-central1/sessionTemplates/test-template", - ) - @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") @@ -2659,69 +2645,5 @@ def test_session_skip_terminated(self, mock_session_controller_client): mock_client.get_session.assert_called_once() -class DeprecatedEnvVarAliasTests(unittest.TestCase): - """Test cases for the MANAGED_SPARK_CONNECT_* / DATAPROC_SPARK_CONNECT_* env var fallback.""" - - def setUp(self): - for name in ( - "MANAGED_SPARK_CONNECT_TEST_VAR", - "DATAPROC_SPARK_CONNECT_TEST_VAR", - ): - os.environ.pop(name, None) - - tearDown = setUp - - def test_new_name_takes_precedence(self): - os.environ["MANAGED_SPARK_CONNECT_TEST_VAR"] = "new" - os.environ["DATAPROC_SPARK_CONNECT_TEST_VAR"] = "old" - self.assertEqual( - _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_TEST_VAR", - "DATAPROC_SPARK_CONNECT_TEST_VAR", - ), - "new", - ) - - def test_old_name_used_with_deprecation_warning(self): - os.environ["DATAPROC_SPARK_CONNECT_TEST_VAR"] = "old" - with self.assertWarns(DeprecationWarning): - value = _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_TEST_VAR", - "DATAPROC_SPARK_CONNECT_TEST_VAR", - ) - self.assertEqual(value, "old") - - def test_default_when_neither_set(self): - self.assertIsNone( - _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_TEST_VAR", - "DATAPROC_SPARK_CONNECT_TEST_VAR", - ) - ) - self.assertEqual( - _getenv_with_deprecated_alias( - "MANAGED_SPARK_CONNECT_TEST_VAR", - "DATAPROC_SPARK_CONNECT_TEST_VAR", - "fallback", - ), - "fallback", - ) - - def test_env_var_set_checks_both_names(self): - self.assertFalse( - _env_var_set( - "MANAGED_SPARK_CONNECT_TEST_VAR", - "DATAPROC_SPARK_CONNECT_TEST_VAR", - ) - ) - os.environ["DATAPROC_SPARK_CONNECT_TEST_VAR"] = "old" - self.assertTrue( - _env_var_set( - "MANAGED_SPARK_CONNECT_TEST_VAR", - "DATAPROC_SPARK_CONNECT_TEST_VAR", - ) - ) - - if __name__ == "__main__": unittest.main() From c15c6c13b2f9a306f892e3670c598baed6ba9f11 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Thu, 27 Aug 2026 10:30:15 -0700 Subject: [PATCH 6/8] docs: expand migration guide with concrete before/after steps Spells out each breaking change individually (package name, imports, sessionTemplate() -> runtimeProfile(), env var prefix) with code examples instead of a flat bullet list, since sessionTemplate -> runtimeProfile in particular isn't obvious from a name-only diff. --- README.md | 62 ++++++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 55 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index c6ba9063..fa5003f6 100644 --- a/README.md +++ b/README.md @@ -170,23 +170,71 @@ pip install google-cloud-managed-spark-connect ## Migrating from dataproc-spark-connect -The `dataproc-spark-connect` package has been renamed to `google-cloud-managed-spark-connect`. This is a breaking change — update your code when you switch to the new package: +The `dataproc-spark-connect` package has been renamed to `google-cloud-managed-spark-connect`. This is a breaking change with no compatibility shims — you need to update your code in the following places when you switch to the new package. -* `pip install dataproc-spark-connect` → `pip install google-cloud-managed-spark-connect` -* `google.cloud.dataproc_spark_connect` → `google.cloud.managed_spark_connect` -* `DataprocSparkSession` → `ManagedSparkSession` -* `DataprocMagics` / `google.cloud.dataproc_magics` → `ManagedSparkMagics` / `google.cloud.managed_spark_magics` -* `.sessionTemplate(...)` builder method → `.runtimeProfile(...)` -* `DATAPROC_SPARK_CONNECT_*` environment variables → `MANAGED_SPARK_CONNECT_*` +### 1. Update the package you install + +```sh +# Before +pip install dataproc-spark-connect + +# After +pip install google-cloud-managed-spark-connect +``` + +### 2. Update your imports and session class + +`google.cloud.dataproc_spark_connect` is now `google.cloud.managed_spark_connect`, and `DataprocSparkSession` is now `ManagedSparkSession`: ```python # Before from google.cloud.dataproc_spark_connect import DataprocSparkSession +spark = DataprocSparkSession.builder.getOrCreate() # After from google.cloud.managed_spark_connect import ManagedSparkSession +spark = ManagedSparkSession.builder.getOrCreate() +``` + +If you use the Jupyter magic commands, `google.cloud.dataproc_magics` is now `google.cloud.managed_spark_magics` and `DataprocMagics` is now `ManagedSparkMagics` (the `%dpip` magic itself is unchanged). + +### 3. Rename `sessionTemplate(...)` calls to `runtimeProfile(...)` + +The builder method used to configure a session template is renamed from `sessionTemplate()` to `runtimeProfile()`. It takes the same argument (the full resource name of the template) and behaves identically — only the method name changes: + +```python +# Before +spark = ( + DataprocSparkSession.builder + .sessionTemplate("projects/my-project/locations/us-central1/sessionTemplates/my-template") + .getOrCreate() +) + +# After +spark = ( + ManagedSparkSession.builder + .runtimeProfile("projects/my-project/locations/us-central1/sessionTemplates/my-template") + .getOrCreate() +) ``` +### 4. Rename any `DATAPROC_SPARK_CONNECT_*` environment variables + +If you set any of the library's own environment variables (as opposed to standard GCP ones like `GOOGLE_CLOUD_PROJECT`), rename the `DATAPROC_SPARK_CONNECT_` prefix to `MANAGED_SPARK_CONNECT_`: + +| Before | After | +|--------|-------| +| `DATAPROC_SPARK_CONNECT_SERVICE_ACCOUNT` | `MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT` | +| `DATAPROC_SPARK_CONNECT_SUBNET` | `MANAGED_SPARK_CONNECT_SUBNET` | +| `DATAPROC_SPARK_CONNECT_AUTH_TYPE` | `MANAGED_SPARK_CONNECT_AUTH_TYPE` | +| `DATAPROC_SPARK_CONNECT_TTL_SECONDS` | `MANAGED_SPARK_CONNECT_TTL_SECONDS` | +| `DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS` | `MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS` | +| `DATAPROC_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT` | `MANAGED_SPARK_CONNECT_SESSION_TERMINATE_AT_EXIT` | +| `DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE` | `MANAGED_SPARK_CONNECT_DEFAULT_DATASOURCE` | +| `DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH` | `MANAGED_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH` | + +Note that `GOOGLE_CLOUD_DATAPROC_API_ENDPOINT` and other variables naming the actual Dataproc API (not this library's own config) are unchanged. + ## Developing For development instructions see [guide](DEVELOPING.md). From 21f244452b540d8cc82d0b861e78de4f1029f2c1 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Thu, 27 Aug 2026 10:50:01 -0700 Subject: [PATCH 7/8] revert: keep sessionTemplate() builder method name Undo the runtimeProfile() rename; the method still sets Session.session_template but keeps its original name. --- README.md | 24 ++----------------- google/cloud/managed_spark_connect/session.py | 4 ++-- tests/unit/test_session.py | 4 ++-- 3 files changed, 6 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index fa5003f6..fc4b25f3 100644 --- a/README.md +++ b/README.md @@ -83,9 +83,9 @@ The `ManagedSparkSession.builder` provides a fluent API to configure the session | `labels(labels)` | Adds multiple labels to the session. | | `location(location)` | Sets the Google Cloud region. | | `projectId(project_id)` | Sets the Google Cloud project ID. | -| `runtimeProfile(profile)` | Sets the Runtime Profile to use. | | `runtimeVersion(version)` | Sets the Managed Spark runtime version (e.g., "3.0"). | | `serviceAccount(account)` | Sets the service account for the session. | +| `sessionTemplate(profile)` | Sets the Session Template to use. | | `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. | @@ -198,27 +198,7 @@ spark = ManagedSparkSession.builder.getOrCreate() If you use the Jupyter magic commands, `google.cloud.dataproc_magics` is now `google.cloud.managed_spark_magics` and `DataprocMagics` is now `ManagedSparkMagics` (the `%dpip` magic itself is unchanged). -### 3. Rename `sessionTemplate(...)` calls to `runtimeProfile(...)` - -The builder method used to configure a session template is renamed from `sessionTemplate()` to `runtimeProfile()`. It takes the same argument (the full resource name of the template) and behaves identically — only the method name changes: - -```python -# Before -spark = ( - DataprocSparkSession.builder - .sessionTemplate("projects/my-project/locations/us-central1/sessionTemplates/my-template") - .getOrCreate() -) - -# After -spark = ( - ManagedSparkSession.builder - .runtimeProfile("projects/my-project/locations/us-central1/sessionTemplates/my-template") - .getOrCreate() -) -``` - -### 4. Rename any `DATAPROC_SPARK_CONNECT_*` environment variables +### 3. Rename any `DATAPROC_SPARK_CONNECT_*` environment variables If you set any of the library's own environment variables (as opposed to standard GCP ones like `GOOGLE_CLOUD_PROJECT`), rename the `DATAPROC_SPARK_CONNECT_` prefix to `MANAGED_SPARK_CONNECT_`: diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index 4425b240..9df7ca05 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -257,8 +257,8 @@ def idleTtlSeconds(self, seconds: int): } return self - def runtimeProfile(self, profile: str): - """Set the Runtime Profile to use for the session.""" + def sessionTemplate(self, profile: str): + """Set the Session Template to use for the session.""" self.dataproc_config.session_template = profile return self diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 02b05b22..e7b35198 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -2077,7 +2077,7 @@ def test_builder_pattern_ttl_with_timedelta( @mock.patch( "google.cloud.managed_spark_connect.session.is_s8s_session_active" ) - def test_builder_pattern_runtime_profile_and_labels( + def test_builder_pattern_session_template_and_labels( self, mock_is_s8s_session_active, mock_session_id, @@ -2098,7 +2098,7 @@ def test_builder_pattern_runtime_profile_and_labels( try: session = ( - ManagedSparkSession.builder.runtimeProfile( + ManagedSparkSession.builder.sessionTemplate( "projects/test-project/locations/us-central1/sessionTemplates/test-template" ) .label("environment", "production") From 7a3acf9ec95cb2b0d139ec6678aef040df04ed3a Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Mon, 31 Aug 2026 09:39:48 -0700 Subject: [PATCH 8/8] docs: revert Runtime Profile terminology to Session Template Follow-up to the sessionTemplate() method revert -- the remaining prose in the README and test docstrings/comments still said "Runtime Profile", which no longer matched the actual API surface. --- README.md | 2 +- tests/integration/test_session.py | 2 +- tests/unit/test_session.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index fc4b25f3..a320c653 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ pip uninstall google-cloud-managed-spark-connect ## Setup This client requires permissions to -manage [Managed Spark Sessions and Runtime Profiles](https://cloud.google.com/dataproc-serverless/docs/concepts/iam). +manage [Managed Spark Sessions and Session Templates](https://cloud.google.com/dataproc-serverless/docs/concepts/iam). If you are running the client outside of Google Cloud, you need to provide authentication credentials. Set the `GOOGLE_APPLICATION_CREDENTIALS` environment diff --git a/tests/integration/test_session.py b/tests/integration/test_session.py index efea4c48..676b3a2f 100644 --- a/tests/integration/test_session.py +++ b/tests/integration/test_session.py @@ -332,7 +332,7 @@ def test_create_spark_session_with_session_template_and_user_provided_dataproc_c session_template_name, session_controller_client, ): - """Test creating a Spark session with a Runtime Profile and user-provided Dataproc configuration.""" + """Test creating a Spark session with a session template and user-provided Dataproc configuration.""" dataproc_config = Session() dataproc_config.environment_config.execution_config.ttl = {"seconds": 64800} dataproc_config.session_template = session_template_name diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index e7b35198..985fce25 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -2107,7 +2107,7 @@ def test_builder_pattern_session_template_and_labels( .getOrCreate() ) - # Verify the session was created with the correct Runtime Profile and labels + # Verify the session was created with the correct session template and labels create_session_request = mock_session_controller_client_instance.create_session.call_args[ 0 ][