[SPARK-42477][CONNECT][PYTHON] accept user_agent in spark connect's connection string

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

Currently, the Spark Connect service's `client_type` attribute (which is really [user
agent]) is set to `_SPARK_CONNECT_PYTHON` to signify PySpark.

With this change, the connection for the Spark Connect remote accepts an optional
`user_agent` parameter which is then passed down to the service.

[user agent]: https://www.w3.org/WAI/UA/work/wiki/Definition_of_User_Agent

### Why are the changes needed?

This enables partners using Spark Connect to set their application as the user agent,
which then allows visibility and measurement of integrations and usages of spark
connect.

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

A new optional `user_agent` parameter is now recognized as part of the Spark Connect
connection string.

### How was this patch tested?

- unit tests attached
- manually running the `pyspark` binary with the `user_agent` connection string set and
   verifying the payload sent to the server. Similar testing for the default.

Closes #40054 from nija-at/user-agent.

Authored-by: Niranjan Jayakar <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>

Author: nija-at

connector/connect/docs/client-connection-string.md

@@ -58,7 +58,8 @@ sc://hostname:port/;param1=value;param2=value
     <td>token</td>
     <td>String</td>
     <td>When this param is set in the URL, it will enable standard
-    bearer token authentication using GRPC. By default this value is not set.</td>
+    bearer token authentication using GRPC. By default this value is not set.
+    Setting this value enables SSL.</td>
     <td><pre>token=ABCDEFGH</pre></td>
   </tr>
   <tr>
@@ -81,6 +82,15 @@ sc://hostname:port/;param1=value;param2=value
     <pre>user_id=Martin</pre>
     </td>
   </tr>
+  <tr>
+    <td>user_agent</td>
+    <td>String</td>
+    <td>The user agent acting on behalf of the user, typically applications
+    that use Spark Connect to implement its functionality and execute Spark
+    requests on behalf of the user.<br/>
+    <i>Default: </i><pre>_SPARK_CONNECT_PYTHON</pre> in the Python client</td>
+    <td><pre>user_agent=my_data_query_app</pre></td>
+  </tr>
 </table>
 
 ## Examples

dev/sparktestsupport/modules.py

@@ -516,6 +516,7 @@ def __hash__(self):
         "pyspark.sql.connect.dataframe",
         "pyspark.sql.connect.functions",
         # unittests
+        "pyspark.sql.tests.connect.test_client",
         "pyspark.sql.tests.connect.test_connect_plan",
         "pyspark.sql.tests.connect.test_connect_basic",
         "pyspark.sql.tests.connect.test_connect_function",

python/pyspark/sql/connect/client.py

@@ -19,6 +19,8 @@
     "SparkConnectClient",
 ]
 
+import string
+
 from pyspark.sql.connect.utils import check_dependencies
 
 check_dependencies(__name__, __file__)
@@ -120,6 +122,7 @@ class ChannelBuilder:
     PARAM_USE_SSL = "use_ssl"
     PARAM_TOKEN = "token"
     PARAM_USER_ID = "user_id"
+    PARAM_USER_AGENT = "user_agent"
 
     @staticmethod
     def default_port() -> int:
@@ -215,6 +218,7 @@ def metadata(self) -> Iterable[Tuple[str, str]]:
                 ChannelBuilder.PARAM_TOKEN,
                 ChannelBuilder.PARAM_USE_SSL,
                 ChannelBuilder.PARAM_USER_ID,
+                ChannelBuilder.PARAM_USER_AGENT,
             ]
         ]
 
@@ -244,6 +248,27 @@ def userId(self) -> Optional[str]:
         """
         return self.params.get(ChannelBuilder.PARAM_USER_ID, None)
 
+    @property
+    def userAgent(self) -> str:
+        """
+        Returns
+        -------
+        user_agent : str
+            The user_agent parameter specified in the connection string,
+            or "_SPARK_CONNECT_PYTHON" when not specified.
+        """
+        user_agent = self.params.get(ChannelBuilder.PARAM_USER_AGENT, "_SPARK_CONNECT_PYTHON")
+        allowed_chars = string.ascii_letters + string.punctuation
+        if len(user_agent) > 200:
+            raise SparkConnectException(
+                "'user_agent' parameter cannot exceed 200 characters in length"
+            )
+        if set(user_agent).difference(allowed_chars):
+            raise SparkConnectException(
+                "Only alphanumeric and common punctuations are allowed for 'user_agent'"
+            )
+        return user_agent
+
     def get(self, key: str) -> Any:
         """
         Parameters
