Skip to content

Commit f0b23a5

Browse files
authored
Support dev and staging workspaces (#514)
## Changes Copies the DatabricksEnvironment abstraction from the Go SDK to the Python SDK. This enables using the SDK for Azure development and staging workspaces, which use an alternate Azure login application ID. Additionally, this PR changes `azure_environment` to accept the same values as the Go SDK and Terraform providers for both Databricks and Azure, i.e. `PUBLIC`, `USGOVERNMENT`, and `CHINA`. The first commit of this PR is a refactor moving config & credentials provider code to separate files, as core.py is getting quite big and difficult to maintain. The second one is much smaller and has just the net new bits. ## Tests - [x] Unit test for the new environment lookup code. - [x] Manual test Azure CLI auth against a staging Azure workspace.
1 parent 26d9eaa commit f0b23a5

File tree

14 files changed

+1216
-1050
lines changed

14 files changed

+1216
-1050
lines changed

.codegen/__init__.py.tmpl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import databricks.sdk.core as client
22
import databricks.sdk.dbutils as dbutils
3+
from databricks.sdk.credentials_provider import CredentialsProvider
34

45
from databricks.sdk.mixins.files import DbfsExt
56
from databricks.sdk.mixins.compute import ClustersExt
@@ -43,7 +44,7 @@ class WorkspaceClient:
4344
debug_headers: bool = None,
4445
product="unknown",
4546
product_version="0.0.0",
46-
credentials_provider: client.CredentialsProvider = None,
47+
credentials_provider: CredentialsProvider = None,
4748
config: client.Config = None):
4849
if not config:
4950
config = client.Config({{range $args}}{{.}}={{.}}, {{end}}
@@ -91,7 +92,7 @@ class AccountClient:
9192
debug_headers: bool = None,
9293
product="unknown",
9394
product_version="0.0.0",
94-
credentials_provider: client.CredentialsProvider = None,
95+
credentials_provider: CredentialsProvider = None,
9596
config: client.Config = None):
9697
if not config:
9798
config = client.Config({{range $args}}{{.}}={{.}}, {{end}}

databricks/sdk/__init__.py

Lines changed: 3 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

databricks/sdk/azure.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,19 +15,15 @@ class AzureEnvironment:
1515
ARM_DATABRICKS_RESOURCE_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
1616

1717
ENVIRONMENTS = dict(
18-
PUBLIC=AzureEnvironment(name="AzurePublicCloud",
18+
PUBLIC=AzureEnvironment(name="PUBLIC",
1919
service_management_endpoint="https://management.core.windows.net/",
2020
resource_manager_endpoint="https://management.azure.com/",
2121
active_directory_endpoint="https://login.microsoftonline.com/"),
22-
GERMAN=AzureEnvironment(name="AzureGermanCloud",
23-
service_management_endpoint="https://management.core.cloudapi.de/",
24-
resource_manager_endpoint="https://management.microsoftazure.de/",
25-
active_directory_endpoint="https://login.microsoftonline.de/"),
26-
USGOVERNMENT=AzureEnvironment(name="AzureUSGovernmentCloud",
22+
USGOVERNMENT=AzureEnvironment(name="USGOVERNMENT",
2723
service_management_endpoint="https://management.core.usgovcloudapi.net/",
2824
resource_manager_endpoint="https://management.usgovcloudapi.net/",
2925
active_directory_endpoint="https://login.microsoftonline.us/"),
30-
CHINA=AzureEnvironment(name="AzureChinaCloud",
26+
CHINA=AzureEnvironment(name="CHINA",
3127
service_management_endpoint="https://management.core.chinacloudapi.cn/",
3228
resource_manager_endpoint="https://management.chinacloudapi.cn/",
3329
active_directory_endpoint="https://login.chinacloudapi.cn/"),

0 commit comments

Comments
 (0)