Add azure_storage_container to Azure filesystem (PrefectHQ#11784)

Co-authored-by: Thomas Frederik Hoeck <[email protected]> Co-authored-by: nate nowack <[email protected]> Co-authored-by: nate nowack <[email protected]>
Point72 · Feb 1, 2024 · c03fe84 · c03fe84
1 parent 090d44d
commit c03fe84
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 3 deletions.
diff --git a/src/prefect/filesystems.py b/src/prefect/filesystems.py
@@ -683,12 +683,27 @@ class Azure(WritableFileSystem, WritableDeploymentStorage):
             " require ADLFS to use DefaultAzureCredentials."
         ),
     )
-
+    azure_storage_container: Optional[SecretStr] = Field(
+        default=None,
+        title="Azure storage container",
+        description=(
+            "Blob Container in Azure Storage Account. If set the 'bucket_path' will"
+            " be interpreted using the following URL format:"
+            "'az://<container>@<storage_account>.dfs.core.windows.net/<bucket_path>'."
+        ),
+    )
     _remote_file_system: RemoteFileSystem = None
 
     @property
     def basepath(self) -> str:
-        return f"az://{self.bucket_path}"
+        if self.azure_storage_container:
+            return (
+                f"az://{self.azure_storage_container.get_secret_value()}"
+                f"@{self.azure_storage_account_name.get_secret_value()}"
+                f".dfs.core.windows.net/{self.bucket_path}"
+            )
+        else:
+            return f"az://{self.bucket_path}"
 
     @property
     def filesystem(self) -> RemoteFileSystem:
@@ -713,7 +728,7 @@ def filesystem(self) -> RemoteFileSystem:
             )
         settings["anon"] = self.azure_storage_anon
         self._remote_file_system = RemoteFileSystem(
-            basepath=f"az://{self.bucket_path}", settings=settings
+            basepath=self.basepath, settings=settings
         )
         return self._remote_file_system
 

diff --git a/tests/test_filesystems.py b/tests/test_filesystems.py
@@ -709,3 +709,28 @@ def test_init_with_anon(self, monkeypatch):
                 "anon": False,
             },
         )
+
+    def test_init_with_container(self, monkeypatch):
+        remote_storage_mock = MagicMock()
+        monkeypatch.setattr("prefect.filesystems.RemoteFileSystem", remote_storage_mock)
+        Azure(
+            azure_storage_tenant_id="tenant",
+            azure_storage_account_name="account",
+            azure_storage_client_id="client_id",
+            azure_storage_account_key="key",
+            azure_storage_client_secret="secret",
+            azure_storage_container="container",
+            bucket_path="bucket",
+            azure_storage_anon=False,
+        ).filesystem
+        remote_storage_mock.assert_called_once_with(
+            basepath="az://[email protected]/bucket",
+            settings={
+                "account_name": "account",
+                "account_key": "key",
+                "tenant_id": "tenant",
+                "client_id": "client_id",
+                "client_secret": "secret",
+                "anon": False,
+            },
+        )