We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[FR] AzureBlobArtifactRepository connection using Service Principal credentials
## Willingness to contribute
The MLflow Community encourages new feature contributions. Would you or another member of your organization be willing to contribute an implementation of this feature (either as an MLflow Plugin or an enhancement to the MLflow code base)?

- [x] Yes. I can contribute this feature independently.
- [x] Yes. I would be willing to contribute this feature with guidance from the MLflow community.
- [ ] No. I cannot contribute this feature at this time.

## Proposal Summary

Allow the user to connect with the AzureBlobArtifactRepository using service principal credentials taken from the following environment variables:
* AZURE_TENANT_ID
* AZURE_CLIENT_ID
* AZURE_CLIENT_SECRET

## Motivation
- What is the use case for this feature?
Having the AZURE_STORAGE_CONNECTION_STRING or AZURE_STORAGE_ACCESS_KEY variables set will give the user of that environment complete control over the whole storage account. This is not always desirable in a situation where a storage account has multiple blob containers and the users of one blob container should not have access to another user's container. 

- Why is this use case valuable to support for MLflow users in general?
It gives the user more options to choose from while connecting with their artifact repository.

- Why is this use case valuable to support for your project(s) or organization
My organizations hosts multiple tracking servers (with corresponding artifact repositories) in which groups of users should be separated. Allowing the connection via service principal credentials gives us the ability assign permissions on a container to a service principal

- Why is it currently difficult to achieve this use case? (please be as specific as possible about why related MLflow features and components are insufficient)
The AzureBlobArtifactRepository requires either the AZURE_STORAGE_CONNECTION_STRING or AZURE_STORAGE_ACCESS_KEY to be set, which gives the user full control of the whole storage account.

### What component(s), interfaces, languages, and integrations does this feature affect?
Components 
- [x] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [x] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: Local serving, model deployment tools, spark UDFs
- [ ] `area/server-infra`: MLflow server, JavaScript dev server
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interfaces
- [ ] `area/uiux`: Front-end, user experience, JavaScript, plotting
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Languages 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

## Details

