azure_databricks_workspace_creation_with_default_credential.py

1import os
2from azure.identity import DefaultAzureCredential
3from azure.mgmt.databricks import AzureDatabricksManagementClient
4
5def main():
6    # Subscription ID and Resource Group details
7    SUBSCRIPTION_ID = os.environ.get("AZURE_SUBSCRIPTION_ID", "your-subscription-id")
8    RESOURCE_GROUP_NAME = "myResourceGroup"
9    WORKSPACE_NAME = "myDatabricksWorkspace"
10    LOCATION = "eastus"
11
12    # Acquire a credential object using CLI/Environment authentication
13    credential = DefaultAzureCredential()
14
15    # Initialize the management client
16    databricks_client = AzureDatabricksManagementClient(credential, SUBSCRIPTION_ID)
17
18    # Provision a Databricks workspace
19    # Standard pricing tier (sku) is used in this example
20    print(f"Creating workspace: {WORKSPACE_NAME}...")
21    
22    poller = databricks_client.workspaces.begin_create_or_update(
23        RESOURCE_GROUP_NAME,
24        WORKSPACE_NAME,
25        {
26            "location": LOCATION,
27            "sku": {"name": "standard"},
28            "managed_resource_group_id": f"/subscriptions/{SUBSCRIPTION_ID}/resourceGroups/{WORKSPACE_NAME}-managed-rg"
29        }
30    )
31
32    workspace_result = poller.result()
33    print(f"Workspace created successfully: {workspace_result.name}")
34
35if __name__ == "__main__":
36    main()