[Fix] Do not specify --tenant flag when fetching managed identity access token from the CLI

databricks/sdk/credentials_provider.py

@@ -437,13 +437,27 @@ def __init__(self, resource: str, subscription: Optional[str] = None, tenant: Op
         if subscription is not None:
             cmd.append("--subscription")
             cmd.append(subscription)
-        if tenant:
+        if tenant and not self.__is_cli_using_managed_identity():
             cmd.extend(["--tenant", tenant])
         super().__init__(cmd=cmd,
                          token_type_field='tokenType',
                          access_token_field='accessToken',
                          expiry_field='expiresOn')
 
+    @staticmethod
+    def __is_cli_using_managed_identity() -> bool:
+        """Checks whether the current CLI session is authenticated using managed identity."""
+        try:
+            out = subprocess.run(["az", "account", "show", "--output", "json"], capture_output=True, check=True)
+            account = json.loads(out.stdout.decode())
+            user = account.get("user")
+            if user is None:
+                return False
+            return user.get("type") == "servicePrincipal" and user.get("name") in ['systemAssignedIdentity', 'userAssignedIdentity']
+        except subprocess.CalledProcessError as e:
+            logger.debug("Failed to get account information from Azure CLI", exc_info=e)
+            return False
+
     def is_human_user(self) -> bool:
         """The UPN claim is the username of the user, but not the Service Principal.
 

tests/test_auth_manual_tests.py

@@ -1,3 +1,5 @@
+import pytest
+
 from databricks.sdk.core import Config
 
 from .conftest import set_az_path, set_home
@@ -60,3 +62,13 @@ def test_azure_cli_with_warning_on_stderr(monkeypatch, mock_tenant):
                  host='https://adb-123.4.azuredatabricks.net',
                  azure_workspace_resource_id=resource_id)
     assert 'X-Databricks-Azure-SP-Management-Token' in cfg.authenticate()
+
+
+@pytest.mark.parametrize('username', ['systemAssignedIdentity', 'userAssignedIdentity'])
+def test_azure_cli_does_not_specify_tenant_id_with_msi(monkeypatch, username):
+    set_home(monkeypatch, '/testdata/azure')
+    set_az_path(monkeypatch)
+    monkeypatch.setenv('FAIL_IF_TENANT_ID_SET', 'true')
+    monkeypatch.setenv('AZ_USER_NAME', username)
+    monkeypatch.setenv('AZ_USER_TYPE', 'servicePrincipal')
+    cfg = Config(auth_type='azure-cli', host='https://adb-123.4.azuredatabricks.net', azure_tenant_id='abc')

tests/testdata/az

@@ -1,7 +1,20 @@
 #!/bin/bash
 
-if [ -n "$WARN" ]; then
-    >&2 /bin/echo "WARNING: ${WARN}"
+# If the arguments are "account show", return the account details.
+if [ "$1" == "account" ] && [ "$2" == "show" ]; then
+    /bin/echo "{
+    \"environmentName\": \"AzureCloud\",
+    \"id\": \"aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee\",
+    \"isDefault\": true,
+    \"name\": \"Pay-As-You-Go\",
+    \"state\": \"Enabled\",
+    \"tenantId\": \"aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee\",
+    \"user\": {
+        \"name\": \"${AZ_USER_NAME:-testuser@databricks.com}\",
+        \"type\": \"${AZ_USER_TYPE:-user}\"
+    }
+}"
+    exit 0
 fi
 
 if [ "yes" == "$FAIL" ]; then
@@ -26,6 +39,21 @@ for arg in "$@"; do
     fi
 done
 
+# Add character to file at $COUNT if it is defined.
+if [ -n "$COUNT" ]; then
+    echo -n x >> "$COUNT"
+fi
+
+# If FAIL_IF_TENANT_ID_SET is set & --tenant-id is passed, fail.
+if [ -n "$FAIL_IF_TENANT_ID_SET" ]; then
+    for arg in "$@"; do
+        if [[ "$arg" == "--tenant" ]]; then
+            echo 1>&2 "ERROR: Tenant shouldn't be specified for managed identity account"
+            exit 1
+        fi
+    done
+fi
+
 # Macos
 EXP="$(/bin/date -v+${EXPIRE:=10S} +'%F %T' 2>/dev/null)"
 if [ -z "${EXP}" ]; then