The solution would entail adding another elif branch in the AzureBlobArtifactRepository __init__ function, that checks for the AZURE_TENANT_ID, AZURE_CLIENT_ID and AZURE_CLIENT_SECRET environment variables, creates a [ClientSecretCredential](https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity.clientsecretcredential?view=azure-python) and uses it to create the BlobServiceClient.



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mlflow/store/artifact/azure_blob_artifact_repo.py`
Content:
```
1 import os
2 import posixpath
3 import re
4 import urllib.parse
5 
6 from mlflow.entities import FileInfo
7 from mlflow.exceptions import MlflowException
8 from mlflow.store.artifact.artifact_repo import ArtifactRepository
9 
10 
11 class AzureBlobArtifactRepository(ArtifactRepository):
12     """
13     Stores artifacts on Azure Blob Storage.
14 
15     This repository is used with URIs of the form
16     ``wasbs://<container-name>@<ystorage-account-name>.blob.core.windows.net/<path>``,
17     following the same URI scheme as Hadoop on Azure blob storage. It requires that your Azure
18     storage access key be available in the environment variable ``AZURE_STORAGE_ACCESS_KEY``.
19     """
20 
21     def __init__(self, artifact_uri, client=None):
22         super().__init__(artifact_uri)
23 
24         # Allow override for testing
25         if client:
26             self.client = client
27             return
28 
29         from azure.storage.blob import BlobServiceClient
30 
31         (_, account, _) = AzureBlobArtifactRepository.parse_wasbs_uri(artifact_uri)
32         if "AZURE_STORAGE_CONNECTION_STRING" in os.environ:
33             self.client = BlobServiceClient.from_connection_string(
34                 conn_str=os.environ.get("AZURE_STORAGE_CONNECTION_STRING")
35             )
36         elif "AZURE_STORAGE_ACCESS_KEY" in os.environ:
37             account_url = "https://{account}.blob.core.windows.net".format(account=account)
38             self.client = BlobServiceClient(
39                 account_url=account_url, credential=os.environ.get("AZURE_STORAGE_ACCESS_KEY")
40             )
41         else:
42             raise Exception(
43                 "You need to set one of AZURE_STORAGE_CONNECTION_STRING or "
44                 "AZURE_STORAGE_ACCESS_KEY to access Azure storage."
45             )
46 
47     @staticmethod
48     def parse_wasbs_uri(uri):
49         """Parse a wasbs:// URI, returning (container, storage_account, path)."""
50         parsed = urllib.parse.urlparse(uri)
51         if parsed.scheme != "wasbs":
52             raise Exception("Not a WASBS URI: %s" % uri)
53         match = re.match("([^@]+)@([^.]+)\\.blob\\.core\\.windows\\.net", parsed.netloc)
54         if match is None:
55             raise Exception(
56                 "WASBS URI must be of the form " "<container>@<account>.blob.core.windows.net"
57             )
58         container = match.group(1)
59         storage_account = match.group(2)
60         path = parsed.path
61         if path.startswith("/"):
62             path = path[1:]
63         return container, storage_account, path
64 
65     def log_artifact(self, local_file, artifact_path=None):
66         (container, _, dest_path) = self.parse_wasbs_uri(self.artifact_uri)
67         container_client = self.client.get_container_client(container)
68         if artifact_path:
69             dest_path = posixpath.join(dest_path, artifact_path)
70         dest_path = posixpath.join(dest_path, os.path.basename(local_file))
71         with open(local_file, "rb") as file:
72             container_client.upload_blob(dest_path, file)
73 
74     def log_artifacts(self, local_dir, artifact_path=None):
75         (container, _, dest_path) = self.parse_wasbs_uri(self.artifact_uri)
76         container_client = self.client.get_container_client(container)
77         if artifact_path:
78             dest_path = posixpath.join(dest_path, artifact_path)
79         local_dir = os.path.abspath(local_dir)
80         for (root, _, filenames) in os.walk(local_dir):
81             upload_path = dest_path
82             if root != local_dir:
83                 rel_path = os.path.relpath(root, local_dir)
84                 upload_path = posixpath.join(dest_path, rel_path)
85             for f in filenames:
86                 remote_file_path = posixpath.join(upload_path, f)
87                 local_file_path = os.path.join(root, f)
88                 with open(local_file_path, "rb") as file:
89                     container_client.upload_blob(remote_file_path, file)
90 
91     def list_artifacts(self, path=None):
92         # Newer versions of `azure-storage-blob` (>= 12.4.0) provide a public
93         # `azure.storage.blob.BlobPrefix` object to signify that a blob is a directory,
94         # while older versions only expose this API internally as
95         # `azure.storage.blob._models.BlobPrefix`
96         try:
97             from azure.storage.blob import BlobPrefix
98         except ImportError:
99             from azure.storage.blob._models import BlobPrefix
100 
101         (container, _, artifact_path) = self.parse_wasbs_uri(self.artifact_uri)
102         container_client = self.client.get_container_client(container)
103         dest_path = artifact_path
104         if path:
105             dest_path = posixpath.join(dest_path, path)
106         infos = []
107         prefix = dest_path if dest_path.endswith("/") else dest_path + "/"
108         results = container_client.walk_blobs(name_starts_with=prefix)
109         for r in results:
110             if not r.name.startswith(artifact_path):
111                 raise MlflowException(
112                     "The name of the listed Azure blob does not begin with the specified"
113                     " artifact path. Artifact path: {artifact_path}. Blob name:"
114                     " {blob_name}".format(artifact_path=artifact_path, blob_name=r.name)
115                 )
116             if isinstance(r, BlobPrefix):  # This is a prefix for items in a subdirectory
117                 subdir = posixpath.relpath(path=r.name, start=artifact_path)
118                 if subdir.endswith("/"):
119                     subdir = subdir[:-1]
120                 infos.append(FileInfo(subdir, True, None))
121             else:  # Just a plain old blob
122                 file_name = posixpath.relpath(path=r.name, start=artifact_path)
123                 infos.append(FileInfo(file_name, False, r.size))
124         # The list_artifacts API expects us to return an empty list if the
125         # the path references a single file.
126         rel_path = dest_path[len(artifact_path) + 1 :]
127         if (len(infos) == 1) and not infos[0].is_dir and (infos[0].path == rel_path):
128             return []
129         return sorted(infos, key=lambda f: f.path)
130 
131     def _download_file(self, remote_file_path, local_path):
132         (container, _, remote_root_path) = self.parse_wasbs_uri(self.artifact_uri)
133         container_client = self.client.get_container_client(container)
134         remote_full_path = posixpath.join(remote_root_path, remote_file_path)
135         with open(local_path, "wb") as file:
136             container_client.download_blob(remote_full_path).readinto(file)
137 
138     def delete_artifacts(self, artifact_path=None):
139         raise MlflowException("Not implemented yet")
140
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mlflow/store/artifact/azure_blob_artifact_repo.py b/mlflow/store/artifact/azure_blob_artifact_repo.py
--- a/mlflow/store/artifact/azure_blob_artifact_repo.py
+++ b/mlflow/store/artifact/azure_blob_artifact_repo.py
@@ -14,8 +14,10 @@
 
     This repository is used with URIs of the form
     ``wasbs://<container-name>@<ystorage-account-name>.blob.core.windows.net/<path>``,
