This is an automated email from the ASF dual-hosted git repository.
LuciferYang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new b85c7e944955 [SPARK-57595][PYTHON] Declarative Pipelines analysis
context should not mask the original error when registration fails
b85c7e944955 is described below
commit b85c7e944955913e95e9b0178213b82d49046ac5
Author: YangJie <[email protected]>
AuthorDate: Wed Jun 24 13:39:36 2026 +0800
[SPARK-57595][PYTHON] Declarative Pipelines analysis context should not
mask the original error when registration fails
### What changes were proposed in this pull request?
`add_pipeline_analysis_context` sets `extension_id = None` up front and
only assigns it after the extension is registered, but the `finally` block
always calls `client.remove_user_context_extension(extension_id)`. If any step
before that assignment fails, `extension_id` is still `None`, and
`remove_user_context_extension(None)` runs `None.find(...)` and raises
`AttributeError`, which hides the real failure.
This change runs the cleanup only when an extension was actually registered:
```python
finally:
if extension_id is not None:
client.remove_user_context_extension(extension_id)
```
### Why are the changes needed?
A failure while setting up the analysis context (for example the proto
import or `add_threadlocal_user_context_extension`) currently surfaces as
`AttributeError: 'NoneType' object has no attribute 'find'` instead of the
error that actually caused it.
### Does this PR introduce _any_ user-facing change?
No. It only changes which exception is raised when setup fails: the
original one rather than the `AttributeError` that was masking it.
### How was this patch tested?
Added a test that makes extension registration raise and checks that the
original error propagates and that no extension is left registered. It fails
without the fix and passes with it.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.8)
Closes #56652 from LuciferYang/sdp-analysis-context-finally-guard.
Authored-by: YangJie <[email protected]>
Signed-off-by: yangjie01 <[email protected]>
(cherry picked from commit 6c2a12678d3ec1c241deff1441b6278465948c94)
Signed-off-by: yangjie01 <[email protected]>
---
.../pipelines/add_pipeline_analysis_context.py | 5 ++-
.../tests/test_add_pipeline_analysis_context.py | 38 ++++++++++++++++++++++
2 files changed, 42 insertions(+), 1 deletion(-)
diff --git a/python/pyspark/pipelines/add_pipeline_analysis_context.py
b/python/pyspark/pipelines/add_pipeline_analysis_context.py
index 6ecabdf43b07..6d0bd4dd7308 100644
--- a/python/pyspark/pipelines/add_pipeline_analysis_context.py
+++ b/python/pyspark/pipelines/add_pipeline_analysis_context.py
@@ -45,4 +45,7 @@ def add_pipeline_analysis_context(
extension_id = client.add_threadlocal_user_context_extension(extension)
yield
finally:
- client.remove_user_context_extension(extension_id)
+ # extension_id stays None if registering the extension above failed;
skip cleanup in that
+ # case so we don't call remove_user_context_extension(None) and mask
the original error.
+ if extension_id is not None:
+ client.remove_user_context_extension(extension_id)
diff --git
a/python/pyspark/pipelines/tests/test_add_pipeline_analysis_context.py
b/python/pyspark/pipelines/tests/test_add_pipeline_analysis_context.py
index 44e60f4597c6..f99e7d56c3d5 100644
--- a/python/pyspark/pipelines/tests/test_add_pipeline_analysis_context.py
+++ b/python/pyspark/pipelines/tests/test_add_pipeline_analysis_context.py
@@ -15,6 +15,7 @@
# limitations under the License.
#
import unittest
+from unittest import mock
from pyspark.testing.connectutils import (
ReusedConnectTestCase,
@@ -89,6 +90,43 @@ class AddPipelineAnalysisContextTests(ReusedConnectTestCase):
thread_local_extensions_after_2 =
self.spark.client.thread_local.user_context_extensions
self.assertEqual(len(thread_local_extensions_after_2), 0)
+ def test_setup_failure_does_not_mask_original_error(self):
+ # If any setup step fails before the extension is registered,
extension_id stays None and
+ # the finally block must skip remove_user_context_extension(None) -
which would raise
+ # AttributeError and mask the original error. Cover each step that can
fail.
+ import pyspark.sql.connect.proto as pb2
+ from google.protobuf import any_pb2
+
+ failing_any = mock.MagicMock()
+ failing_any.Pack.side_effect = ValueError("boom")
+
+ failure_points = {
+ "context construction": mock.patch.object(
+ pb2, "PipelineAnalysisContext", side_effect=ValueError("boom")
+ ),
+ "extension packing": mock.patch.object(any_pb2, "Any",
return_value=failing_any),
+ "extension registration": mock.patch.object(
+ self.spark.client,
+ "add_threadlocal_user_context_extension",
+ side_effect=ValueError("boom"),
+ ),
+ }
+ for step, patcher in failure_points.items():
+ with self.subTest(failing_step=step):
+ with patcher:
+ with self.assertRaises(ValueError):
+ with add_pipeline_analysis_context(
+ self.spark, "test_dataflow_graph_id", None
+ ):
+ pass
+ # A failed setup should not leave an extension registered.
Read lazily, since the
+ # attribute is only created on first successful registration,
which never happens
+ # here.
+ thread_local_extensions = getattr(
+ self.spark.client.thread_local, "user_context_extensions",
[]
+ )
+ self.assertEqual(len(thread_local_extensions), 0)
+
if __name__ == "__main__":
from pyspark.testing import main
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
