[SPARK-47986][CONNECT][PYTHON] Unable to create a new session when th…

…e default session is closed by the server

### What changes were proposed in this pull request?

When the server closes a session, usually after a cluster restart,
the client is unaware of this until it receives an error.

At this point, the client in unable to create a new session to the
same connect endpoint, since the stale session is still recorded
as the active and default session.

With this change, when the server communicates that the session
has changed via a GRPC error, the session and the respective client
are marked as stale. A new default connection can be created
via the session builder.

### Why are the changes needed?

See section above.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Attached unit tests

### Was this patch authored or co-authored using generative AI tooling?

No.

Closes apache#46221 from nija-at/session-expires.

Authored-by: Niranjan Jayakar <nija@databricks.com>
Signed-off-by: Ruifeng Zheng <ruifengz@apache.org>

python/pyspark/sql/connect/client/core.py

@@ -1763,6 +1763,9 @@ def _handle_rpc_error(self, rpc_error: grpc.RpcError) -> NoReturn:
                     info = error_details_pb2.ErrorInfo()
                     d.Unpack(info)
 
+                    if info.metadata["errorClass"] == "INVALID_HANDLE.SESSION_CHANGED":
+                        self._closed = True
+
                     raise convert_exception(
                         info,
                         status.message,

python/pyspark/sql/connect/session.py

@@ -237,9 +237,9 @@ def create(self) -> "SparkSession":
         def getOrCreate(self) -> "SparkSession":
             with SparkSession._lock:
                 session = SparkSession.getActiveSession()
-                if session is None:
+                if session is None or session.is_stopped:
                     session = SparkSession._default_session
-                    if session is None:
+                    if session is None or session.is_stopped:
                         session = self.create()
                 self._apply_options(session)
                 return session

python/pyspark/sql/tests/connect/test_connect_session.py

@@ -242,6 +242,20 @@ def toChannel(self):
         session = RemoteSparkSession.builder.channelBuilder(CustomChannelBuilder()).create()
         session.sql("select 1 + 1")
 
+    def test_reset_when_server_session_changes(self):
+        session = RemoteSparkSession.builder.remote("sc://localhost").getOrCreate()
+        # run a simple query so the session id is synchronized.
+        session.range(3).collect()
+
+        # trigger a mismatch between client session id and server session id.
+        session._client._session_id = str(uuid.uuid4())
+        with self.assertRaises(SparkConnectException):
+            session.range(3).collect()
+
+        # assert that getOrCreate() generates a new session
+        session = RemoteSparkSession.builder.remote("sc://localhost").getOrCreate()
+        session.range(3).collect()
+
 
 class SparkConnectSessionWithOptionsTest(unittest.TestCase):
     def setUp(self) -> None: