Skip to content
Permalink
Browse files
wasb hook: user defaultAzureCredentials instead of managedIdentity (#…
…23394)

Co-authored-by: Sanjay Pillai <sanjaypillai11 [at] gmail.com>
  • Loading branch information
snjypl committed May 8, 2022
1 parent fcfaa83 commit 8f181c10344bd319ac5f6aeb102ee3c06e1f1637
Showing 2 changed files with 6 additions and 6 deletions.
@@ -29,7 +29,7 @@
from typing import Any, Dict, List, Optional

from azure.core.exceptions import HttpResponseError, ResourceExistsError, ResourceNotFoundError
from azure.identity import ClientSecretCredential, ManagedIdentityCredential
from azure.identity import ClientSecretCredential, DefaultAzureCredential
from azure.storage.blob import BlobClient, BlobServiceClient, ContainerClient, StorageStreamDownloader

from airflow.exceptions import AirflowException
@@ -46,7 +46,7 @@ class WasbHook(BaseHook):
passed to the `BlockBlockService()` constructor. For example, authenticate
using a SAS token by adding {"sas_token": "YOUR_TOKEN"}.
If no authentication configuration is provided, managed identity will be used (applicable
If no authentication configuration is provided, DefaultAzureCredential will be used (applicable
when using Azure compute infrastructure).
:param wasb_conn_id: Reference to the :ref:`wasb connection <howto/connection:wasb>`.
@@ -151,8 +151,8 @@ def get_conn(self) -> BlobServiceClient:
# Fall back to old auth (password) or use managed identity if not provided.
credential = conn.password
if not credential:
credential = ManagedIdentityCredential()
self.log.info("Using managed identity as credential")
credential = DefaultAzureCredential()
self.log.info("Using DefaultAzureCredential as credential")
return BlobServiceClient(
account_url=f"https://{conn.login}.blob.core.windows.net/",
credential=credential,
@@ -22,7 +22,7 @@
from unittest import mock

import pytest
from azure.identity import ManagedIdentityCredential
from azure.identity import DefaultAzureCredential
from azure.storage.blob import BlobServiceClient

from airflow.exceptions import AirflowException
@@ -140,7 +140,7 @@ def test_shared_key_connection(self):
def test_managed_identity(self):
hook = WasbHook(wasb_conn_id=self.managed_identity_conn_id)
assert isinstance(hook.get_conn(), BlobServiceClient)
assert isinstance(hook.get_conn().credential, ManagedIdentityCredential)
assert isinstance(hook.get_conn().credential, DefaultAzureCredential)

@pytest.mark.parametrize(
argnames="conn_id_str, extra_key",

0 comments on commit 8f181c1

Please sign in to comment.