update after MI changes

Azure · Oct 8, 2024 · 9cda34c · 9cda34c
1 parent b51b89b
commit 9cda34c
Show file tree

Hide file tree

Showing 6 changed files with 70 additions and 32 deletions.
diff --git a/azurefunctions-extensions-bindings-blob/azurefunctions/extensions/bindings/blob/__init__.py b/azurefunctions-extensions-bindings-blob/azurefunctions/extensions/bindings/blob/__init__.py
@@ -5,12 +5,15 @@
 from .blobClientConverter import BlobClientConverter
 from .containerClient import ContainerClient
 from .storageStreamDownloader import StorageStreamDownloader
+from .utils import get_connection_string, using_managed_identity
 
 __all__ = [
     "BlobClient",
     "ContainerClient",
     "StorageStreamDownloader",
     "BlobClientConverter",
+    "get_connection_string",
+    "using_managed_identity",
 ]
 
 __version__ = "1.0.0b1"
diff --git a/...unctions-extensions-bindings-blob/azurefunctions/extensions/bindings/blob/aio/__init__.py b/...unctions-extensions-bindings-blob/azurefunctions/extensions/bindings/blob/aio/__init__.py
@@ -5,6 +5,10 @@
 from .containerClient import ContainerClient
 from .storageStreamDownloader import StorageStreamDownloader
 
-__all__ = ["BlobClient", "ContainerClient", "StorageStreamDownloader"]
+__all__ = [
+    "BlobClient",
+    "ContainerClient",
+    "StorageStreamDownloader",
+]
 
-__version__ = "1.0.0b2"
+__version__ = "1.0.0b1"
diff --git a/...ctions-extensions-bindings-blob/azurefunctions/extensions/bindings/blob/aio/blobClient.py b/...ctions-extensions-bindings-blob/azurefunctions/extensions/bindings/blob/aio/blobClient.py
@@ -2,17 +2,20 @@
 #  Licensed under the MIT License.
 
 import json
-import os
 from typing import Union
 
-from azure.storage.blob.aio import BlobClient as AioBlobClientSdk
+from azure.storage.blob.aio import BlobServiceClient
+from azurefunctions.extensions.bindings.blob import (
+    get_connection_string,
+    using_managed_identity,
+)
 from azurefunctions.extensions.base import Datum, SdkType
 
 
 class BlobClient(SdkType):
     def __init__(self, *, data: Union[bytes, Datum]) -> None:
-        # model_binding_data properties
         self._data = data
+        self._using_managed_identity = False
         self._version = None
         self._source = None
         self._content_type = None
@@ -24,17 +27,24 @@ def __init__(self, *, data: Union[bytes, Datum]) -> None:
             self._source = data.source
             self._content_type = data.content_type
             content_json = json.loads(data.content)
-            self._connection = os.getenv(content_json["Connection"])
-            self._containerName = content_json["ContainerName"]
-            self._blobName = content_json["BlobName"]
+            self._connection = get_connection_string(content_json.get("Connection"))
+            self._using_managed_identity = using_managed_identity(
+                content_json.get("Connection")
+            )
+            self._containerName = content_json.get("ContainerName")
+            self._blobName = content_json.get("BlobName")
 
     # Returns a BlobClient
     async def get_sdk_type(self):
         if self._data:
-            return AioBlobClientSdk.from_connection_string(
-                conn_str=self._connection,
-                container_name=self._containerName,
-                blob_name=self._blobName,
+            blob_service_client = (
+                BlobServiceClient(account_url=self._connection)
+                if self._using_managed_identity
+                else BlobServiceClient.from_connection_string(self._connection)
+            )
+            return blob_service_client.get_blob_client(
+                container=self._containerName,
+                blob=self._blobName,
             )
         else:
             return None
diff --git a/...s-extensions-bindings-blob/azurefunctions/extensions/bindings/blob/aio/containerClient.py b/...s-extensions-bindings-blob/azurefunctions/extensions/bindings/blob/aio/containerClient.py
@@ -2,17 +2,21 @@
 #  Licensed under the MIT License.
 
 import json
-import os
 from typing import Union
 
-from azure.storage.blob.aio import ContainerClient as AioContainerClientSdk
+from azure.storage.blob.aio import BlobServiceClient
+from azurefunctions.extensions.bindings.blob import (
+    get_connection_string,
+    using_managed_identity,
+)
 from azurefunctions.extensions.base import Datum, SdkType
 
 
 class ContainerClient(SdkType):
     def __init__(self, *, data: Union[bytes, Datum]) -> None:
         # model_binding_data properties
         self._data = data
+        self._using_managed_identity = False
         self._version = ""
         self._source = ""
         self._content_type = ""
@@ -24,15 +28,23 @@ def __init__(self, *, data: Union[bytes, Datum]) -> None:
             self._source = data.source
             self._content_type = data.content_type
             content_json = json.loads(data.content)
-            self._connection = os.getenv(content_json["Connection"])
-            self._containerName = content_json["ContainerName"]
-            self._blobName = content_json["BlobName"]
+            self._connection = get_connection_string(content_json.get("Connection"))
+            self._using_managed_identity = using_managed_identity(
+                content_json.get("Connection")
+            )
+            self._containerName = content_json.get("ContainerName")
+            self._blobName = content_json.get("BlobName")
 
     # Returns a ContainerClient
     async def get_sdk_type(self):
         if self._data:
-            return AioContainerClientSdk.from_connection_string(
-                conn_str=self._connection, container_name=self._containerName
+            blob_service_client = (
+                BlobServiceClient(account_url=self._connection)
+                if self._using_managed_identity
+                else BlobServiceClient.from_connection_string(self._connection)
+            )
+            return blob_service_client.get_container_client(
+                container=self._containerName
             )
         else:
             return None
diff --git a/...ions-bindings-blob/azurefunctions/extensions/bindings/blob/aio/storageStreamDownloader.py b/...ions-bindings-blob/azurefunctions/extensions/bindings/blob/aio/storageStreamDownloader.py
@@ -2,17 +2,21 @@
 #  Licensed under the MIT License.
 
 import json
-import os
 from typing import Union
 
-from azure.storage.blob.aio import BlobClient as AioBlobClientSdk
+from azure.storage.blob.aio import BlobServiceClient
+from azurefunctions.extensions.bindings.blob import (
+    get_connection_string,
+    using_managed_identity,
+)
 from azurefunctions.extensions.base import Datum, SdkType
 
 
 class StorageStreamDownloader(SdkType):
     def __init__(self, *, data: Union[bytes, Datum]) -> None:
         # model_binding_data properties
-        self._data = data or {}
+        self._data = data
+        self._using_managed_identity = False
         self._version = ""
         self._source = ""
         self._content_type = ""
@@ -24,20 +28,25 @@ def __init__(self, *, data: Union[bytes, Datum]) -> None:
             self._source = data.source
             self._content_type = data.content_type
             content_json = json.loads(data.content)
-            self._connection = os.getenv(content_json["Connection"])
-            self._containerName = content_json["ContainerName"]
-            self._blobName = content_json["BlobName"]
+            self._connection = get_connection_string(content_json.get("Connection"))
+            self._using_managed_identity = using_managed_identity(
+                content_json.get("Connection")
+            )
+            self._containerName = content_json.get("ContainerName")
+            self._blobName = content_json.get("BlobName")
 
     # Returns a StorageStreamDownloader
     async def get_sdk_type(self):
         if self._data:
-            # Create BlobClient
-            blob_client = AioBlobClientSdk.from_connection_string(
-                conn_str=self._connection,
-                container_name=self._containerName,
-                blob_name=self._blobName,
+            blob_service_client = (
+                BlobServiceClient(account_url=self._connection)
+                if self._using_managed_identity
+                else BlobServiceClient.from_connection_string(self._connection)
             )
             # download_blob() returns a StorageStreamDownloader object
-            return await blob_client.download_blob()
+            return blob_service_client.get_blob_client(
+                container=self._containerName,
+                blob=self._blobName,
+            ).download_blob()
         else:
             return None
diff --git a/azurefunctions-extensions-bindings-blob/pyproject.toml b/azurefunctions-extensions-bindings-blob/pyproject.toml
@@ -37,7 +37,7 @@ dev = [
     'pre-commit'
     ]
 aio = [
-    'azure-storage-blob[aio]==12.20.0'
+    'azure-storage-blob[aio]==12.23.1'
     ]
 
 [tool.setuptools.dynamic]