fix: pysparkvalue error

Author: zhiweilin-svg

google/cloud/dataproc_spark_connect/session.py

@@ -472,6 +472,28 @@ def create_session_pbar():
                     session_response, dataproc_config.name
                 )
 
+        # def _wait_for_spark_connect_endpoint(
+        #     self, session_name: str, timeout: int = 300
+        # ) -> Session:
+        #     """Waits for the Spark Connect endpoint to be available in the 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"https://console.cloud.google.com/dataproc/interactive/{self._region}/{session_id}?project={self._project_id}"
             plain_message = f"Creating Dataproc Session: {session_url}"
@@ -537,6 +559,10 @@ def _get_exiting_active_session(
                 )
                 self._display_view_session_details_button(s8s_session_id)
                 if session is None:
+                    # Wait for the Spark Connect endpoint to be available
+                    # session_response = self._wait_for_spark_connect_endpoint(
+                    #     session_name
+                    # )
                     session = self.__create_spark_connect_session_from_s8s(
                         session_response, session_name
                     )
@@ -559,6 +585,13 @@ def getOrCreate(self) -> "DataprocSparkSession":
                 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):
@@ -1162,6 +1195,20 @@ def stop(self) -> None:
                     )
 
                 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