add extension

Author: cookiedough77

python/pyspark/pipelines/block_connect_access.py

@@ -53,9 +53,29 @@ def blocked_method(*args: object, **kwargs: object) -> NoReturn:
     extension_id = None
     try:
         # Apply our custom __getattribute__ method
+        extension_id = add_user_context_extension(spark, definition_path="pipeline.path")
         setattr(SparkConnectServiceStub, "__getattribute__", blocked_getattr)
         yield
     finally:
         # Restore the original __getattribute__ method
+        if extension_id is not None:
+            spark.removeUserContextExtension(extension_id)
         setattr(SparkConnectServiceStub, "__getattribute__", original_getattr)
 
+def add_user_context_extension(spark: SparkSession, **context_fields: Any) -> str:
+    """
+    Adds a user context extension to Spark's thread-local user context and
+    returns the extension ID.
+    """
+    if not context_fields:
+        raise ValueError("At least one field must be provided for PipelineAnalysisContext.")
+
+    from pyspark.sql.connect.proto import pipelines_pb2
+
+    analysis_context = pipelines_pb2.PipelineAnalysisContext(**context_fields)
+
+    from google.protobuf import any_pb2
+
+    extension = any_pb2.Any()
+
+    return spark.addThreadlocalUserContextExtension(extension)