From 8436b2842d8574be009149a85bd9ffd8a8a7939a Mon Sep 17 00:00:00 2001
From: Andrew Ma <136692+ajma@users.noreply.github.com>
Date: Tue, 7 Apr 2026 15:48:51 -0700
Subject: [PATCH] fix: Clean up leaked progress bars on session stop

_execution_progress_bar entries are only removed when operations complete
with done=True. If operations are interrupted or fail before the handler
is called with done=True, their tqdm objects leak. Now clears all
remaining progress bars during stop().
---
 google/cloud/dataproc_spark_connect/session.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/google/cloud/dataproc_spark_connect/session.py b/google/cloud/dataproc_spark_connect/session.py
index 7ad24fe..b5e5451 100644
--- a/google/cloud/dataproc_spark_connect/session.py
+++ b/google/cloud/dataproc_spark_connect/session.py
@@ -1259,6 +1259,14 @@ def stop(self, terminate: Optional[bool] = None) -> None:
         >>> spark.stop(terminate=False)
         """
         with DataprocSparkSession._lock:
+            # Clean up any remaining progress bars from interrupted operations
+            for pbar in DataprocSparkSession._execution_progress_bar.values():
+                try:
+                    pbar.close()
+                except Exception:
+                    pass
+            DataprocSparkSession._execution_progress_bar.clear()
+
             if DataprocSparkSession._active_s8s_session_id is not None:
                 # Determine if we should terminate the server-side session
                 if terminate is None: