Skip to content

Commit

Permalink
Support dod in python sdk (#212)
Browse files Browse the repository at this point in the history
## Changes
This PR adds support for the new department of defense accounts
endpoints for Databricks.

## Tests
- [x] Unit tests added for is_account_client.
  • Loading branch information
mgyucht authored Jun 30, 2023
1 parent c27940d commit befbb42
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 1 deletion.
2 changes: 1 addition & 1 deletion databricks/sdk/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
15 changes: 15 additions & 0 deletions tests/test_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit befbb42

Please sign in to comment.