Fix Databricks OAuth M2M on Azure (#513)

## Changes
Databricks OAuth M2M is now also available on Azure. This PR fixes OAuth
for Azure.

## Tests
* manually tested with an Azure Workspace
*
https://github.com/databricks/eng-dev-ecosystem/actions/runs/7568420047/job/20609591977

- [x] `make test` run locally
- [x] `make fmt` applied
- [x] relevant integration tests applied

README.md

@@ -151,12 +151,12 @@ The Databricks SDK for Python picks up an Azure CLI token, if you've previously
 
 To authenticate as an Azure Active Directory (Azure AD) service principal, you must provide one of the following. See also [Add a service principal to your Azure Databricks account](https://learn.microsoft.com/azure/databricks/administration-guide/users-groups/service-principals#add-sp-account):
 
-- `azure_resource_id`, `azure_client_secret`, `azure_client_id`, and `azure_tenant_id`; or their environment variable or `.databrickscfg` file field equivalents.
-- `azure_resource_id` and `azure_use_msi`; or their environment variable or `.databrickscfg` file field equivalents.
+- `azure_workspace_resource_id`, `azure_client_secret`, `azure_client_id`, and `azure_tenant_id`; or their environment variable or `.databrickscfg` file field equivalents.
+- `azure_workspace_resource_id` and `azure_use_msi`; or their environment variable or `.databrickscfg` file field equivalents.
 
 | Argument              | Description | Environment variable |
 |-----------------------|-------------|----------------------|
-| `azure_resource_id`   | _(String)_ The Azure Resource Manager ID for the Azure Databricks workspace, which is exchanged for a Databricks host URL. | `DATABRICKS_AZURE_RESOURCE_ID` |
+| `azure_workspace_resource_id`   | _(String)_ The Azure Resource Manager ID for the Azure Databricks workspace, which is exchanged for a Databricks host URL. | `DATABRICKS_AZURE_RESOURCE_ID` |
 | `azure_use_msi`       | _(Boolean)_ `true` to use Azure Managed Service Identity passwordless authentication flow for service principals. _This feature is not yet implemented in the Databricks SDK for Python._ | `ARM_USE_MSI` |
 | `azure_client_secret` | _(String)_ The Azure AD service principal's client secret. | `ARM_CLIENT_SECRET` |
 | `azure_client_id`     | _(String)_ The Azure AD service principal's application ID. | `ARM_CLIENT_ID` |

databricks/sdk/core.py

@@ -123,12 +123,10 @@ def runtime_native_auth(cfg: 'Config') -> Optional[HeaderFactory]:
     return None
 
 
-@credentials_provider('oauth-m2m', ['is_aws', 'host', 'client_id', 'client_secret'])
+@credentials_provider('oauth-m2m', ['host', 'client_id', 'client_secret'])
 def oauth_service_principal(cfg: 'Config') -> Optional[HeaderFactory]:
     """ Adds refreshed Databricks machine-to-machine OAuth Bearer token to every request,
     if /oidc/.well-known/oauth-authorization-server is available on the given host. """
-    # TODO: Azure returns 404 for UC workspace after redirecting to
-    # https://login.microsoftonline.com/{cfg.azure_tenant_id}/.well-known/oauth-authorization-server
     oidc = cfg.oidc_endpoints
     if oidc is None:
         return None
@@ -838,7 +836,7 @@ def oidc_endpoints(self) -> Optional[OidcEndpoints]:
         self._fix_host_if_needed()
         if not self.host:
             return None
-        if self.is_azure:
+        if self.is_azure and self.azure_client_id:
             # Retrieve authorize endpoint to retrieve token endpoint after
             res = requests.get(f'{self.host}/oidc/oauth2/v2.0/authorize', allow_redirects=False)
             real_auth_url = res.headers.get('location')

docs/authentication.md

@@ -84,13 +84,13 @@ The Databricks SDK for Python picks up an Azure CLI token, if you've previously
 
 To authenticate as an Azure Active Directory (Azure AD) service principal, you must provide one of the following. See also [Add a service principal to your Azure Databricks account](https://learn.microsoft.com/azure/databricks/administration-guide/users-groups/service-principals#add-sp-account):
 
-- `azure_resource_id`, `azure_client_secret`, `azure_client_id`, and `azure_tenant_id`; or their environment variable or `.databrickscfg` file field equivalents.
-- `azure_resource_id` and `azure_use_msi`; or their environment variable or `.databrickscfg` file field equivalents.
+- `azure_workspace_resource_id`, `azure_client_secret`, `azure_client_id`, and `azure_tenant_id`; or their environment variable or `.databrickscfg` file field equivalents.
+- `azure_workspace_resource_id` and `azure_use_msi`; or their environment variable or `.databrickscfg` file field equivalents.
 
 | Argument              | Description | Environment variable |
 |-----------------------|-------------|----------------------|
-| `azure_resource_id`   | _(String)_ The Azure Resource Manager ID for the Azure Databricks workspace, which is exchanged for a Databricks host URL. | `DATABRICKS_AZURE_RESOURCE_ID` |
-| `azure_use_msi`       | _(Boolean)_ `true` to use Azure Managed Service Identity passwordless authentication flow for service principals. _This feature is not yet implemented in the Databricks SDK for Python._ | `ARM_USE_MSI` |
+| `azure_workspace_resource_id`   | _(String)_ The Azure Resource Manager ID for the Azure Databricks workspace, which is exchanged for a Databricks host URL. | `DATABRICKS_AZURE_RESOURCE_ID` |
+| `azure_use_msi`       | _(Boolean)_ `true` to use Azure Managed Service Identity passwordless authentication flow for service principals. | `ARM_USE_MSI` |
 | `azure_client_secret` | _(String)_ The Azure AD service principal's client secret. | `ARM_CLIENT_SECRET` |
 | `azure_client_id`     | _(String)_ The Azure AD service principal's application ID. | `ARM_CLIENT_ID` |
 | `azure_tenant_id`     | _(String)_ The Azure AD service principal's tenant ID. | `ARM_TENANT_ID` |