make user_agent_header part of public API

Author: shivam2680

examples/set_user_agent.py

@@ -5,7 +5,7 @@
     server_hostname=os.getenv("DATABRICKS_SERVER_HOSTNAME"),
     http_path=os.getenv("DATABRICKS_HTTP_PATH"),
     access_token=os.getenv("DATABRICKS_TOKEN"),
-    _user_agent_entry="ExamplePartnerTag",
+    user_agent_entry="ExamplePartnerTag",
 ) as connection:
 
     with connection.cursor() as cursor:

src/databricks/sql/client.py

@@ -122,6 +122,9 @@ def __init__(
                 port of the oauth redirect uri (localhost). This is required when custom oauth client_id
                 `oauth_client_id` is set
 
+            user_agent_entry: `str`, optional
+                Tag to add to User-Agent header. For use by partners. If not specified, it will use the default user agent PyDatabricksSqlConnector
+
             experimental_oauth_persistence: configures preferred storage for persisting oauth tokens.
                 This has to be a class implementing `OAuthPersistence`.
                 When `auth_type` is set to `databricks-oauth` or `azure-oauth` without persisting the oauth token in a
@@ -176,7 +179,7 @@ def read(self) -> Optional[OAuthToken]:
         """
 
         # Internal arguments in **kwargs:
-        # _user_agent_entry
+        # user_agent_entry
         #   Tag to add to User-Agent header. For use by partners.
         # _use_cert_as_auth
         #  Use a TLS cert instead of a token
@@ -227,11 +230,11 @@ def read(self) -> Optional[OAuthToken]:
             server_hostname, **kwargs
         )
 
-        if not kwargs.get("_user_agent_entry"):
+        if not kwargs.get("user_agent_entry"):
             useragent_header = "{}/{}".format(USER_AGENT_NAME, __version__)
         else:
             useragent_header = "{}/{} ({})".format(
-                USER_AGENT_NAME, __version__, kwargs.get("_user_agent_entry")
+                USER_AGENT_NAME, __version__, kwargs.get("user_agent_entry")
             )
 
         base_headers = [("User-Agent", useragent_header)]

src/databricks/sql/thrift_backend.py

@@ -95,8 +95,6 @@ def __init__(
         **kwargs,
     ):
         # Internal arguments in **kwargs:
-        # _user_agent_entry
-        #   Tag to add to User-Agent header. For use by partners.
         # _username, _password
         #   Username and password Basic authentication (no official support)
         # _connection_uri

tests/unit/test_client.py

@@ -157,7 +157,7 @@ def test_useragent_header(self, mock_client_class):
         )
         self.assertIn(user_agent_header, http_headers)
 
-        databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS, _user_agent_entry="foobar")
+        databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS, user_agent_entry="foobar")
         user_agent_header_with_entry = (
             "User-Agent",
             "{}/{} ({})".format(