Fix URL construction in the ApiClient class (#193)

databricks_cli/sdk/api_client.py

@@ -39,6 +39,7 @@
 from . import version
 
 from requests.adapters import HTTPAdapter
+from six.moves.urllib.parse import urlparse
 
 try:
     from requests.packages.urllib3.poolmanager import PoolManager
@@ -72,7 +73,10 @@ def __init__(self, user=None, password=None, host=None, token=None,
         self.session = requests.Session()
         self.session.mount('https://', TlsV1HttpAdapter())
 
-        self.url = "%s/api/%s" % (host, apiVersion)
+        parsed_url = urlparse(host)
+        scheme = parsed_url.scheme
+        hostname = parsed_url.hostname
+        self.url = "%s://%s/api/%s" % (scheme, hostname, apiVersion)
         if user is not None and password is not None:
             encoded_auth = (user + ":" + password).encode()
             user_header_data = "Basic " + base64.standard_b64encode(encoded_auth).decode()
@@ -108,7 +112,6 @@ def perform_query(self, method, path, data = {}, headers = None):
             warnings.simplefilter("ignore", exceptions.InsecureRequestWarning)
             resp = self.session.request(method, self.url + path, data = json.dumps(data),
                 verify = self.verify, headers = headers)
-
         try:
             resp.raise_for_status()
         except requests.exceptions.HTTPError as e:

tests/configure/test_config.py

@@ -136,7 +136,7 @@ def inner_test_group():
     def test_command(api_client, x): # noqa
         click.echo(json.dumps(api_client.default_headers))
 
-    with mock.patch("databricks_cli.configure.provider.DatabricksConfig") as config_mock:
+    with mock.patch("databricks_cli.configure.config.get_config") as config_mock:
         with mock.patch("uuid.uuid1") as uuid_mock:
             config_mock.return_value = DatabricksConfig.from_token(TEST_HOST, TEST_TOKEN)
             uuid_mock.return_value = '1234'

tests/sdk/test_api_client.py

@@ -60,3 +60,18 @@ def test_content_from_server_on_error(m):
     with pytest.raises(requests.exceptions.HTTPError) as e:
         client.perform_query('GET', '/endpoint')
         assert error_message_contains in e.value.message
+
+def test_api_client_url_parsing():
+    client = ApiClient(host='https://databricks.com')
+    assert client.url == 'https://databricks.com/api/2.0'
+
+    client = ApiClient(host='https://databricks.com/?o=123')
+    assert client.url == 'https://databricks.com/api/2.0'
+
+    client = ApiClient(host='https://databricks.com?o=123')
+    assert client.url == 'https://databricks.com/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')
+    assert client.url == 'http://databricks.com/api/2.0'