From 44805acc5014f4aa1ce30da31e4540e1dc348d58 Mon Sep 17 00:00:00 2001
From: Alexander Streed <desertaxle@users.noreply.github.com>
Date: Thu, 22 Jun 2023 12:30:22 -0500
Subject: [PATCH] Use aio version of credentials only with blob storage (#106)

* Use aio version of credentials only with blob storage

* Add aiohttp dependency

* Fixes test

* Updates changelog
---
 CHANGELOG.md                     | 8 ++++++++
 prefect_azure/credentials.py     | 9 +++++----
 requirements.txt                 | 2 ++
 tests/test_aci_infrastructure.py | 2 +-
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cca2907..06c81a0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Security
 
+## 0.2.10
+
+Released on June 22nd, 2023.
+
+### Fixed
+
+- Fixed credentials imports to ensure async credential classes are only used with async clients - [#106](https://github.com/PrefectHQ/prefect-azure/pull/106)
+
 ## 0.2.9
 
 Released on June 21st, 2023.
diff --git a/prefect_azure/credentials.py b/prefect_azure/credentials.py
index 9e410e7..c8046c6 100644
--- a/prefect_azure/credentials.py
+++ b/prefect_azure/credentials.py
@@ -3,7 +3,8 @@
 import functools
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
-from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
+from azure.identity import ClientSecretCredential, DefaultAzureCredential
+from azure.identity.aio import DefaultAzureCredential as ADefaultAzureCredential
 from azure.mgmt.containerinstance import ContainerInstanceManagementClient
 from azure.mgmt.resource import ResourceManagementClient
 from pydantic import Field, SecretStr, root_validator
@@ -158,7 +159,7 @@ async def example_get_client_flow():
         if self.connection_string is None:
             return BlobServiceClient(
                 account_url=self.account_url,
-                credential=DefaultAzureCredential(),
+                credential=ADefaultAzureCredential(),
             )
 
         return BlobServiceClient.from_connection_string(
@@ -202,7 +203,7 @@ async def example_get_blob_client_flow():
             return BlobClient(
                 account_url=self.account_url,
                 container_name=container,
-                credential=DefaultAzureCredential(),
+                credential=ADefaultAzureCredential(),
                 blob_name=blob,
             )
 
@@ -247,7 +248,7 @@ async def example_get_container_client_flow():
             return ContainerClient(
                 account_url=self.account_url,
                 container_name=container,
-                credential=DefaultAzureCredential(),
+                credential=ADefaultAzureCredential(),
             )
 
         container_client = ContainerClient.from_connection_string(
diff --git a/requirements.txt b/requirements.txt
index 6e7c907..a0f0240 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,5 @@ prefect>=2.10.5
 azure_mgmt_containerinstance>=10.0
 azure_identity>=1.10
 azure-mgmt-resource>=21.2
+aiohttp
+
diff --git a/tests/test_aci_infrastructure.py b/tests/test_aci_infrastructure.py
index 5ba25bc..74ee6b4 100644
--- a/tests/test_aci_infrastructure.py
+++ b/tests/test_aci_infrastructure.py
@@ -7,7 +7,7 @@
 import pytest
 from anyio.abc import TaskStatus
 from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
-from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
+from azure.identity import ClientSecretCredential, DefaultAzureCredential
 from azure.mgmt.containerinstance.models import (
     EnvironmentVariable,
     ImageRegistryCredential,