Skip to content

Commit b13042b

Browse files
mgyuchtmgyucht
authored
Better error message when private link enabled workspaces reject requests (#647)
## Changes This PR ports databricks/databricks-sdk-go#924 to the Python SDK. When a user tries to access a Private Link-enabled workspace configured with no public internet access from a different network than the VPC endpoint belongs to, the Private Link backend redirects the user to the login page, rather than outright rejecting the request. The login page, however, is not a JSON document and cannot be parsed by the SDK, resulting in this error message: ``` $ databricks current-user me Error: unexpected error handling request: invalid character '<' looking for beginning of value. This is likely a bug in the Databricks SDK for Go or the underlying REST API. Please report this issue with the following debugging information to the SDK issue tracker at https://github.com/databricks/databricks-sdk-go/issues. Request log: GET /login.html?error=private-link-validation-error:<WSID> > * Host: > * Accept: application/json > * Authorization: REDACTED > * Referer: https://adb-<WSID>.azuredatabricks.net/api/2.0/preview/scim/v2/Me > * User-Agent: cli/0.0.0-dev+5ed10bb8ccc1 databricks-sdk-go/0.39.0 go/1.22.2 os/darwin cmd/current-user_me auth/pat < HTTP/2.0 200 OK < * Cache-Control: no-cache, no-store, must-revalidate < * Content-Security-Policy: default-src *; font-src * data:; frame-src * blob:; img-src * blob: data:; media-src * data:; object-src 'none'; style-src * 'unsafe-inline'; worker-src * blob:; script-src 'self' 'unsafe-eval' 'unsafe-hashes' 'report-sample' https://*.databricks.com https://databricks.github.io/debug-bookmarklet/ https://widget.intercom.io https://js.intercomcdn.com https://ajax.googleapis.com/ajax/libs/jquery/2.2.0/jquery.min.js https://databricks-ui-assets.azureedge.net https://ui-serving-cdn-testing.azureedge.net https://uiserviceprodwestus-cdn-endpoint.azureedge.net https://databricks-ui-infra.s3.us-west-2.amazonaws.com 'sha256-47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=' 'sha256-YOlue469P2BtTMZYUFLupA2aOUsgc6B/TDewH7/Qz7s=' 'sha256-Lh4yp7cr3YOJ3MOn6erNz3E3WI0JA20mWV+0RuuviFM=' 'sha256-0jMhpY6PB/BTRDLWtfcjdwiHOpE+6rFk3ohvY6fbuHU='; report-uri /ui-csp-reports; frame-ancestors *.vocareum.com *.docebosaas.com *.edx.org *.deloitte.com *.cloudlabs.ai *.databricks.com *.myteksi.net < * Content-Type: text/html; charset=utf-8 < * Date: Fri, 17 May 2024 07:47:38 GMT < * Server: databricks < * Set-Cookie: enable-armeria-workspace-server-for-ui-flags=false; Max-Age=1800; Expires=Fri, 17 May 2024 08:17:38 GMT; Secure; HTTPOnly; SameSite=Strict < * Strict-Transport-Security: max-age=31536000; includeSubDomains; preload < * X-Content-Type-Options: nosniff < * X-Ui-Svc: true < * X-Xss-Protection: 1; mode=block < <!doctype html> < <html> < <head> < <meta charset="utf-8"> < <meta http-equiv="Content-Language" content="en"> < <title>Databricks - Sign In</title> < <meta name="viewport" content="width=960"> < <link rel="icon" type="image/png" href="https://databricks-ui-assets.azureedge.net/favicon.ico"> < <meta http-equiv="content-type" content="text/html; charset=UTF8"> < <script id="__databricks_react_script"></script> < <script>window.__DATABRICKS_SAFE_FLAGS__={"databricks.infra.showErrorModalOnFetchError":true,"databricks.fe.infra.useReact18":true,"databricks.fe.infra.useReact18NewAPI":false,"databricks.fe.infra.fixConfigPrefetch":true},window.__DATABRICKS_CONFIG__={"publicPath":{"mlflow":"https://databricks-ui-assets.azureedge.net/","dbsql":"https://databricks-ui-assets.azureedge.net/","feature-store":"https://databricks-ui-assets.azureedge.net/","monolith":"https://databricks-ui-assets.azureedge.net/","jaws":"https://databricks-ui-assets.azureedge.net/"}}</script> < <link rel="icon" href="https://databricks-ui-assets.azureedge.net/favicon.ico"> < <script> < function setNoCdnAndReload() { < document.cookie = `x-databricks-cdn-inaccessible=true; path=/; max-age=86400`; < const metric = 'cdnFallbackOccurred'; < const browserUserAgent = navigator.userAgent; < const browserTabId = window.browserTabId; < const performanceEntry = performance.getEntriesByType('resource').filter(e => e.initiatorType === 'script').slice(-1)[0] < sessionStorage.setItem('databricks-cdn-fallback-telemetry-key', JSON.stringify({ tags: { browserUserAgent, browserTabId }, performanceEntry})); < window.location.reload(); < } < </script> < <script> < // Set a manual timeout for dropped packets to CDN < function loadScriptWithTimeout(src, timeout) { < return new Promise((resolve, reject) => { < const script = document.createElement('script'); < script.defer = true; < script.src = src; < script.onload = resolve; < script.onerror = reject; < document.head.appendChild(script); < setTimeout(() => { < reject(new Error('Script load timeout')); < }, timeout); < }); < } < loadScriptWithTimeout('https://databricks-ui-assets.azureedge.net/static/js/login/login.8a983ca2.js', 10000).catch(setNoCdnAndReload); < </script> < </head> < <body class="light-mode"> < <uses-legacy-bootstrap> < <div id="login-page"></div> < </uses-legacy-bootstrap> < </body> < </html> ``` To address this, I add one additional check in the error mapper logic to inspect whether the user was redirected to the login page with the private link validation error response code. If so, we return a custom error, `PrivateLinkValidationError`, with error code `PRIVATE_LINK_VALIDATION_ERROR` that inherits from PermissionDenied and has a mock 403 status code. After this change, users will see an error message like this: ``` databricks.sdk.errors.private_link.PrivateLinkValidationError: The requested workspace has Azure Private Link enabled and is not accessible from the current network. Ensure that Azure Private Link is properly configured and that your device has access to the Azure Private Link endpoint. For more information, see https://learn.microsoft.com/en-us/azure/databricks/security/network/classic/private-link-standard#authentication-troubleshooting. ``` The error message is tuned to the specific cloud so that we can redirect users to the appropriate documentation, the cloud being inferred from the request URI. ## Tests Unit tests cover the private link error message mapping. To manually test this, I created a private link workspace in Azure, created an access token, restricted access to the workspace, then ran the `last_job_runs.py` example using the host & token: ``` /Users/miles/databricks-cli/.venv/bin/python /Users/miles/databricks-sdk-py/examples/last_job_runs.py 2024-05-17 11:43:32,529 [databricks.sdk][INFO] loading DEFAULT profile from ~/.databrickscfg: host, token Traceback (most recent call last): File "/Users/miles/databricks-sdk-py/examples/last_job_runs.py", line 20, in <module> for job in w.jobs.list(): File "/Users/miles/databricks-sdk-py/databricks/sdk/service/jobs.py", line 5453, in list json = self._api.do('GET', '/api/2.1/jobs/list', query=query, headers=headers) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/miles/databricks-sdk-py/databricks/sdk/core.py", line 131, in do response = retryable(self._perform)(method, ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/miles/databricks-sdk-py/databricks/sdk/retries.py", line 54, in wrapper raise err File "/Users/miles/databricks-sdk-py/databricks/sdk/retries.py", line 33, in wrapper return func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/Users/miles/databricks-sdk-py/databricks/sdk/core.py", line 245, in _perform raise self._make_nicer_error(response=response) from None databricks.sdk.errors.private_link.PrivateLinkValidationError: The requested workspace has Azure Private Link enabled and is not accessible from the current network. Ensure that Azure Private Link is properly configured and that your device has access to the Azure Private Link endpoint. For more information, see https://learn.microsoft.com/en-us/azure/databricks/security/network/classic/private-link-standard#authentication-troubleshooting. ``` - [ ] `make test` run locally - [ ] `make fmt` applied - [ ] relevant integration tests applied
1 parent fabe7c4 commit b13042b

File tree

11 files changed

+143
-52
lines changed

11 files changed

+143
-52
lines changed

‎databricks/sdk/azure.py‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,36 +1,9 @@
1-
from dataclasses import dataclass
21
from typing import Dict
32

43
from .oauth import TokenSource
54
from .service.provisioning import Workspace
65

76

8-
@dataclass
9-
class AzureEnvironment:
10-
name: str
11-
service_management_endpoint: str
12-
resource_manager_endpoint: str
13-
active_directory_endpoint: str
14-
15-
16-
ARM_DATABRICKS_RESOURCE_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
17-
18-
ENVIRONMENTS = dict(
19-
PUBLIC=AzureEnvironment(name="PUBLIC",
20-
service_management_endpoint="https://management.core.windows.net/",
21-
resource_manager_endpoint="https://management.azure.com/",
22-
active_directory_endpoint="https://login.microsoftonline.com/"),
23-
USGOVERNMENT=AzureEnvironment(name="USGOVERNMENT",
24-
service_management_endpoint="https://management.core.usgovcloudapi.net/",
25-
resource_manager_endpoint="https://management.usgovcloudapi.net/",
26-
active_directory_endpoint="https://login.microsoftonline.us/"),
27-
CHINA=AzureEnvironment(name="CHINA",
28-
service_management_endpoint="https://management.core.chinacloudapi.cn/",
29-
resource_manager_endpoint="https://management.chinacloudapi.cn/",
30-
active_directory_endpoint="https://login.chinacloudapi.cn/"),
31-
)
32-
33-
347
def add_workspace_id_header(cfg: 'Config', headers: Dict[str, str]):
358
if cfg.azure_workspace_resource_id:
369
headers["X-Databricks-Azure-Workspace-Resource-Id"] = cfg.azure_workspace_resource_id

‎databricks/sdk/config.py‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,10 @@
1010

1111
import requests
1212

13-
from .azure import AzureEnvironment
1413
from .clock import Clock, RealClock
1514
from .credentials_provider import CredentialsProvider, DefaultCredentials
16-
from .environments import (ALL_ENVS, DEFAULT_ENVIRONMENT, Cloud,
17-
DatabricksEnvironment)
15+
from .environments import (ALL_ENVS, AzureEnvironment, Cloud,
16+
DatabricksEnvironment, get_environment_for_hostname)
1817
from .oauth import OidcEndpoints
1918
from .version import __version__
2019

@@ -154,11 +153,7 @@ def environment(self) -> DatabricksEnvironment:
154153
"""Returns the environment based on configuration."""
155154
if self.databricks_environment:
156155
return self.databricks_environment
157-
if self.host:
158-
for environment in ALL_ENVS:
159-
if self.host.endswith(environment.dns_zone):
160-
return environment
161-
if self.azure_workspace_resource_id:
156+
if not self.host and self.azure_workspace_resource_id:
162157
azure_env = self._get_azure_environment_name()
163158
for environment in ALL_ENVS:
164159
if environment.cloud != Cloud.AZURE:
@@ -168,10 +163,12 @@ def environment(self) -> DatabricksEnvironment:
168163
if environment.dns_zone.startswith(".dev") or environment.dns_zone.startswith(".staging"):
169164
continue
170165
return environment
171-
return DEFAULT_ENVIRONMENT
166+
return get_environment_for_hostname(self.host)
172167

173168
@property
174169
def is_azure(self) -> bool:
170+
if self.azure_workspace_resource_id:
171+
return True
175172
return self.environment.cloud == Cloud.AZURE
176173

177174
@property

‎databricks/sdk/core.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# To preserve backwards compatibility (as these definitions were previously in this module)
1313
from .credentials_provider import *
1414
from .errors import DatabricksError, error_mapper
15+
from .errors.private_link import _is_private_link_redirect
1516
from .retries import retried
1617

1718
__all__ = ['Config', 'DatabricksError']
@@ -239,6 +240,10 @@ def _perform(self,
239240
# See https://stackoverflow.com/a/58821552/277035
240241
payload = response.json()
241242
raise self._make_nicer_error(response=response, **payload) from None
243+
# Private link failures happen via a redirect to the login page. From a requests-perspective, the request
244+
# is successful, but the response is not what we expect. We need to handle this case separately.
245+
if _is_private_link_redirect(response):
246+
raise self._make_nicer_error(response=response) from None
242247
return response
243248
except requests.exceptions.JSONDecodeError:
244249
message = self._make_sense_from_html(response.text)

‎databricks/sdk/environments.py‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,31 @@
22
from enum import Enum
33
from typing import Optional
44

5-
from .azure import ARM_DATABRICKS_RESOURCE_ID, ENVIRONMENTS, AzureEnvironment
5+
6+
@dataclass
7+
class AzureEnvironment:
8+
name: str
9+
service_management_endpoint: str
10+
resource_manager_endpoint: str
11+
active_directory_endpoint: str
12+
13+
14+
ARM_DATABRICKS_RESOURCE_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
15+
16+
ENVIRONMENTS = dict(
17+
PUBLIC=AzureEnvironment(name="PUBLIC",
18+
service_management_endpoint="https://management.core.windows.net/",
19+
resource_manager_endpoint="https://management.azure.com/",
20+
active_directory_endpoint="https://login.microsoftonline.com/"),
21+
USGOVERNMENT=AzureEnvironment(name="USGOVERNMENT",
22+
service_management_endpoint="https://management.core.usgovcloudapi.net/",
23+
resource_manager_endpoint="https://management.usgovcloudapi.net/",
24+
active_directory_endpoint="https://login.microsoftonline.us/"),
25+
CHINA=AzureEnvironment(name="CHINA",
26+
service_management_endpoint="https://management.core.chinacloudapi.cn/",
27+
resource_manager_endpoint="https://management.chinacloudapi.cn/",
28+
active_directory_endpoint="https://login.chinacloudapi.cn/"),
29+
)
630

731

832
class Cloud(Enum):
@@ -70,3 +94,12 @@ def azure_active_directory_endpoint(self) -> Optional[str]:
7094
DatabricksEnvironment(Cloud.GCP, ".staging.gcp.databricks.com"),
7195
DatabricksEnvironment(Cloud.GCP, ".gcp.databricks.com")
7296
]
97+
98+
99+
def get_environment_for_hostname(hostname: str) -> DatabricksEnvironment:
100+
if not hostname:
101+
return DEFAULT_ENVIRONMENT
102+
for env in ALL_ENVS:
103+
if hostname.endswith(env.dns_zone):
104+
return env
105+
return DEFAULT_ENVIRONMENT

‎databricks/sdk/errors/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from .base import DatabricksError, ErrorDetail
22
from .mapper import error_mapper
33
from .platform import *
4+
from .private_link import PrivateLinkValidationError
45
from .sdk import *

‎databricks/sdk/errors/mapper.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
from databricks.sdk.errors.base import DatabricksError
55

66
from .overrides import _ALL_OVERRIDES
7+
from .private_link import (_get_private_link_validation_error,
8+
_is_private_link_redirect)
79

810

911
def error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
@@ -21,6 +23,8 @@ def error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
2123
# where there's a default exception class per HTTP status code, and we do
2224
# rely on Databricks platform exception mapper to do the right thing.
2325
return platform.STATUS_CODE_MAPPING[status_code](**raw)
26+
if _is_private_link_redirect(response):
27+
return _get_private_link_validation_error(response.url)
2428

2529
# backwards-compatible error creation for cases like using older versions of
2630
# the SDK on way never releases of the platform.

‎databricks/sdk/errors/private_link.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
from dataclasses import dataclass
2+
from urllib import parse
3+
4+
import requests
5+
6+
from ..environments import Cloud, get_environment_for_hostname
7+
from .platform import PermissionDenied
8+
9+
10+
@dataclass
11+
class _PrivateLinkInfo:
12+
serviceName: str
13+
endpointName: str
14+
referencePage: str
15+
16+
def error_message(self):
17+
return (
18+
f'The requested workspace has {self.serviceName} enabled and is not accessible from the current network. '
19+
f'Ensure that {self.serviceName} is properly configured and that your device has access to the '
20+
f'{self.endpointName}. For more information, see {self.referencePage}.')
21+
22+
23+
_private_link_info_map = {
24+
Cloud.AWS:
25+
_PrivateLinkInfo(serviceName='AWS PrivateLink',
26+
endpointName='AWS VPC endpoint',
27+
referencePage='https://docs.databricks.com/en/security/network/classic/privatelink.html',
28+
),
29+
Cloud.AZURE:
30+
_PrivateLinkInfo(
31+
serviceName='Azure Private Link',
32+
endpointName='Azure Private Link endpoint',
33+
referencePage='https://learn.microsoft.com/en-us/azure/databricks/security/network/classic/private-link-standard#authentication-troubleshooting',
34+
),
35+
Cloud.GCP:
36+
_PrivateLinkInfo(
37+
serviceName='Private Service Connect',
38+
endpointName='GCP VPC endpoint',
39+
referencePage='https://docs.gcp.databricks.com/en/security/network/classic/private-service-connect.html',
40+
)
41+
}
42+
43+
44+
class PrivateLinkValidationError(PermissionDenied):
45+
"""Raised when a user tries to access a Private Link-enabled workspace, but the user's network does not have access
46+
to the workspace."""
47+
48+
49+
def _is_private_link_redirect(resp: requests.Response) -> bool:
50+
parsed = parse.urlparse(resp.url)
51+
return parsed.path == '/login.html' and 'error=private-link-validation-error' in parsed.query
52+
53+
54+
def _get_private_link_validation_error(url: str) -> _PrivateLinkInfo:
55+
parsed = parse.urlparse(url)
56+
env = get_environment_for_hostname(parsed.hostname)
57+
return PrivateLinkValidationError(message=_private_link_info_map[env.cloud].error_message(),
58+
error_code='PRIVATE_LINK_VALIDATION_ERROR',
59+
status_code=403,
60+
)

‎tests/test_auth.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -196,10 +196,10 @@ def test_config_azure_pat():
196196
def test_config_azure_cli_host(monkeypatch):
197197
monkeypatch.setenv('HOME', __tests__ + '/testdata/azure')
198198
monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
199-
cfg = Config(host='x', azure_workspace_resource_id='/sub/rg/ws')
199+
cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws')
200200

201201
assert cfg.auth_type == 'azure-cli'
202-
assert cfg.host == 'https://x'
202+
assert cfg.host == 'https://adb-123.4.azuredatabricks.net'
203203
assert cfg.is_azure
204204

205205

@@ -232,32 +232,32 @@ def test_config_azure_cli_host_pat_conflict_with_config_file_present_without_def
232232
def test_config_azure_cli_host_and_resource_id(monkeypatch):
233233
monkeypatch.setenv('HOME', __tests__ + '/testdata')
234234
monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
235-
cfg = Config(host='x', azure_workspace_resource_id='/sub/rg/ws')
235+
cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws')
236236

237237
assert cfg.auth_type == 'azure-cli'
238-
assert cfg.host == 'https://x'
238+
assert cfg.host == 'https://adb-123.4.azuredatabricks.net'
239239
assert cfg.is_azure
240240

241241

242242
def test_config_azure_cli_host_and_resource_i_d_configuration_precedence(monkeypatch):
243243
monkeypatch.setenv('DATABRICKS_CONFIG_PROFILE', 'justhost')
244244
monkeypatch.setenv('HOME', __tests__ + '/testdata/azure')
245245
monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
246-
cfg = Config(host='x', azure_workspace_resource_id='/sub/rg/ws')
246+
cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws')
247247

248248
assert cfg.auth_type == 'azure-cli'
249-
assert cfg.host == 'https://x'
249+
assert cfg.host == 'https://adb-123.4.azuredatabricks.net'
250250
assert cfg.is_azure
251251

252252

253253
@raises(
254-
"validate: more than one authorization method configured: azure and basic. Config: host=https://x, username=x, azure_workspace_resource_id=/sub/rg/ws. Env: DATABRICKS_USERNAME"
254+
"validate: more than one authorization method configured: azure and basic. Config: host=https://adb-123.4.azuredatabricks.net, username=x, azure_workspace_resource_id=/sub/rg/ws. Env: DATABRICKS_USERNAME"
255255
)
256256
def test_config_azure_and_password_conflict(monkeypatch):
257257
monkeypatch.setenv('DATABRICKS_USERNAME', 'x')
258258
monkeypatch.setenv('HOME', __tests__ + '/testdata/azure')
259259
monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
260-
cfg = Config(host='x', azure_workspace_resource_id='/sub/rg/ws')
260+
cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws')
261261

262262

263263
@raises(

‎tests/test_auth_manual_tests.py‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@ def test_azure_cli_workspace_header_present(monkeypatch):
77
monkeypatch.setenv('HOME', __tests__ + '/testdata/azure')
88
monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
99
resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
10-
cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id)
10+
cfg = Config(auth_type='azure-cli',
11+
host='https://adb-123.4.azuredatabricks.net',
12+
azure_workspace_resource_id=resource_id)
1113
assert 'X-Databricks-Azure-Workspace-Resource-Id' in cfg.authenticate()
1214
assert cfg.authenticate()['X-Databricks-Azure-Workspace-Resource-Id'] == resource_id
1315

@@ -16,7 +18,9 @@ def test_azure_cli_user_with_management_access(monkeypatch):
1618
monkeypatch.setenv('HOME', __tests__ + '/testdata/azure')
1719
monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
1820
resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
19-
cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id)
21+
cfg = Config(auth_type='azure-cli',
22+
host='https://adb-123.4.azuredatabricks.net',
23+
azure_workspace_resource_id=resource_id)
2024
assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate()
2125