-    following the same URI scheme as Hadoop on Azure blob storage. It requires that your Azure
-    storage access key be available in the environment variable ``AZURE_STORAGE_ACCESS_KEY``.
+    following the same URI scheme as Hadoop on Azure blob storage. It requires either that:
+    - Azure storage connection string is in the env var ``AZURE_STORAGE_CONNECTION_STRING``
+    - Azure storage access key is in the env var ``AZURE_STORAGE_ACCESS_KEY``
+    - DefaultAzureCredential is configured
     """
 
     def __init__(self, artifact_uri, client=None):
@@ -39,9 +41,17 @@
                 account_url=account_url, credential=os.environ.get("AZURE_STORAGE_ACCESS_KEY")
             )
         else:
-            raise Exception(
-                "You need to set one of AZURE_STORAGE_CONNECTION_STRING or "
-                "AZURE_STORAGE_ACCESS_KEY to access Azure storage."
+            try:
+                from azure.identity import DefaultAzureCredential
+            except ImportError as exc:
+                raise ImportError(
+                    "Using DefaultAzureCredential requires the azure-identity package. "
+                    "Please install it via: pip install azure-identity"
+                ) from exc
+
+            account_url = "https://{account}.blob.core.windows.net".format(account=account)
+            self.client = BlobServiceClient(
+                account_url=account_url, credential=DefaultAzureCredential()
             )
 
     @staticmethod

{"golden_diff": "diff --git a/mlflow/store/artifact/azure_blob_artifact_repo.py b/mlflow/store/artifact/azure_blob_artifact_repo.py\n--- a/mlflow/store/artifact/azure_blob_artifact_repo.py\n+++ b/mlflow/store/artifact/azure_blob_artifact_repo.py\n@@ -14,8 +14,10 @@\n \n     This repository is used with URIs of the form\n     ``wasbs://<container-name>@<ystorage-account-name>.blob.core.windows.net/<path>``,\n-    following the same URI scheme as Hadoop on Azure blob storage. It requires that your Azure\n-    storage access key be available in the environment variable ``AZURE_STORAGE_ACCESS_KEY``.\n+    following the same URI scheme as Hadoop on Azure blob storage. It requires either that:\n+    - Azure storage connection string is in the env var ``AZURE_STORAGE_CONNECTION_STRING``\n+    - Azure storage access key is in the env var ``AZURE_STORAGE_ACCESS_KEY``\n+    - DefaultAzureCredential is configured\n     \"\"\"\n \n     def __init__(self, artifact_uri, client=None):\n@@ -39,9 +41,17 @@\n                 account_url=account_url, credential=os.environ.get(\"AZURE_STORAGE_ACCESS_KEY\")\n             )\n         else:\n-            raise Exception(\n-                \"You need to set one of AZURE_STORAGE_CONNECTION_STRING or \"\n-                \"AZURE_STORAGE_ACCESS_KEY to access Azure storage.\"\n+            try:\n+                from azure.identity import DefaultAzureCredential\n+            except ImportError as exc:\n+                raise ImportError(\n+                    \"Using DefaultAzureCredential requires the azure-identity package. \"\n+                    \"Please install it via: pip install azure-identity\"\n+                ) from exc\n+\n+            account_url = \"https://{account}.blob.core.windows.net\".format(account=account)\n+            self.client = BlobServiceClient(\n+                account_url=account_url, credential=DefaultAzureCredential()\n             )\n \n     @staticmethod\n", "issue": "[FR] AzureBlobArtifactRepository connection using Service Principal credentials\n## Willingness to contribute\r\nThe MLflow Community encourages new feature contributions. Would you or another member of your organization be willing to contribute an implementation of this feature (either as an MLflow Plugin or an enhancement to the MLflow code base)?\r\n\r\n- [x] Yes. I can contribute this feature independently.\r\n- [x] Yes. I would be willing to contribute this feature with guidance from the MLflow community.\r\n- [ ] No. I cannot contribute this feature at this time.\r\n\r\n## Proposal Summary\r\n\r\nAllow the user to connect with the AzureBlobArtifactRepository using service principal credentials taken from the following environment variables:\r\n* AZURE_TENANT_ID\r\n* AZURE_CLIENT_ID\r\n* AZURE_CLIENT_SECRET\r\n\r\n## Motivation\r\n- What is the use case for this feature?\r\nHaving the AZURE_STORAGE_CONNECTION_STRING or AZURE_STORAGE_ACCESS_KEY variables set will give the user of that environment complete control over the whole storage account. This is not always desirable in a situation where a storage account has multiple blob containers and the users of one blob container should not have access to another user's container. \r\n\r\n- Why is this use case valuable to support for MLflow users in general?\r\nIt gives the user more options to choose from while connecting with their artifact repository.\r\n\r\n- Why is this use case valuable to support for your project(s) or organization\r\nMy organizations hosts multiple tracking servers (with corresponding artifact repositories) in which groups of users should be separated. Allowing the connection via service principal credentials gives us the ability assign permissions on a container to a service principal\r\n\r\n- Why is it currently difficult to achieve this use case? (please be as specific as possible about why related MLflow features and components are insufficient)\r\nThe AzureBlobArtifactRepository requires either the AZURE_STORAGE_CONNECTION_STRING or AZURE_STORAGE_ACCESS_KEY to be set, which gives the user full control of the whole storage account.\r\n\r\n### What component(s), interfaces, languages, and integrations does this feature affect?\r\nComponents \r\n- [x] `area/artifacts`: Artifact stores and artifact logging\r\n- [ ] `area/build`: Build and test infrastructure for MLflow\r\n- [x] `area/docs`: MLflow documentation pages\r\n- [ ] `area/examples`: Example code\r\n- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry\r\n- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors\r\n- [ ] `area/projects`: MLproject format, project running backends\r\n- [ ] `area/scoring`: Local serving, model deployment tools, spark UDFs\r\n- [ ] `area/server-infra`: MLflow server, JavaScript dev server\r\n- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging\r\n\r\nInterfaces\r\n- [ ] `area/uiux`: Front-end, user experience, JavaScript, plotting\r\n- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models\r\n- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry\r\n- [ ] `area/windows`: Windows support\r\n\r\nLanguages \r\n- [ ] `language/r`: R APIs and clients\r\n- [ ] `language/java`: Java APIs and clients\r\n- [ ] `language/new`: Proposals for new client languages\r\n\r\nIntegrations\r\n- [ ] `integrations/azure`: Azure and Azure ML integrations\r\n- [ ] `integrations/sagemaker`: SageMaker integrations\r\n- [ ] `integrations/databricks`: Databricks integrations\r\n\r\n## Details\r\n\r\nThe solution would entail adding another elif branch in the AzureBlobArtifactRepository __init__ function, that checks for the AZURE_TENANT_ID, AZURE_CLIENT_ID and AZURE_CLIENT_SECRET environment variables, creates a [ClientSecretCredential](https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity.clientsecretcredential?view=azure-python) and uses it to create the BlobServiceClient.\r\n\r\n\n", "before_files": [{"content": "import os\nimport posixpath\nimport re\nimport urllib.parse\n\nfrom mlflow.entities import FileInfo\nfrom mlflow.exceptions import MlflowException\nfrom mlflow.store.artifact.artifact_repo import ArtifactRepository\n\n\nclass AzureBlobArtifactRepository(ArtifactRepository):\n    \"\"\"\n    Stores artifacts on Azure Blob Storage.\n\n    This repository is used with URIs of the form\n    ``wasbs://<container-name>@<ystorage-account-name>.blob.core.windows.net/<path>``,\n    following the same URI scheme as Hadoop on Azure blob storage. It requires that your Azure\n    storage access key be available in the environment variable ``AZURE_STORAGE_ACCESS_KEY``.\n    \"\"\"\n\n    def __init__(self, artifact_uri, client=None):\n        super().__init__(artifact_uri)\n\n        # Allow override for testing\n        if client:\n            self.client = client\n            return\n\n        from azure.storage.blob import BlobServiceClient\n\n        (_, account, _) = AzureBlobArtifactRepository.parse_wasbs_uri(artifact_uri)\n        if \"AZURE_STORAGE_CONNECTION_STRING\" in os.environ:\n            self.client = BlobServiceClient.from_connection_string(\n                conn_str=os.environ.get(\"AZURE_STORAGE_CONNECTION_STRING\")\n            )\n        elif \"AZURE_STORAGE_ACCESS_KEY\" in os.environ:\n            account_url = \"https://{account}.blob.core.windows.net\".format(account=account)\n            self.client = BlobServiceClient(\n                account_url=account_url, credential=os.environ.get(\"AZURE_STORAGE_ACCESS_KEY\")\n            )\n        else:\n            raise Exception(\n                \"You need to set one of AZURE_STORAGE_CONNECTION_STRING or \"\n                \"AZURE_STORAGE_ACCESS_KEY to access Azure storage.\"\n            )\n\n    @staticmethod\n    def parse_wasbs_uri(uri):\n        \"\"\"Parse a wasbs:// URI, returning (container, storage_account, path).\"\"\"\n        parsed = urllib.parse.urlparse(uri)\n        if parsed.scheme != \"wasbs\":\n            raise Exception(\"Not a WASBS URI: %s\" % uri)\n        match = re.match(\"([^@]+)@([^.]+)\\\\.blob\\\\.core\\\\.windows\\\\.net\", parsed.netloc)\n        if match is None:\n            raise Exception(\n                \"WASBS URI must be of the form \" \"<container>@<account>.blob.core.windows.net\"\n            )\n        container = match.group(1)\n        storage_account = match.group(2)\n        path = parsed.path\n        if path.startswith(\"/\"):\n            path = path[1:]\n        return container, storage_account, path\n\n    def log_artifact(self, local_file, artifact_path=None):\n        (container, _, dest_path) = self.parse_wasbs_uri(self.artifact_uri)\n        container_client = self.client.get_container_client(container)\n        if artifact_path:\n            dest_path = posixpath.join(dest_path, artifact_path)\n        dest_path = posixpath.join(dest_path, os.path.basename(local_file))\n        with open(local_file, \"rb\") as file:\n            container_client.upload_blob(dest_path, file)\n\n    def log_artifacts(self, local_dir, artifact_path=None):\n        (container, _, dest_path) = self.parse_wasbs_uri(self.artifact_uri)\n        container_client = self.client.get_container_client(container)\n        if artifact_path:\n            dest_path = posixpath.join(dest_path, artifact_path)\n        local_dir = os.path.abspath(local_dir)\n        for (root, _, filenames) in os.walk(local_dir):\n            upload_path = dest_path\n            if root != local_dir:\n                rel_path = os.path.relpath(root, local_dir)\n                upload_path = posixpath.join(dest_path, rel_path)\n            for f in filenames:\n                remote_file_path = posixpath.join(upload_path, f)\n                local_file_path = os.path.join(root, f)\n                with open(local_file_path, \"rb\") as file:\n                    container_client.upload_blob(remote_file_path, file)\n\n    def list_artifacts(self, path=None):\n        # Newer versions of `azure-storage-blob` (>= 12.4.0) provide a public\n        # `azure.storage.blob.BlobPrefix` object to signify that a blob is a directory,\n        # while older versions only expose this API internally as\n        # `azure.storage.blob._models.BlobPrefix`\n        try:\n            from azure.storage.blob import BlobPrefix\n        except ImportError:\n            from azure.storage.blob._models import BlobPrefix\n\n        (container, _, artifact_path) = self.parse_wasbs_uri(self.artifact_uri)\n        container_client = self.client.get_container_client(container)\n        dest_path = artifact_path\n        if path:\n            dest_path = posixpath.join(dest_path, path)\n        infos = []\n        prefix = dest_path if dest_path.endswith(\"/\") else dest_path + \"/\"\n        results = container_client.walk_blobs(name_starts_with=prefix)\n        for r in results:\n            if not r.name.startswith(artifact_path):\n                raise MlflowException(\n                    \"The name of the listed Azure blob does not begin with the specified\"\n                    \" artifact path. Artifact path: {artifact_path}. Blob name:\"\n                    \" {blob_name}\".format(artifact_path=artifact_path, blob_name=r.name)\n                )\n            if isinstance(r, BlobPrefix):  # This is a prefix for items in a subdirectory\n                subdir = posixpath.relpath(path=r.name, start=artifact_path)\n                if subdir.endswith(\"/\"):\n                    subdir = subdir[:-1]\n                infos.append(FileInfo(subdir, True, None))\n            else:  # Just a plain old blob\n                file_name = posixpath.relpath(path=r.name, start=artifact_path)\n                infos.append(FileInfo(file_name, False, r.size))\n        # The list_artifacts API expects us to return an empty list if the\n        # the path references a single file.\n        rel_path = dest_path[len(artifact_path) + 1 :]\n        if (len(infos) == 1) and not infos[0].is_dir and (infos[0].path == rel_path):\n            return []\n        return sorted(infos, key=lambda f: f.path)\n\n    def _download_file(self, remote_file_path, local_path):\n        (container, _, remote_root_path) = self.parse_wasbs_uri(self.artifact_uri)\n        container_client = self.client.get_container_client(container)\n        remote_full_path = posixpath.join(remote_root_path, remote_file_path)\n        with open(local_path, \"wb\") as file:\n            container_client.download_blob(remote_full_path).readinto(file)\n\n    def delete_artifacts(self, artifact_path=None):\n        raise MlflowException(\"Not implemented yet\")\n", "path": "mlflow/store/artifact/azure_blob_artifact_repo.py"}], "after_files": [{"content": "import os\nimport posixpath\nimport re\nimport urllib.parse\n\nfrom mlflow.entities import FileInfo\nfrom mlflow.exceptions import MlflowException\nfrom mlflow.store.artifact.artifact_repo import ArtifactRepository\n\n\nclass AzureBlobArtifactRepository(ArtifactRepository):\n    \"\"\"\n    Stores artifacts on Azure Blob Storage.\n\n    This repository is used with URIs of the form\n    ``wasbs://<container-name>@<ystorage-account-name>.blob.core.windows.net/<path>``,\n    following the same URI scheme as Hadoop on Azure blob storage. It requires either that:\n    - Azure storage connection string is in the env var ``AZURE_STORAGE_CONNECTION_STRING``\n    - Azure storage access key is in the env var ``AZURE_STORAGE_ACCESS_KEY``\n    - DefaultAzureCredential is configured\n    \"\"\"\n\n    def __init__(self, artifact_uri, client=None):\n        super().__init__(artifact_uri)\n\n        # Allow override for testing\n        if client:\n            self.client = client\n            return\n\n        from azure.storage.blob import BlobServiceClient\n\n        (_, account, _) = AzureBlobArtifactRepository.parse_wasbs_uri(artifact_uri)\n        if \"AZURE_STORAGE_CONNECTION_STRING\" in os.environ:\n            self.client = BlobServiceClient.from_connection_string(\n                conn_str=os.environ.get(\"AZURE_STORAGE_CONNECTION_STRING\")\n            )\n        elif \"AZURE_STORAGE_ACCESS_KEY\" in os.environ:\n            account_url = \"https://{account}.blob.core.windows.net\".format(account=account)\n            self.client = BlobServiceClient(\n                account_url=account_url, credential=os.environ.get(\"AZURE_STORAGE_ACCESS_KEY\")\n            )\n        else:\n            try:\n                from azure.identity import DefaultAzureCredential\n            except ImportError as exc:\n                raise ImportError(\n                    \"Using DefaultAzureCredential requires the azure-identity package. \"\n                    \"Please install it via: pip install azure-identity\"\n                ) from exc\n\n            account_url = \"https://{account}.blob.core.windows.net\".format(account=account)\n            self.client = BlobServiceClient(\n                account_url=account_url, credential=DefaultAzureCredential()\n            )\n\n    @staticmethod\n    def parse_wasbs_uri(uri):\n        \"\"\"Parse a wasbs:// URI, returning (container, storage_account, path).\"\"\"\n        parsed = urllib.parse.urlparse(uri)\n        if parsed.scheme != \"wasbs\":\n            raise Exception(\"Not a WASBS URI: %s\" % uri)\n        match = re.match(\"([^@]+)@([^.]+)\\\\.blob\\\\.core\\\\.windows\\\\.net\", parsed.netloc)\n        if match is None:\n            raise Exception(\n                \"WASBS URI must be of the form \" \"<container>@<account>.blob.core.windows.net\"\n            )\n        container = match.group(1)\n        storage_account = match.group(2)\n        path = parsed.path\n        if path.startswith(\"/\"):\n            path = path[1:]\n        return container, storage_account, path\n\n    def log_artifact(self, local_file, artifact_path=None):\n        (container, _, dest_path) = self.parse_wasbs_uri(self.artifact_uri)\n        container_client = self.client.get_container_client(container)\n        if artifact_path:\n            dest_path = posixpath.join(dest_path, artifact_path)\n        dest_path = posixpath.join(dest_path, os.path.basename(local_file))\n        with open(local_file, \"rb\") as file:\n            container_client.upload_blob(dest_path, file)\n\n    def log_artifacts(self, local_dir, artifact_path=None):\n        (container, _, dest_path) = self.parse_wasbs_uri(self.artifact_uri)\n        container_client = self.client.get_container_client(container)\n        if artifact_path:\n            dest_path = posixpath.join(dest_path, artifact_path)\n        local_dir = os.path.abspath(local_dir)\n        for (root, _, filenames) in os.walk(local_dir):\n            upload_path = dest_path\n            if root != local_dir:\n                rel_path = os.path.relpath(root, local_dir)\n                upload_path = posixpath.join(dest_path, rel_path)\n            for f in filenames:\n                remote_file_path = posixpath.join(upload_path, f)\n                local_file_path = os.path.join(root, f)\n                with open(local_file_path, \"rb\") as file:\n                    container_client.upload_blob(remote_file_path, file)\n\n    def list_artifacts(self, path=None):\n        # Newer versions of `azure-storage-blob` (>= 12.4.0) provide a public\n        # `azure.storage.blob.BlobPrefix` object to signify that a blob is a directory,\n        # while older versions only expose this API internally as\n        # `azure.storage.blob._models.BlobPrefix`\n        try:\n            from azure.storage.blob import BlobPrefix\n        except ImportError:\n            from azure.storage.blob._models import BlobPrefix\n\n        (container, _, artifact_path) = self.parse_wasbs_uri(self.artifact_uri)\n        container_client = self.client.get_container_client(container)\n        dest_path = artifact_path\n        if path:\n            dest_path = posixpath.join(dest_path, path)\n        infos = []\n        prefix = dest_path if dest_path.endswith(\"/\") else dest_path + \"/\"\n        results = container_client.walk_blobs(name_starts_with=prefix)\n        for r in results:\n            if not r.name.startswith(artifact_path):\n                raise MlflowException(\n                    \"The name of the listed Azure blob does not begin with the specified\"\n                    \" artifact path. Artifact path: {artifact_path}. Blob name:\"\n                    \" {blob_name}\".format(artifact_path=artifact_path, blob_name=r.name)\n                )\n            if isinstance(r, BlobPrefix):  # This is a prefix for items in a subdirectory\n                subdir = posixpath.relpath(path=r.name, start=artifact_path)\n                if subdir.endswith(\"/\"):\n                    subdir = subdir[:-1]\n                infos.append(FileInfo(subdir, True, None))\n            else:  # Just a plain old blob\n                file_name = posixpath.relpath(path=r.name, start=artifact_path)\n                infos.append(FileInfo(file_name, False, r.size))\n        # The list_artifacts API expects us to return an empty list if the\n        # the path references a single file.\n        rel_path = dest_path[len(artifact_path) + 1 :]\n        if (len(infos) == 1) and not infos[0].is_dir and (infos[0].path == rel_path):\n            return []\n        return sorted(infos, key=lambda f: f.path)\n\n    def _download_file(self, remote_file_path, local_path):\n        (container, _, remote_root_path) = self.parse_wasbs_uri(self.artifact_uri)\n        container_client = self.client.get_container_client(container)\n        remote_full_path = posixpath.join(remote_root_path, remote_file_path)\n        with open(local_path, \"wb\") as file:\n            container_client.download_blob(remote_full_path).readinto(file)\n\n    def delete_artifacts(self, artifact_path=None):\n        raise MlflowException(\"Not implemented yet\")\n", "path": "mlflow/store/artifact/azure_blob_artifact_repo.py"}]}