From 0eb56f112036c9ad34861b288f4ecaeb61e50b1b Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Fri, 7 Mar 2025 11:29:04 +0530 Subject: [PATCH 1/4] make user_agent_header part of public API --- examples/set_user_agent.py | 2 +- src/databricks/sql/client.py | 9 ++++++--- src/databricks/sql/thrift_backend.py | 2 -- tests/unit/test_client.py | 2 +- 4 files changed, 8 insertions(+), 7 deletions(-) diff --git a/examples/set_user_agent.py b/examples/set_user_agent.py index 93eb2e0b3..093f03bd5 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 dca286ef4..65b57d58e 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 + 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)] diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index f76350a21..45fc5d6dd 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 0ff660d5c..c39aeb524 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( From 16cb0735cf273deba1f1026a02ac1e12d8ce406f Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Mon, 10 Mar 2025 12:38:19 +0530 Subject: [PATCH 2/4] removed user_agent_entry from list of internal params --- src/databricks/sql/client.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 608fb9cf6..5049443b4 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -179,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 From ebffb969135f22a18c06a917fef0a19337af46d6 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Tue, 11 Mar 2025 10:51:03 +0530 Subject: [PATCH 3/4] add backward compatibility --- src/databricks/sql/client.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 5049443b4..655cbb831 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -123,7 +123,7 @@ def __init__( `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 + 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`. @@ -228,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 next release." + ) + + 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)] From b66d2738bdaa00679a038d448683a927213c56e7 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Wed, 12 Mar 2025 09:50:10 +0530 Subject: [PATCH 4/4] minor change to warning --- src/databricks/sql/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 655cbb831..93e3975ad 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -234,7 +234,7 @@ def read(self) -> Optional[OAuthToken]: 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 next release." + "This parameter will be removed in the upcoming releases." ) if user_agent_entry: