diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py index dabd7280e3d2d..c363c227690ed 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py @@ -30,6 +30,7 @@ import os from typing import TYPE_CHECKING, Any, cast +from azure.core.credentials import AzureSasCredential from azure.core.exceptions import HttpResponseError, ResourceExistsError, ResourceNotFoundError from azure.identity import ClientSecretCredential from azure.identity.aio import ( @@ -203,7 +204,9 @@ def get_conn(self) -> BlobServiceClient: if sas_token: if sas_token.startswith("https"): return BlobServiceClient(account_url=sas_token, **extra) - return BlobServiceClient(account_url=f"{account_url.rstrip('/')}/{sas_token}", **extra) + return BlobServiceClient( + account_url=account_url, credential=AzureSasCredential(sas_token), **extra + ) # Fall back to old auth (password) or use managed identity if not provided. credential: str | TokenCredential | None = conn.password @@ -671,7 +674,7 @@ async def get_async_conn(self) -> AsyncBlobServiceClient: self.blob_service_client = AsyncBlobServiceClient(account_url=sas_token, **extra) else: self.blob_service_client = AsyncBlobServiceClient( - account_url=f"{account_url.rstrip('/')}/{sas_token}", **extra + account_url=account_url, credential=AzureSasCredential(sas_token), **extra ) return self.blob_service_client diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_wasb.py b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_wasb.py index 40f72229950f1..f890cb680e105 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_wasb.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_wasb.py @@ -23,6 +23,7 @@ from unittest.mock import create_autospec import pytest +from azure.core.credentials import AzureSasCredential from azure.core.exceptions import ResourceNotFoundError from azure.storage.blob import BlobServiceClient, ContainerClient from azure.storage.blob._models import BlobProperties @@ -304,10 +305,16 @@ def test_sas_token_provided_and_active_directory_id_used_as_host( self, mocked_connection, mocked_blob_service_client ): WasbHook(wasb_conn_id="testconn").get_conn() - mocked_blob_service_client.assert_called_once_with( - account_url="https://testaccountname.blob.core.windows.net/SAStoken", - sas_token="SAStoken", + assert mocked_blob_service_client.call_args == ( + (), + { + "account_url": "https://testaccountname.blob.core.windows.net/", + "credential": mock.ANY, + }, ) + called_credential = mocked_blob_service_client.call_args[1]["credential"] + assert isinstance(called_credential, AzureSasCredential) + assert called_credential.signature == "SAStoken" @pytest.mark.parametrize( "mocked_connection", @@ -362,9 +369,9 @@ def test_sas_token_connection(self, conn_id_str, extra_key): sas_token = hook_conn.extra_dejson[extra_key] assert isinstance(conn, BlobServiceClient) assert conn.url.startswith("https://") - if hook_conn.login: - assert hook_conn.login in conn.url - assert conn.url.endswith(sas_token + "/") + assert isinstance(conn.credential, AzureSasCredential) + assert conn.credential.signature == sas_token + assert hook_conn.login in conn.url @pytest.mark.parametrize( argnames="conn_id_str",