diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/fs/msgraph.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/fs/msgraph.py index af3fabe3d6427..f0988f961b6f9 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/fs/msgraph.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/fs/msgraph.py @@ -105,6 +105,12 @@ def get_fs(conn_id: str | None, storage_options: dict[str, Any] | None = None) - if param in options: oauth2_client_params[param] = options[param] + # Construct default token_endpoint from tenant_id if not explicitly provided + if "token_endpoint" not in oauth2_client_params and tenant_id: + oauth2_client_params["token_endpoint"] = ( + f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token" + ) + # Determine which filesystem to return based on drive_id drive_id = options.get("drive_id") diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/fs/test_msgraph.py b/providers/microsoft/azure/tests/unit/microsoft/azure/fs/test_msgraph.py index 235219b6dcb30..d95a772dda296 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/fs/test_msgraph.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/fs/test_msgraph.py @@ -64,10 +64,27 @@ def test_get_fs_with_drive_id(self, mock_msgdrivefs, mock_get_connection, mock_c "client_id": "test_client_id", "client_secret": "test_client_secret", "tenant_id": "test_tenant_id", + "token_endpoint": "https://login.microsoftonline.com/test_tenant_id/oauth2/v2.0/token", }, ) assert result == mock_fs_instance + @patch("airflow.providers.microsoft.azure.fs.msgraph.BaseHook.get_connection") + @patch("msgraphfs.MSGDriveFS") + def test_get_fs_constructs_token_endpoint(self, mock_msgdrivefs, mock_get_connection, mock_connection): + """token_endpoint is auto-constructed from tenant_id when not in extras.""" + mock_get_connection.return_value = mock_connection + mock_fs_instance = MagicMock() + mock_msgdrivefs.return_value = mock_fs_instance + + result = get_fs("msgraph_default") + + actual_params = mock_msgdrivefs.call_args[1]["oauth2_client_params"] + assert actual_params["token_endpoint"] == ( + "https://login.microsoftonline.com/test_tenant_id/oauth2/v2.0/token" + ) + assert result == mock_fs_instance + @patch("msgraphfs.MSGDriveFS") def test_get_fs_no_connection(self, mock_msgdrivefs): mock_fs_instance = MagicMock() @@ -127,6 +144,7 @@ def test_get_fs_with_storage_options(self, mock_msgdrivefs, mock_get_connection, "client_id": "test_client_id", "client_secret": "test_client_secret", "tenant_id": "test_tenant_id", + "token_endpoint": "https://login.microsoftonline.com/test_tenant_id/oauth2/v2.0/token", "scope": "custom.scope", } mock_msgdrivefs.assert_called_once_with(