diff --git a/examples/set_user_agent.py b/examples/set_user_agent.py index 93eb2e0b..093f03bd 100644 --- a/examples/set_user_agent.py +++ b/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: diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index b3cafea0..93e3975a 100755 --- a/src/databricks/sql/client.py +++ b/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 + A custom tag to append to the User-Agent header. This is typically used by partners to identify their applications.. 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,8 +179,6 @@ def read(self) -> Optional[OAuthToken]: """ # Internal arguments in **kwargs: - # _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 # _enable_ssl @@ -227,12 +228,21 @@ def read(self) -> Optional[OAuthToken]: server_hostname, **kwargs ) - if not kwargs.get("_user_agent_entry"): - useragent_header = "{}/{}".format(USER_AGENT_NAME, __version__) - else: + user_agent_entry = kwargs.get("user_agent_entry") + if user_agent_entry is None: + user_agent_entry = kwargs.get("_user_agent_entry") + if user_agent_entry is not None: + logger.warning( + "[WARN] Parameter '_user_agent_entry' is deprecated; use 'user_agent_entry' instead. " + "This parameter will be removed in the upcoming releases." + ) + + if user_agent_entry: useragent_header = "{}/{} ({})".format( - USER_AGENT_NAME, __version__, kwargs.get("_user_agent_entry") + USER_AGENT_NAME, __version__, user_agent_entry ) + else: + useragent_header = "{}/{}".format(USER_AGENT_NAME, __version__) base_headers = [("User-Agent", useragent_header)] diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index 972d9a9b..e1e27273 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/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 diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 0ff660d5..c39aeb52 100644 --- a/tests/unit/test_client.py +++ b/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(