From 741049fea1dd5177f9d3c52c471eb349089b256e Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 17 May 2024 11:34:52 +0200 Subject: [PATCH 1/5] Private link --- databricks/sdk/azure.py | 26 ------------ databricks/sdk/config.py | 13 ++---- databricks/sdk/environments.py | 33 ++++++++++++++- databricks/sdk/errors/__init__.py | 1 + databricks/sdk/errors/mapper.py | 5 +++ databricks/sdk/errors/private_link.py | 58 +++++++++++++++++++++++++++ tests/test_errors.py | 8 ++++ 7 files changed, 108 insertions(+), 36 deletions(-) create mode 100644 databricks/sdk/errors/private_link.py diff --git a/databricks/sdk/azure.py b/databricks/sdk/azure.py index ec084cf22..62ffce756 100644 --- a/databricks/sdk/azure.py +++ b/databricks/sdk/azure.py @@ -5,32 +5,6 @@ from .service.provisioning import Workspace -@dataclass -class AzureEnvironment: - name: str - service_management_endpoint: str - resource_manager_endpoint: str - active_directory_endpoint: str - - -ARM_DATABRICKS_RESOURCE_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d" - -ENVIRONMENTS = dict( - PUBLIC=AzureEnvironment(name="PUBLIC", - service_management_endpoint="https://management.core.windows.net/", - resource_manager_endpoint="https://management.azure.com/", - active_directory_endpoint="https://login.microsoftonline.com/"), - USGOVERNMENT=AzureEnvironment(name="USGOVERNMENT", - service_management_endpoint="https://management.core.usgovcloudapi.net/", - resource_manager_endpoint="https://management.usgovcloudapi.net/", - active_directory_endpoint="https://login.microsoftonline.us/"), - CHINA=AzureEnvironment(name="CHINA", - service_management_endpoint="https://management.core.chinacloudapi.cn/", - resource_manager_endpoint="https://management.chinacloudapi.cn/", - active_directory_endpoint="https://login.chinacloudapi.cn/"), -) - - def add_workspace_id_header(cfg: 'Config', headers: Dict[str, str]): if cfg.azure_workspace_resource_id: headers["X-Databricks-Azure-Workspace-Resource-Id"] = cfg.azure_workspace_resource_id diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index a5763e232..d084da018 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -10,11 +10,10 @@ import requests -from .azure import AzureEnvironment from .clock import Clock, RealClock from .credentials_provider import CredentialsProvider, DefaultCredentials -from .environments import (ALL_ENVS, DEFAULT_ENVIRONMENT, Cloud, - DatabricksEnvironment) +from .environments import (ALL_ENVS, AzureEnvironment, Cloud, + DatabricksEnvironment, get_environment_for_hostname) from .oauth import OidcEndpoints from .version import __version__ @@ -154,11 +153,7 @@ def environment(self) -> DatabricksEnvironment: """Returns the environment based on configuration.""" if self.databricks_environment: return self.databricks_environment - if self.host: - for environment in ALL_ENVS: - if self.host.endswith(environment.dns_zone): - return environment - if self.azure_workspace_resource_id: + if not self.host and self.azure_workspace_resource_id: azure_env = self._get_azure_environment_name() for environment in ALL_ENVS: if environment.cloud != Cloud.AZURE: @@ -168,7 +163,7 @@ def environment(self) -> DatabricksEnvironment: if environment.dns_zone.startswith(".dev") or environment.dns_zone.startswith(".staging"): continue return environment - return DEFAULT_ENVIRONMENT + return get_environment_for_hostname(self.host) @property def is_azure(self) -> bool: diff --git a/databricks/sdk/environments.py b/databricks/sdk/environments.py index 78aa4655d..644eba6f5 100644 --- a/databricks/sdk/environments.py +++ b/databricks/sdk/environments.py @@ -2,7 +2,31 @@ from enum import Enum from typing import Optional -from .azure import ARM_DATABRICKS_RESOURCE_ID, ENVIRONMENTS, AzureEnvironment + +@dataclass +class AzureEnvironment: + name: str + service_management_endpoint: str + resource_manager_endpoint: str + active_directory_endpoint: str + + +ARM_DATABRICKS_RESOURCE_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d" + +ENVIRONMENTS = dict( + PUBLIC=AzureEnvironment(name="PUBLIC", + service_management_endpoint="https://management.core.windows.net/", + resource_manager_endpoint="https://management.azure.com/", + active_directory_endpoint="https://login.microsoftonline.com/"), + USGOVERNMENT=AzureEnvironment(name="USGOVERNMENT", + service_management_endpoint="https://management.core.usgovcloudapi.net/", + resource_manager_endpoint="https://management.usgovcloudapi.net/", + active_directory_endpoint="https://login.microsoftonline.us/"), + CHINA=AzureEnvironment(name="CHINA", + service_management_endpoint="https://management.core.chinacloudapi.cn/", + resource_manager_endpoint="https://management.chinacloudapi.cn/", + active_directory_endpoint="https://login.chinacloudapi.cn/"), +) class Cloud(Enum): @@ -70,3 +94,10 @@ def azure_active_directory_endpoint(self) -> Optional[str]: DatabricksEnvironment(Cloud.GCP, ".staging.gcp.databricks.com"), DatabricksEnvironment(Cloud.GCP, ".gcp.databricks.com") ] + + +def get_environment_for_hostname(hostname: str) -> DatabricksEnvironment: + for env in ALL_ENVS: + if hostname.endswith(env.dns_zone): + return env + return DEFAULT_ENVIRONMENT diff --git a/databricks/sdk/errors/__init__.py b/databricks/sdk/errors/__init__.py index d6ea05cc7..35d4f16a4 100644 --- a/databricks/sdk/errors/__init__.py +++ b/databricks/sdk/errors/__init__.py @@ -2,3 +2,4 @@ from .mapper import error_mapper from .platform import * from .sdk import * +from .private_link import PrivateLinkValidationError diff --git a/databricks/sdk/errors/mapper.py b/databricks/sdk/errors/mapper.py index c4ba0d55a..52ebacdda 100644 --- a/databricks/sdk/errors/mapper.py +++ b/databricks/sdk/errors/mapper.py @@ -4,6 +4,7 @@ from databricks.sdk.errors.base import DatabricksError from .overrides import _ALL_OVERRIDES +from .private_link import _is_private_link_redirect, _get_private_link_validation_error def error_mapper(response: requests.Response, raw: dict) -> DatabricksError: @@ -21,7 +22,11 @@ def error_mapper(response: requests.Response, raw: dict) -> DatabricksError: # where there's a default exception class per HTTP status code, and we do # rely on Databricks platform exception mapper to do the right thing. return platform.STATUS_CODE_MAPPING[status_code](**raw) + if _is_private_link_redirect(response): + return _get_private_link_validation_error(response.url) # backwards-compatible error creation for cases like using older versions of # the SDK on way never releases of the platform. return DatabricksError(**raw) + + diff --git a/databricks/sdk/errors/private_link.py b/databricks/sdk/errors/private_link.py new file mode 100644 index 000000000..f01fc9bbb --- /dev/null +++ b/databricks/sdk/errors/private_link.py @@ -0,0 +1,58 @@ +from urllib import parse +from dataclasses import dataclass + +import requests + +from .platform import PermissionDenied +from ..environments import Cloud, get_environment_for_hostname + +@dataclass +class _PrivateLinkInfo: + serviceName: str + endpointName: str + referencePage: str + + def error_message(self): + return ( + f'The requested workspace has {self.serviceName} enabled and is not accessible from the current network. ' + f'Ensure that {self.serviceName} is properly configured and that your device has access to the ' + f'{self.endpointName}. For more information, see {self.referencePage}.' + ) + +_private_link_info_map = { + Cloud.AWS: _PrivateLinkInfo( + serviceName='AWS PrivateLink', + endpointName='AWS VPC endpoint', + referencePage='https://docs.databricks.com/en/security/network/classic/privatelink.html', + ), + Cloud.AZURE: _PrivateLinkInfo( + serviceName='Azure Private Link', + endpointName='Azure Private Link endpoint', + referencePage='https://learn.microsoft.com/en-us/azure/databricks/security/network/classic/private-link-standard#authentication-troubleshooting', + ), + Cloud.GCP: _PrivateLinkInfo( + serviceName='Private Service Connect', + endpointName='GCP VPC endpoint', + referencePage='https://docs.gcp.databricks.com/en/security/network/classic/private-service-connect.html', + ) +} + + +class PrivateLinkValidationError(PermissionDenied): + """Raised when a user tries to access a Private Link-enabled workspace, but the user's network does not have access + to the workspace.""" + + +def _is_private_link_redirect(resp: requests.Response) -> bool: + parsed = parse.urlparse(resp.url) + return parsed.path == '/login.html' and 'error=private-link-validation-error' in parsed.query + + +def _get_private_link_validation_error(url: str) -> _PrivateLinkInfo: + parsed = parse.urlparse(url) + env = get_environment_for_hostname(parsed.hostname) + return PrivateLinkValidationError( + message=_private_link_info_map[env.cloud].error_message(), + error_code='PRIVATE_LINK_VALIDATION_ERROR', + status_code=403, + ) diff --git a/tests/test_errors.py b/tests/test_errors.py index d9243615a..87111c4b4 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -34,6 +34,14 @@ def test_missing_error_code(): assert errors.DatabricksError == type(err) +def test_private_link_error(): + resp = requests.Response() + resp.url = 'https://databricks.com/login.html?error=private-link-validation-error' + resp.request = requests.Request('GET', 'https://databricks.com/api/2.0/service').prepare() + err = errors.error_mapper(resp, {}) + assert errors.PrivateLinkValidationError == type(err) + + @pytest.mark.parametrize('status_code, error_code, klass', [(400, ..., errors.BadRequest), (400, 'INVALID_PARAMETER_VALUE', errors.BadRequest), (400, 'INVALID_PARAMETER_VALUE', errors.InvalidParameterValue), From 0e245e74d218d34b20ad9c25d41b0bb414f740d9 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 17 May 2024 11:41:33 +0200 Subject: [PATCH 2/5] trigger error mapping --- databricks/sdk/core.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index ebf3aa185..af07a69f9 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -12,6 +12,7 @@ # To preserve backwards compatibility (as these definitions were previously in this module) from .credentials_provider import * from .errors import DatabricksError, error_mapper +from .errors.private_link import _is_private_link_redirect from .retries import retried __all__ = ['Config', 'DatabricksError'] @@ -238,7 +239,10 @@ def _perform(self, # TODO: experiment with traceback pruning for better readability # See https://stackoverflow.com/a/58821552/277035 payload = response.json() - raise self._make_nicer_error(response=response, **payload) from None + # Private link failures happen via a redirect to the login page. From a requests-perspective, the request + # is successful, but the response is not what we expect. We need to handle this case separately. + if _is_private_link_redirect(response): + raise self._make_nicer_error(response=response) from None return response except requests.exceptions.JSONDecodeError: message = self._make_sense_from_html(response.text) From 453c827e0af9d32c87d7f139452559930c45ee25 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 17 May 2024 11:44:31 +0200 Subject: [PATCH 3/5] fmt --- databricks/sdk/azure.py | 1 - databricks/sdk/core.py | 2 +- databricks/sdk/errors/__init__.py | 2 +- databricks/sdk/errors/mapper.py | 5 ++-- databricks/sdk/errors/private_link.py | 34 ++++++++++++++------------- 5 files changed, 22 insertions(+), 22 deletions(-) diff --git a/databricks/sdk/azure.py b/databricks/sdk/azure.py index 62ffce756..372669d61 100644 --- a/databricks/sdk/azure.py +++ b/databricks/sdk/azure.py @@ -1,4 +1,3 @@ -from dataclasses import dataclass from typing import Dict from .oauth import TokenSource diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index af07a69f9..3656cdded 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -238,7 +238,7 @@ def _perform(self, if not response.ok: # internally calls response.raise_for_status() # TODO: experiment with traceback pruning for better readability # See https://stackoverflow.com/a/58821552/277035 - payload = response.json() + response.json() # Private link failures happen via a redirect to the login page. From a requests-perspective, the request # is successful, but the response is not what we expect. We need to handle this case separately. if _is_private_link_redirect(response): diff --git a/databricks/sdk/errors/__init__.py b/databricks/sdk/errors/__init__.py index 35d4f16a4..749c95116 100644 --- a/databricks/sdk/errors/__init__.py +++ b/databricks/sdk/errors/__init__.py @@ -1,5 +1,5 @@ from .base import DatabricksError, ErrorDetail from .mapper import error_mapper from .platform import * -from .sdk import * from .private_link import PrivateLinkValidationError +from .sdk import * diff --git a/databricks/sdk/errors/mapper.py b/databricks/sdk/errors/mapper.py index 52ebacdda..0b809eb7e 100644 --- a/databricks/sdk/errors/mapper.py +++ b/databricks/sdk/errors/mapper.py @@ -4,7 +4,8 @@ from databricks.sdk.errors.base import DatabricksError from .overrides import _ALL_OVERRIDES -from .private_link import _is_private_link_redirect, _get_private_link_validation_error +from .private_link import (_get_private_link_validation_error, + _is_private_link_redirect) def error_mapper(response: requests.Response, raw: dict) -> DatabricksError: @@ -28,5 +29,3 @@ def error_mapper(response: requests.Response, raw: dict) -> DatabricksError: # backwards-compatible error creation for cases like using older versions of # the SDK on way never releases of the platform. return DatabricksError(**raw) - - diff --git a/databricks/sdk/errors/private_link.py b/databricks/sdk/errors/private_link.py index f01fc9bbb..e8cc5eadf 100644 --- a/databricks/sdk/errors/private_link.py +++ b/databricks/sdk/errors/private_link.py @@ -1,10 +1,11 @@ -from urllib import parse from dataclasses import dataclass +from urllib import parse import requests -from .platform import PermissionDenied from ..environments import Cloud, get_environment_for_hostname +from .platform import PermissionDenied + @dataclass class _PrivateLinkInfo: @@ -16,21 +17,23 @@ def error_message(self): return ( f'The requested workspace has {self.serviceName} enabled and is not accessible from the current network. ' f'Ensure that {self.serviceName} is properly configured and that your device has access to the ' - f'{self.endpointName}. For more information, see {self.referencePage}.' - ) + f'{self.endpointName}. For more information, see {self.referencePage}.') + _private_link_info_map = { - Cloud.AWS: _PrivateLinkInfo( - serviceName='AWS PrivateLink', - endpointName='AWS VPC endpoint', - referencePage='https://docs.databricks.com/en/security/network/classic/privatelink.html', - ), - Cloud.AZURE: _PrivateLinkInfo( + Cloud.AWS: + _PrivateLinkInfo(serviceName='AWS PrivateLink', + endpointName='AWS VPC endpoint', + referencePage='https://docs.databricks.com/en/security/network/classic/privatelink.html', + ), + Cloud.AZURE: + _PrivateLinkInfo( serviceName='Azure Private Link', endpointName='Azure Private Link endpoint', referencePage='https://learn.microsoft.com/en-us/azure/databricks/security/network/classic/private-link-standard#authentication-troubleshooting', ), - Cloud.GCP: _PrivateLinkInfo( + Cloud.GCP: + _PrivateLinkInfo( serviceName='Private Service Connect', endpointName='GCP VPC endpoint', referencePage='https://docs.gcp.databricks.com/en/security/network/classic/private-service-connect.html', @@ -51,8 +54,7 @@ def _is_private_link_redirect(resp: requests.Response) -> bool: def _get_private_link_validation_error(url: str) -> _PrivateLinkInfo: parsed = parse.urlparse(url) env = get_environment_for_hostname(parsed.hostname) - return PrivateLinkValidationError( - message=_private_link_info_map[env.cloud].error_message(), - error_code='PRIVATE_LINK_VALIDATION_ERROR', - status_code=403, - ) + return PrivateLinkValidationError(message=_private_link_info_map[env.cloud].error_message(), + error_code='PRIVATE_LINK_VALIDATION_ERROR', + status_code=403, + ) From 6422cb51b81e913fb836cbab13013c6a5cdbde73 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 17 May 2024 15:19:11 +0200 Subject: [PATCH 4/5] work --- databricks/sdk/config.py | 2 ++ databricks/sdk/core.py | 3 ++- databricks/sdk/environments.py | 2 ++ tests/test_auth.py | 16 ++++++++-------- tests/test_auth_manual_tests.py | 10 +++++----- tests/test_core.py | 3 +-- 6 files changed, 20 insertions(+), 16 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index d084da018..d7d45c20d 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -167,6 +167,8 @@ def environment(self) -> DatabricksEnvironment: @property def is_azure(self) -> bool: + if self.azure_workspace_resource_id: + return True return self.environment.cloud == Cloud.AZURE @property diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index 3656cdded..658ed2190 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -238,7 +238,8 @@ def _perform(self, if not response.ok: # internally calls response.raise_for_status() # TODO: experiment with traceback pruning for better readability # See https://stackoverflow.com/a/58821552/277035 - response.json() + payload = response.json() + raise self._make_nicer_error(response=response, **payload) from None # Private link failures happen via a redirect to the login page. From a requests-perspective, the request # is successful, but the response is not what we expect. We need to handle this case separately. if _is_private_link_redirect(response): diff --git a/databricks/sdk/environments.py b/databricks/sdk/environments.py index 644eba6f5..10d647708 100644 --- a/databricks/sdk/environments.py +++ b/databricks/sdk/environments.py @@ -97,6 +97,8 @@ def azure_active_directory_endpoint(self) -> Optional[str]: def get_environment_for_hostname(hostname: str) -> DatabricksEnvironment: + if not hostname: + return DEFAULT_ENVIRONMENT for env in ALL_ENVS: if hostname.endswith(env.dns_zone): return env diff --git a/tests/test_auth.py b/tests/test_auth.py index 504e14439..02535c39e 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -196,10 +196,10 @@ def test_config_azure_pat(): def test_config_azure_cli_host(monkeypatch): monkeypatch.setenv('HOME', __tests__ + '/testdata/azure') monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') - cfg = Config(host='x', azure_workspace_resource_id='/sub/rg/ws') + cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws') assert cfg.auth_type == 'azure-cli' - assert cfg.host == 'https://x' + assert cfg.host == 'https://adb-123.4.azuredatabricks.net' assert cfg.is_azure @@ -232,10 +232,10 @@ def test_config_azure_cli_host_pat_conflict_with_config_file_present_without_def def test_config_azure_cli_host_and_resource_id(monkeypatch): monkeypatch.setenv('HOME', __tests__ + '/testdata') monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') - cfg = Config(host='x', azure_workspace_resource_id='/sub/rg/ws') + cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws') assert cfg.auth_type == 'azure-cli' - assert cfg.host == 'https://x' + assert cfg.host == 'https://adb-123.4.azuredatabricks.net' assert cfg.is_azure @@ -243,21 +243,21 @@ def test_config_azure_cli_host_and_resource_i_d_configuration_precedence(monkeyp monkeypatch.setenv('DATABRICKS_CONFIG_PROFILE', 'justhost') monkeypatch.setenv('HOME', __tests__ + '/testdata/azure') monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') - cfg = Config(host='x', azure_workspace_resource_id='/sub/rg/ws') + cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws') assert cfg.auth_type == 'azure-cli' - assert cfg.host == 'https://x' + assert cfg.host == 'https://adb-123.4.azuredatabricks.net' assert cfg.is_azure @raises( - "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" + "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" ) def test_config_azure_and_password_conflict(monkeypatch): monkeypatch.setenv('DATABRICKS_USERNAME', 'x') monkeypatch.setenv('HOME', __tests__ + '/testdata/azure') monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') - cfg = Config(host='x', azure_workspace_resource_id='/sub/rg/ws') + cfg = Config(host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id='/sub/rg/ws') @raises( diff --git a/tests/test_auth_manual_tests.py b/tests/test_auth_manual_tests.py index e7562e266..cf4fff4cb 100644 --- a/tests/test_auth_manual_tests.py +++ b/tests/test_auth_manual_tests.py @@ -7,7 +7,7 @@ def test_azure_cli_workspace_header_present(monkeypatch): monkeypatch.setenv('HOME', __tests__ + '/testdata/azure') monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-Workspace-Resource-Id' in cfg.authenticate() assert cfg.authenticate()['X-Databricks-Azure-Workspace-Resource-Id'] == resource_id @@ -16,7 +16,7 @@ def test_azure_cli_user_with_management_access(monkeypatch): monkeypatch.setenv('HOME', __tests__ + '/testdata/azure') monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate() @@ -25,7 +25,7 @@ def test_azure_cli_user_no_management_access(monkeypatch): monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') monkeypatch.setenv('FAIL_IF', 'https://management.core.windows.net/') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-SP-Management-Token' not in cfg.authenticate() @@ -34,7 +34,7 @@ def test_azure_cli_fallback(monkeypatch): monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') monkeypatch.setenv('FAIL_IF', 'subscription') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate() @@ -43,5 +43,5 @@ def test_azure_cli_with_warning_on_stderr(monkeypatch): monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') monkeypatch.setenv('WARN', 'this is a warning') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate() diff --git a/tests/test_core.py b/tests/test_core.py index 3ca6ca830..e000b13c2 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -14,14 +14,13 @@ import requests from databricks.sdk import WorkspaceClient -from databricks.sdk.azure import ENVIRONMENTS, AzureEnvironment from databricks.sdk.core import (ApiClient, Config, DatabricksError, StreamingResponse) from databricks.sdk.credentials_provider import (CliTokenSource, CredentialsProvider, DatabricksCliTokenSource, HeaderFactory, databricks_cli) -from databricks.sdk.environments import Cloud, DatabricksEnvironment +from databricks.sdk.environments import Cloud, DatabricksEnvironment, ENVIRONMENTS, AzureEnvironment from databricks.sdk.service.catalog import PermissionsChange from databricks.sdk.service.iam import AccessControlRequest from databricks.sdk.version import __version__ From bf1644cb363a798a98c8fd0c9d8430f8c38925de Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 17 May 2024 15:59:32 +0200 Subject: [PATCH 5/5] fmt --- tests/test_auth_manual_tests.py | 20 +++++++++++++++----- tests/test_core.py | 3 ++- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/tests/test_auth_manual_tests.py b/tests/test_auth_manual_tests.py index cf4fff4cb..07250c532 100644 --- a/tests/test_auth_manual_tests.py +++ b/tests/test_auth_manual_tests.py @@ -7,7 +7,9 @@ def test_azure_cli_workspace_header_present(monkeypatch): monkeypatch.setenv('HOME', __tests__ + '/testdata/azure') monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', + host='https://adb-123.4.azuredatabricks.net', + azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-Workspace-Resource-Id' in cfg.authenticate() assert cfg.authenticate()['X-Databricks-Azure-Workspace-Resource-Id'] == resource_id @@ -16,7 +18,9 @@ def test_azure_cli_user_with_management_access(monkeypatch): monkeypatch.setenv('HOME', __tests__ + '/testdata/azure') monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', + host='https://adb-123.4.azuredatabricks.net', + azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate() @@ -25,7 +29,9 @@ def test_azure_cli_user_no_management_access(monkeypatch): monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') monkeypatch.setenv('FAIL_IF', 'https://management.core.windows.net/') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', + host='https://adb-123.4.azuredatabricks.net', + azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-SP-Management-Token' not in cfg.authenticate() @@ -34,7 +40,9 @@ def test_azure_cli_fallback(monkeypatch): monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') monkeypatch.setenv('FAIL_IF', 'subscription') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', + host='https://adb-123.4.azuredatabricks.net', + azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate() @@ -43,5 +51,7 @@ def test_azure_cli_with_warning_on_stderr(monkeypatch): monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin') monkeypatch.setenv('WARN', 'this is a warning') resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123' - cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_workspace_resource_id=resource_id) + cfg = Config(auth_type='azure-cli', + host='https://adb-123.4.azuredatabricks.net', + azure_workspace_resource_id=resource_id) assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate() diff --git a/tests/test_core.py b/tests/test_core.py index e000b13c2..fd2f12ad4 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -20,7 +20,8 @@ CredentialsProvider, DatabricksCliTokenSource, HeaderFactory, databricks_cli) -from databricks.sdk.environments import Cloud, DatabricksEnvironment, ENVIRONMENTS, AzureEnvironment +from databricks.sdk.environments import (ENVIRONMENTS, AzureEnvironment, Cloud, + DatabricksEnvironment) from databricks.sdk.service.catalog import PermissionsChange from databricks.sdk.service.iam import AccessControlRequest from databricks.sdk.version import __version__