dataproc-spark-connect 0.8.0__py2.py3-none-any.whl → 0.8.2__py2.py3-none-any.whl

{dataproc_spark_connect-0.8.0.dist-info → dataproc_spark_connect-0.8.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataproc-spark-connect
-Version: 0.8.0
+Version: 0.8.2
 Summary: Dataproc client library for Spark Connect
 Home-page: https://github.com/GoogleCloudDataproc/dataproc-spark-connect-python
 Author: Google LLC

{dataproc_spark_connect-0.8.0.dist-info → dataproc_spark_connect-0.8.2.dist-info}/RECORD

@@ -1,12 +1,12 @@
-dataproc_spark_connect-0.8.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+dataproc_spark_connect-0.8.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 google/cloud/dataproc_spark_connect/__init__.py,sha256=dIqHNWVWWrSuRf26x11kX5e9yMKSHCtmI_GBj1-FDdE,1101
 google/cloud/dataproc_spark_connect/exceptions.py,sha256=WF-qdzgdofRwILCriIkjjsmjObZfF0P3Ecg4lv-Hmec,968
 google/cloud/dataproc_spark_connect/pypi_artifacts.py,sha256=gd-VMwiVP-EJuPp9Vf9Shx8pqps3oSKp0hBcSSZQS-A,1575
-google/cloud/dataproc_spark_connect/session.py,sha256=sX6uLFsm6wJWnE6zSNFCh5P5lJCoOhZpSfyUzFgwthw,30023
+google/cloud/dataproc_spark_connect/session.py,sha256=Sr9ISKIJ6U5dJ13FzKQ8UC_pGeFXbchc7X3d9U5Hj48,32144
 google/cloud/dataproc_spark_connect/client/__init__.py,sha256=6hCNSsgYlie6GuVpc5gjFsPnyeMTScTpXSPYqp1fplY,615
 google/cloud/dataproc_spark_connect/client/core.py,sha256=m3oXTKBm3sBy6jhDu9GRecrxLb5CdEM53SgMlnJb6ag,4616
 google/cloud/dataproc_spark_connect/client/proxy.py,sha256=qUZXvVY1yn934vE6nlO495XUZ53AUx9O74a9ozkGI9U,8976
-dataproc_spark_connect-0.8.0.dist-info/METADATA,sha256=xgmdQGsugNmxvYq9NwKS-FonzZmYUyK4HyWFA_xZnv0,3465
-dataproc_spark_connect-0.8.0.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-dataproc_spark_connect-0.8.0.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
-dataproc_spark_connect-0.8.0.dist-info/RECORD,,
+dataproc_spark_connect-0.8.2.dist-info/METADATA,sha256=2PCMrKtuuab4232elYKFHiTdaJcqiM4N38ceD_AhS-E,3465
+dataproc_spark_connect-0.8.2.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
+dataproc_spark_connect-0.8.2.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
+dataproc_spark_connect-0.8.2.dist-info/RECORD,,

google/cloud/dataproc_spark_connect/session.py

@@ -234,6 +234,9 @@ class DataprocSparkSession(SparkSession):
                     target=create_session_pbar
                 )
 
+                # Activate Spark Connect mode for Spark client
+                os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1"
+
                 try:
                     if (
                         os.getenv(
@@ -253,10 +256,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_view_session_details_button(session_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(
@@ -269,7 +271,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()
                     )
@@ -314,6 +316,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"]:
@@ -334,7 +376,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}"
                 )
-                self._display_view_session_details_button(s8s_session_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
@@ -453,6 +496,11 @@ class DataprocSparkSession(SparkSession):
         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(