remove unnecessary keyword

Author: cookiedough77

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

@@ -709,8 +709,8 @@ def __init__(
         # cleanup ml cache if possible
         atexit.register(self._cleanup_ml_cache)
 
-        self.global_user_context_extensions = []  # EDGE
-        self.global_user_context_extensions_lock = threading.Lock()  # EDGE
+        self.global_user_context_extensions = []
+        self.global_user_context_extensions_lock = threading.Lock()
 
     @property
     def _stub(self) -> grpc_lib.SparkConnectServiceStub:
@@ -1243,7 +1243,6 @@ def token(self) -> Optional[str]:
         """
         return self._builder.token
 
-    # BEGIN-EDGE
     def _update_request_with_user_context_extensions(
         self,
         req: Union[
@@ -1262,8 +1261,6 @@ def _update_request_with_user_context_extensions(
         for _, extension in self.thread_local.user_context_extensions:
             req.user_context.extensions.append(extension)
 
-    # END-EDGE
-
     def _execute_plan_request_with_metadata(
         self, operation_id: Optional[str] = None
     ) -> pb2.ExecutePlanRequest:
@@ -1294,9 +1291,7 @@ def _execute_plan_request_with_metadata(
                     messageParameters={"arg_name": "operation_id", "origin": str(ve)},
                 )
             req.operation_id = operation_id
-        # BEGIN-EDGE
         self._update_request_with_user_context_extensions(req)
-        # END-EDGE
         return req
 
     def _analyze_plan_request_with_metadata(self) -> pb2.AnalyzePlanRequest:
@@ -1307,9 +1302,7 @@ def _analyze_plan_request_with_metadata(self) -> pb2.AnalyzePlanRequest:
         req.client_type = self._builder.userAgent
         if self._user_id:
             req.user_context.user_id = self._user_id
-        # BEGIN-EDGE
         self._update_request_with_user_context_extensions(req)
-        # END-EDGE
         return req
 
     def _analyze(self, method: str, **kwargs: Any) -> AnalyzeResult:
@@ -1724,9 +1717,7 @@ def _config_request_with_metadata(self) -> pb2.ConfigRequest:
         req.client_type = self._builder.userAgent
         if self._user_id:
             req.user_context.user_id = self._user_id
-        # BEGIN-EDGE
         self._update_request_with_user_context_extensions(req)
-        # END-EDGE
         return req
 
     def get_configs(self, *keys: str) -> Tuple[Optional[str], ...]:
@@ -1902,7 +1893,6 @@ def _throw_if_invalid_tag(self, tag: str) -> None:
                 messageParameters={"arg_name": "Spark Connect tag", "arg_value": tag},
             )
 
-    # BEGIN-EDGE
     def add_threadlocal_user_context_extension(self, extension: any_pb2.Any) -> str:
         if not hasattr(self.thread_local, "user_context_extensions"):
             self.thread_local.user_context_extensions = list()
@@ -1935,8 +1925,6 @@ def clear_user_context_extensions(self) -> None:
         with self.global_user_context_extensions_lock:
             self.global_user_context_extensions = list()
 
-    # END-EDGE
-
     def _handle_error(self, error: Exception) -> NoReturn:
         """
         Handle errors that occur during RPC calls.
@@ -1977,9 +1965,7 @@ def _fetch_enriched_error(self, info: "ErrorInfo") -> Optional[pb2.FetchErrorDet
             req.client_observed_server_side_session_id = self._server_session_id
         if self._user_id:
             req.user_context.user_id = self._user_id
-        # BEGIN-EDGE
         self._update_request_with_user_context_extensions(req)
-        # END-EDGE
         try:
             return self._stub.FetchErrorDetails(req, metadata=self._builder.metadata())
         except grpc.RpcError:

python/pyspark/sql/connect/session.py

@@ -42,7 +42,7 @@
     ClassVar,
 )
 
-import google.protobuf.any_pb2 as any_pb2  # EDGE
+import google.protobuf.any_pb2 as any_pb2
 import numpy as np
 import pandas as pd
 import pyarrow as pa
@@ -895,7 +895,6 @@ def clearTags(self) -> None:
 
     clearTags.__doc__ = PySparkSession.clearTags.__doc__
 
-    # BEGIN-EDGE
     def addThreadlocalUserContextExtension(self, extension: any_pb2.Any) -> str:
         """
         Add a user context extension to the current session in the current thread.
@@ -950,8 +949,6 @@ def clearUserContextExtensions(self) -> None:
         """
         self.client.clear_user_context_extensions()
 
-    # END-EDGE
-
     def stop(self) -> None:
         """
         Release the current session and close the GRPC connection to the Spark Connect server.

python/pyspark/sql/tests/connect/client/test_client.py

@@ -136,11 +136,11 @@ class MockService:
         def __init__(self, session_id: str):
             self._session_id = session_id
             self.req = None
-            self.client_user_context_extensions = []  # EDGE
+            self.client_user_context_extensions = []
 
         def ExecutePlan(self, req: proto.ExecutePlanRequest, metadata):
             self.req = req
-            self.client_user_context_extensions = req.user_context.extensions  # EDGE
+            self.client_user_context_extensions = req.user_context.extensions
             resp = proto.ExecutePlanResponse()
             resp.session_id = self._session_id
             resp.operation_id = req.operation_id
@@ -161,14 +161,14 @@ def ExecutePlan(self, req: proto.ExecutePlanRequest, metadata):
 
         def Interrupt(self, req: proto.InterruptRequest, metadata):
             self.req = req
-            self.client_user_context_extensions = req.user_context.extensions  # EDGE
+            self.client_user_context_extensions = req.user_context.extensions
             resp = proto.InterruptResponse()
             resp.session_id = self._session_id
             return resp
 
         def Config(self, req: proto.ConfigRequest, metadata):
             self.req = req
-            self.client_user_context_extensions = req.user_context.extensions  # EDGE
+            self.client_user_context_extensions = req.user_context.extensions
             resp = proto.ConfigResponse()
             resp.session_id = self._session_id
             if req.operation.HasField("get"):
@@ -233,7 +233,6 @@ def userId(self) -> Optional[str]:
 
         self.assertEqual(client._user_id, "abc")
 
-    # BEGIN-EDGE
     def test_user_context_extension(self):
         client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
         mock = MockService(client._session_id)
@@ -321,8 +320,6 @@ def test_user_context_extension(self):
         self.assertFalse(exlocal2 in mock.client_user_context_extensions)
         self.assertFalse(exglobal2 in mock.client_user_context_extensions)
 
-    # END-EDGE
-
     def test_interrupt_all(self):
         client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
         mock = MockService(client._session_id)