From befbb42329a42727c530b6ac8101e0ff20be5013 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 30 Jun 2023 18:00:48 +0200 Subject: [PATCH] Support dod in python sdk (#212) ## Changes This PR adds support for the new department of defense accounts endpoints for Databricks. ## Tests - [x] Unit tests added for is_account_client. --- databricks/sdk/core.py | 2 +- tests/test_core.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index 1a619072..0c58d44c 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -566,7 +566,7 @@ def is_aws(self) -> bool: def is_account_client(self) -> bool: if not self.host: return False - return "https://accounts." in self.host + return self.host.startswith("https://accounts.") or self.host.startswith("https://accounts-dod.") @property def arm_environment(self) -> AzureEnvironment: diff --git a/tests/test_core.py b/tests/test_core.py index 142499d9..f00c0b74 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -192,3 +192,18 @@ def test_config_copy_deep_copies_user_agent_other_info(config): config_copy.with_user_agent_extra("test", "test2") assert "test/test2" in config_copy.user_agent assert "test/test2" not in config.user_agent + + +def test_config_accounts_aws_is_accounts_host(config): + config.host = "https://accounts.cloud.databricks.com" + assert config.is_account_client + + +def test_config_accounts_dod_is_accounts_host(config): + config.host = "https://accounts-dod.cloud.databricks.us" + assert config.is_account_client + + +def test_config_workspace_is_not_accounts_host(config): + config.host = "https://westeurope.azuredatabricks.net" + assert not config.is_account_client