Add methods to provide extra user agent and upstream user agent to sdk config

databricks/sdk/core.py

@@ -499,6 +499,7 @@ def __init__(self,
                  product_version="0.0.0",
                  **kwargs):
         self._inner = {}
+        self._user_agent_other_info = []
         self._credentials_provider = credentials_provider if credentials_provider else DefaultCredentials()
         try:
             self._set_inner_config(kwargs)
@@ -602,8 +603,29 @@ def user_agent(self):
         """ Returns User-Agent header used by this SDK """
         py_version = platform.python_version()
         os_name = platform.uname().system.lower()
-        return (f"{self._product}/{self._product_version} databricks-sdk-py/{__version__}"
-                f" python/{py_version} os/{os_name} auth/{self.auth_type}")
+
+        ua = [
+            f"{self._product}/{self._product_version}", f"databricks-sdk-py/{__version__}",
+            f"python/{py_version}", f"os/{os_name}", f"auth/{self.auth_type}",
+        ]
+        if len(self._user_agent_other_info) > 0:
+            ua.append(' '.join(self._user_agent_other_info))
+        if len(self._upstream_user_agent) > 0:
+            ua.append(self._upstream_user_agent)
+
+        return ' '.join(ua)
+
+    @property
+    def _upstream_user_agent(self) -> str:
+        product = os.environ.get('DATABRICKS_SDK_UPSTREAM', None)
+        product_version = os.environ.get('DATABRICKS_SDK_UPSTREAM_VERSION', None)
+        if product is not None and product_version is not None:
+            return f"upstream/{product} upstream-version/{product_version}"
+        return ""
+
+    def with_user_agent_extra(self, key: str, value: str) -> 'Config':
+        self._user_agent_other_info.append(f"{key}/{value}")
+        return self
 
     @property
     def oidc_endpoints(self) -> Optional[OidcEndpoints]:

tests/test_core.py

@@ -1,12 +1,14 @@
 import os
 import pathlib
+import platform
 import random
 import string
 
 import pytest
 
 from databricks.sdk.core import (Config, DatabricksCliTokenSource,
                                  databricks_cli)
+from databricks.sdk.version import __version__
 
 
 def test_parse_dsn():
@@ -123,3 +125,25 @@ def test_databricks_cli_credential_provider_installed_new(config, monkeypatch, t
     write_large_dummy_executable(tmp_path)
     monkeypatch.setenv('PATH', str(os.pathsep).join([tmp_path.as_posix(), os.environ['PATH']]))
     assert databricks_cli(config) is not None
+
+
+def test_extra_and_upstream_user_agent(monkeypatch):
+
+    class MockUname:
+
+        @property
+        def system(self):
+            return 'TestOS'
+
+    monkeypatch.setattr(platform, 'python_version', lambda: '3.0.0')
+    monkeypatch.setattr(platform, 'uname', MockUname)
+    monkeypatch.setenv('DATABRICKS_SDK_UPSTREAM', "upstream-product")
+    monkeypatch.setenv('DATABRICKS_SDK_UPSTREAM_VERSION', "0.0.1")
+
+    config = Config(host='http://localhost', username="something", password="something", product='test', product_version='0.0.0')\
+        .with_user_agent_extra('test-extra-1', '1')\
+        .with_user_agent_extra('test-extra-2', '2')
+
+    assert config.user_agent == (
+        f"test/0.0.0 databricks-sdk-py/{__version__} python/3.0.0 os/testos auth/basic"
+        f" test-extra-1/1 test-extra-2/2 upstream/upstream-product upstream-version/0.0.1")