dataproc-spark-connect 0.7.5__tar.gz → 0.8.1__tar.gz

{dataproc_spark_connect-0.7.5 → dataproc_spark_connect-0.8.1}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: dataproc-spark-connect
-Version: 0.7.5
+Version: 0.8.1
 Summary: Dataproc client library for Spark Connect
 Home-page: https://github.com/GoogleCloudDataproc/dataproc-spark-connect-python
 Author: Google LLC
@@ -12,6 +12,13 @@ Requires-Dist: packaging>=20.0
 Requires-Dist: pyspark[connect]~=3.5.1
 Requires-Dist: tqdm>=4.67
 Requires-Dist: websockets>=14.0
+Dynamic: author
+Dynamic: description
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary
 
 # Dataproc Spark Connect Client
 

{dataproc_spark_connect-0.7.5 → dataproc_spark_connect-0.8.1}/dataproc_spark_connect.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: dataproc-spark-connect
-Version: 0.7.5
+Version: 0.8.1
 Summary: Dataproc client library for Spark Connect
 Home-page: https://github.com/GoogleCloudDataproc/dataproc-spark-connect-python
 Author: Google LLC
@@ -12,6 +12,13 @@ Requires-Dist: packaging>=20.0
 Requires-Dist: pyspark[connect]~=3.5.1
 Requires-Dist: tqdm>=4.67
 Requires-Dist: websockets>=14.0
+Dynamic: author
+Dynamic: description
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary
 
 # Dataproc Spark Connect Client
 

{dataproc_spark_connect-0.7.5 → dataproc_spark_connect-0.8.1}/google/cloud/dataproc_spark_connect/exceptions.py

@@ -16,7 +16,7 @@
 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.h
+    doesn't provide any additional information.
     """
 
     def __init__(self, message):

{dataproc_spark_connect-0.7.5 → dataproc_spark_connect-0.8.1}/google/cloud/dataproc_spark_connect/session.py

@@ -18,14 +18,23 @@ import json
 import logging
 import os
 import random
+import re
 import string
 import threading
 import time
+from typing import Any, cast, ClassVar, Dict, Optional, Union
+import uuid
 import tqdm
 
 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.exceptions import (
+    Aborted,
+    FailedPrecondition,
+    InvalidArgument,
+    NotFound,
+    PermissionDenied,
+)
 from google.api_core.future.polling import POLLING_PREDICATE
 from google.cloud.dataproc_spark_connect.client import DataprocChannelBuilder
 from google.cloud.dataproc_spark_connect.exceptions import DataprocSparkConnectException
@@ -41,13 +50,32 @@ from google.cloud.dataproc_v1 import (
 from google.cloud.dataproc_v1.types import sessions
 from pyspark.sql.connect.session import SparkSession
 from pyspark.sql.utils import to_str
-from typing import Any, cast, ClassVar, Dict, Optional
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
+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))
+
+
 class DataprocSparkSession(SparkSession):
     """The entry point to programming Spark with the Dataset and DataFrame API.
 
@@ -77,16 +105,6 @@ class DataprocSparkSession(SparkSession):
 
     class Builder(SparkSession.Builder):
 
-        _session_static_configs = [
-            "spark.executor.cores",
-            "spark.executor.memoryOverhead",
-            "spark.executor.memory",
-            "spark.driver.memory",
-            "spark.driver.cores",
-            "spark.eventLog.dir",
-            "spark.history.fs.logDirectory",
-        ]
-
         def __init__(self):
             self._options: Dict[str, Any] = {}
             self._channel_builder: Optional[DataprocChannelBuilder] = None
@@ -100,15 +118,6 @@ class DataprocSparkSession(SparkSession):
                 )
             )
 
-        def __apply_options(self, session: "SparkSession") -> None:
-            with self._lock:
-                self._options = {
-                    key: value
-                    for key, value in self._options.items()
-                    if key not in self._session_static_configs
-                }
-                self._apply_options(session)
-
         def projectId(self, project_id):
             self._project_id = project_id
             return self
@@ -166,7 +175,6 @@ class DataprocSparkSession(SparkSession):
             session = DataprocSparkSession(connection=self._channel_builder)
 
             DataprocSparkSession._set_default_and_active_session(session)
-            self.__apply_options(session)
             return session
 
         def __create(self) -> "DataprocSparkSession":
@@ -245,9 +253,9 @@ class DataprocSparkSession(SparkSession):
                     operation = SessionControllerClient(
                         client_options=self._client_options
                     ).create_session(session_request)
-                    print(
-                        f"Creating Dataproc Session: https://console.cloud.google.com/dataproc/interactive/{self._region}/{session_id}?project={self._project_id}"
-                    )
+                    self._display_session_link_on_creation(session_id)
+                    # TODO: Add the 'View Session Details' button once the UI changes are done.
+                    # self._display_view_session_details_button(session_id)
                     create_session_pbar_thread.start()
                     session_response: Session = operation.result(
                         polling=retry.Retry(
@@ -260,7 +268,7 @@ class DataprocSparkSession(SparkSession):
                     )
                     stop_create_session_pbar_event.set()
                     create_session_pbar_thread.join()
-                    print("Dataproc Session was successfully created")
+                    self._print_session_created_message()
                     file_path = (
                         DataprocSparkSession._get_active_session_file_path()
                     )
@@ -305,6 +313,46 @@ class DataprocSparkSession(SparkSession):
                     session_response, dataproc_config.name
                 )
 
+        def _display_session_link_on_creation(self, session_id):
+            session_url = f"https://console.cloud.google.com/dataproc/interactive/{self._region}/{session_id}?project={self._project_id}"
+            plain_message = f"Creating Dataproc Session: {session_url}"
+            html_element = f"""
+                <div>
+                    <p>Creating Dataproc Spark Session<p>
+                    <p><a href="{session_url}">Dataproc Session</a></p>
+                </div>
+            """
+
+            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"<div><p>{plain_message}</p></div>"
+
+            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
+            """
+            try:
+                from IPython.display import display, HTML
+                from IPython.core.interactiveshell import InteractiveShell
+
+                if not InteractiveShell.initialized():
+                    raise DataprocSparkConnectException(
+                        "Not in an Interactive IPython Environment"
+                    )
+                display(HTML(html_element))
+            except (ImportError, DataprocSparkConnectException):
+                print(plain_message)
+
         def _get_exiting_active_session(
             self,
         ) -> Optional["DataprocSparkSession"]:
@@ -325,6 +373,8 @@ class DataprocSparkSession(SparkSession):
                 print(
                     f"Using existing Dataproc Session (configuration changes may not be applied): https://console.cloud.google.com/dataproc/interactive/{self._region}/{s8s_session_id}?project={self._project_id}"
                 )
+                # TODO: Add the 'View Session Details' button once the UI changes are done.
+                # self._display_view_session_details_button(s8s_session_id)
                 if session is None:
                     session = self.__create_spark_connect_session_from_s8s(
                         session_response, session_name
@@ -344,8 +394,6 @@ class DataprocSparkSession(SparkSession):
                 session = self._get_exiting_active_session()
                 if session is None:
                     session = self.__create()
-                if session:
-                    self.__apply_options(session)
                 return session
 
         def _get_dataproc_config(self):
@@ -400,14 +448,22 @@ class DataprocSparkSession(SparkSession):
                         os.getenv("DATAPROC_SPARK_CONNECT_IDLE_TTL_SECONDS")
                     )
                 }
-            if "COLAB_NOTEBOOK_RUNTIME_ID" in os.environ:
-                dataproc_config.labels["colab-notebook-runtime-id"] = (
-                    os.environ["COLAB_NOTEBOOK_RUNTIME_ID"]
-                )
-            if "COLAB_NOTEBOOK_KERNEL_ID" in os.environ:
-                dataproc_config.labels["colab-notebook-kernel-id"] = os.environ[
-                    "COLAB_NOTEBOOK_KERNEL_ID"
-                ]
+            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"Skipping notebook ID label."
+                    )
             default_datasource = os.getenv(
                 "DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE"
             )
@@ -434,6 +490,22 @@ class DataprocSparkSession(SparkSession):
                     )
             return dataproc_config
 
+        def _display_view_session_details_button(self, session_id):
+            try:
+                session_url = f"https://console.cloud.google.com/dataproc/interactive/sessions/{session_id}/locations/{self._region}?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}")
+
         @staticmethod
         def generate_dataproc_session_id():
             timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
@@ -445,6 +517,41 @@ class DataprocSparkSession(SparkSession):
             )
             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)
+
+        base_method = self.client._execute_plan_request_with_metadata
+
+        def wrapped_method(*args, **kwargs):
+            req = 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}"
+                )
+            self._display_operation_link(req.operation_id)
+            return req
+
+        self.client._execute_plan_request_with_metadata = wrapped_method
+
     def _repr_html_(self) -> str:
         if not self._active_s8s_session_id:
             return """
@@ -462,6 +569,37 @@ class DataprocSparkSession(SparkSession):
         </div>
         """
 
+    def _display_operation_link(self, operation_id: str):
+        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"https://console.cloud.google.com/dataproc/interactive/{self._region}/"
+            f"{self._active_s8s_session_id}/sparkApplications/application;"
+            f"associatedSqlOperationId={operation_id}?project={self._project_id}"
+        )
+
+        try:
+            from IPython.display import display, HTML
+            from IPython.core.interactiveshell import InteractiveShell
+
+            if not InteractiveShell.initialized():
+                return
+            html_element = f"""
+              <div>
+                  <p><a href="{url}">Spark UI</a> (Operation: {operation_id})</p>
+              </div>
+              """
+            display(HTML(html_element))
+        except ImportError:
+            return
+
     @staticmethod
     def _remove_stopped_session_from_file():
         file_path = DataprocSparkSession._get_active_session_file_path()

{dataproc_spark_connect-0.7.5 → dataproc_spark_connect-0.8.1}/setup.py

@@ -20,7 +20,7 @@ long_description = (this_directory / "README.md").read_text()
 
 setup(
     name="dataproc-spark-connect",
-    version="0.7.5",
+    version="0.8.1",
     description="Dataproc client library for Spark Connect",
     long_description=long_description,
     author="Google LLC",