databricks-cli
databricks-cli copied to clipboard
Support for custom host in ApiClient
This is to allow enterprise users who are forced to interact with Databricks via a proxy to set a custom host for ApiClient.
Example
Before
client = ApiClient(host='https://company.cloud.databricks.com?o=123')
print(client.get_url("/"))
client = ApiClient(host='https://mydatabricksproxyservice.com/databricks/proxy')
print(client.get_url("/"))
-------------
https://company.cloud.databricks.com/api/2.0/
https://mydatabricksproxyservice.com/api/2.0/
After
client = ApiClient(host='https://company.cloud.databricks.com?o=123')
print(client.get_url("/"))
client = ApiClient(host='https://mydatabricksproxyservice.com/databricks/proxy')
print(client.get_url("/"))
-------------
https://company.cloud.databricks.com/api/2.0/
https://mydatabricksproxyservice.com/databricks/proxy/api/2.0/
Thanks for submitting the PR.
I'm hesitant to merge this because:
- It might break existing users who happened to have a path component in their host parameter
- It doesn't carry over to other tooling, e.g. Terraform
What kind of tools are you using to interact with Databricks where you need to use this proxy? E.g. do you use the ApiClient directly or also the CLI itself? Perhaps other custom tools?
If you're using ApiClient directly, you could subclass it and prepend your path prefix, for example:
class ProxyApiClient(ApiClient):
def __init__(self, path_prefix=None, **kwargs):
super().__init__(**kwargs)
self.path_prefix = path_prefix
def get_url(self, path, version=None):
url = super().get_url(path, version)
if self.path_prefix:
o = urlparse(url)
o = o._replace(path=self.path_prefix + o.path)
url = o.geturl()
return url
Thanks for review @pietern!
We interact with Databricks via the CLI directly so we can't modify code unless we fork it. Do you have an idea of a way of implementing this thats acceptable to merge? I could implement it if so
Would this change be more palatable if a DATABRICKS_BASE_URL or something similar was added. This would allow us to support this feature without breaking backcompat, but does seem like a bit of a hack 😅.
if base_url := os.environ.get("DATABRICKS_BASE_URL"): # this would probably be moved to config
self.url = base_url
else:
... original logic