@@ -559,15 +584,15 @@ def close(self) -> None:
     def _execute_plan_request_with_metadata(self) -> pb2.ExecutePlanRequest:
         req = pb2.ExecutePlanRequest()
         req.client_id = self._session_id
-        req.client_type = "_SPARK_CONNECT_PYTHON"
+        req.client_type = self._builder.userAgent
         if self._user_id:
             req.user_context.user_id = self._user_id
         return req
 
     def _analyze_plan_request_with_metadata(self) -> pb2.AnalyzePlanRequest:
         req = pb2.AnalyzePlanRequest()
         req.client_id = self._session_id
-        req.client_type = "_SPARK_CONNECT_PYTHON"
+        req.client_type = self._builder.userAgent
         if self._user_id:
             req.user_context.user_id = self._user_id
         return req

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

@@ -0,0 +1,67 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+from typing import Optional
+
+from pyspark.sql.connect.client import SparkConnectClient
+import pyspark.sql.connect.proto as proto
+
+
+class SparkConnectClientTestCase(unittest.TestCase):
+    def test_user_agent_passthrough(self):
+        client = SparkConnectClient("sc://foo/;user_agent=bar")
+        mock = MockService(client._session_id)
+        client._stub = mock
+
+        command = proto.Command()
+        client.execute_command(command)
+
+        self.assertIsNotNone(mock.req, "ExecutePlan API was not called when expected")
+        self.assertEqual(mock.req.client_type, "bar")
+
+    def test_user_agent_default(self):
+        client = SparkConnectClient("sc://foo/")
+        mock = MockService(client._session_id)
+        client._stub = mock
+
+        command = proto.Command()
+        client.execute_command(command)
+
+        self.assertIsNotNone(mock.req, "ExecutePlan API was not called when expected")
+        self.assertEqual(mock.req.client_type, "_SPARK_CONNECT_PYTHON")
+
+
+class MockService:
+    # Simplest mock of the SparkConnectService.
+    # If this needs more complex logic, it needs to be replaced with Python mocking.
+
+    req: Optional[proto.ExecutePlanRequest]
+
+    def __init__(self, session_id: str):
+        self._session_id = session_id
+        self.req = None
+
+    def ExecutePlan(self, req: proto.ExecutePlanRequest, metadata):
+        self.req = req
+        resp = proto.ExecutePlanResponse()
+        resp.client_id = self._session_id
+        return [resp]
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -2949,10 +2949,33 @@ def test_sensible_defaults(self):
 
         chan = ChannelBuilder("sc://host/;token=abcs")
         self.assertTrue(chan.secure, "specifying a token must set the channel to secure")
-
+        self.assertEqual(chan.userAgent, "_SPARK_CONNECT_PYTHON")
         chan = ChannelBuilder("sc://host/;use_ssl=abcs")
         self.assertFalse(chan.secure, "Garbage in, false out")
 
+    def test_invalid_user_agent_charset(self):
+        # fmt: off
+        invalid_user_agents = [
+            "agent»",  # non standard symbol
+            "age nt",  # whitespace
+            "ägent",   # non-ascii alphabet
+        ]
+        # fmt: on
+        for user_agent in invalid_user_agents:
+            with self.subTest(user_agent=user_agent):
+                chan = ChannelBuilder(f"sc://host/;user_agent={user_agent}")
+                with self.assertRaises(SparkConnectException) as err:
+                    chan.userAgent
+
+                self.assertRegex(err.exception.message, "alphanumeric and common punctuations")
+
+    def test_invalid_user_agent_len(self):
+        user_agent = "x" * 201
+        chan = ChannelBuilder(f"sc://host/;user_agent={user_agent}")
+        with self.assertRaises(SparkConnectException) as err:
+            chan.userAgent
+        self.assertRegex(err.exception.message, "characters in length")
+
     def test_valid_channel_creation(self):
         chan = ChannelBuilder("sc://host").toChannel()
         self.assertIsInstance(chan, grpc.Channel)
@@ -2965,8 +2988,9 @@ def test_valid_channel_creation(self):
         self.assertIsInstance(chan, grpc.Channel)
 
     def test_channel_properties(self):
-        chan = ChannelBuilder("sc://host/;use_ssl=true;token=abc;param1=120%2021")
+        chan = ChannelBuilder("sc://host/;use_ssl=true;token=abc;user_agent=foo;param1=120%2021")
         self.assertEqual("host:15002", chan.endpoint)
+        self.assertEqual("foo", chan.userAgent)
         self.assertEqual(True, chan.secure)
         self.assertEqual("120 21", chan.get("param1"))