2226

@@ -25,7 +29,9 @@ def test_azure_cli_user_no_management_access(monkeypatch):
2529
monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
2630
monkeypatch.setenv('FAIL_IF', 'https://management.core.windows.net/')
2731
resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
28-
cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id)
32+
cfg = Config(auth_type='azure-cli',
33+
host='https://adb-123.4.azuredatabricks.net',
34+
azure_workspace_resource_id=resource_id)
2935
assert 'X-Databricks-Azure-SP-Management-Token' not in cfg.authenticate()
3036

3137

@@ -34,7 +40,9 @@ def test_azure_cli_fallback(monkeypatch):
3440
monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
3541
monkeypatch.setenv('FAIL_IF', 'subscription')
3642
resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
37-
cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id)
43+
cfg = Config(auth_type='azure-cli',
44+
host='https://adb-123.4.azuredatabricks.net',
45+
azure_workspace_resource_id=resource_id)
3846
assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate()
3947

4048

@@ -43,5 +51,7 @@ def test_azure_cli_with_warning_on_stderr(monkeypatch):
4351
monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
4452
monkeypatch.setenv('WARN', 'this is a warning')
4553
resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
46-
cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id)
54+
cfg = Config(auth_type='azure-cli',
55+
host='https://adb-123.4.azuredatabricks.net',
56+
azure_workspace_resource_id=resource_id)
4757
assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate()

‎tests/test_core.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,14 @@
1414
import requests
1515

1616
from databricks.sdk import WorkspaceClient
17-
from databricks.sdk.azure import ENVIRONMENTS, AzureEnvironment
1817
from databricks.sdk.core import (ApiClient, Config, DatabricksError,
1918
StreamingResponse)
2019
from databricks.sdk.credentials_provider import (CliTokenSource,
2120
CredentialsProvider,
2221
DatabricksCliTokenSource,
2322
HeaderFactory, databricks_cli)
24-
from databricks.sdk.environments import Cloud, DatabricksEnvironment
23+
from databricks.sdk.environments import (ENVIRONMENTS, AzureEnvironment, Cloud,
24+
DatabricksEnvironment)
2525
from databricks.sdk.service.catalog import PermissionsChange
2626
from databricks.sdk.service.iam import AccessControlRequest
2727
from databricks.sdk.version import __version__

0 commit comments

Comments
 (0)