diff --git a/databricks_cli/sdk/api_client.py b/databricks_cli/sdk/api_client.py index 2246023b..6939e690 100644 --- a/databricks_cli/sdk/api_client.py +++ b/databricks_cli/sdk/api_client.py @@ -110,10 +110,15 @@ def __init__(self, user=None, password=None, host=None, token=None, self.session = requests.Session() self.session.auth = FallbackNetrcAuth() self.session.mount('https://', TlsV1HttpAdapter(max_retries=retries)) - + parsed_url = urlparse(host) scheme = parsed_url.scheme - hostname = parsed_url.hostname + # this is to allow enterprise users to specify a custom host if they are forced to interact with Databricks via a proxy + if "databricks.com" in host: + hostname = parsed_url.hostname + else: + hostname = parsed_url.hostname + parsed_url.path + self.url = "%s://%s/api/" % (scheme, hostname) if user is not None and password is not None: encoded_auth = (user + ":" + password).encode() diff --git a/tests/sdk/test_api_client.py b/tests/sdk/test_api_client.py index f188d4b7..dc60a90f 100644 --- a/tests/sdk/test_api_client.py +++ b/tests/sdk/test_api_client.py @@ -133,6 +133,10 @@ def test_api_client_url_parsing(): client = ApiClient(host='https://databricks.com?o=123') assert client.get_url('') == 'https://databricks.com/api/2.0' + # this is to test that users can set a custom host for the apiclient if they are forced to interact with Databricks via a proxy + client = ApiClient(host='https://mydatabricksproxyservice.com/databricks/proxy') + assert client.get_url('') == 'https://mydatabricksproxyservice.com/databricks/proxy/api/2.0' + # NOTE: this technically is not possible since we validate that the "host" has a prefix of https:// in # databricks_cli.configure.cli client = ApiClient(host='http://databricks.com')