fix(azure): update DatastoreType import

src/promptflow-azure/promptflow/azure/_storage/blob/client.py

@@ -6,14 +6,14 @@
 
 from azure.ai.ml import MLClient
 from azure.ai.ml._azure_environments import _get_storage_endpoint_from_metadata
-from azure.ai.ml._restclient.v2022_10_01.models import DatastoreType
 from azure.ai.ml.constants._common import LONG_URI_FORMAT, STORAGE_ACCOUNT_URLS
 from azure.ai.ml.entities._credentials import AccountKeyConfiguration
 from azure.ai.ml.entities._datastore.datastore import Datastore
 from azure.ai.ml.operations import DatastoreOperations
 from azure.storage.blob import ContainerClient
 
 from promptflow.exceptions import UserErrorException
+from promptflow.azure._utils._datastore_type import DatastoreType
 
 _datastore_cache = {}
 _thread_lock = threading.Lock()

src/promptflow-azure/promptflow/azure/_utils/_datastore_type.py

@@ -0,0 +1,10 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+try:
+    from azure.ai.ml._restclient.v2022_10_01_preview.models import DatastoreType
+except ImportError:
+    from azure.ai.ml._restclient.v2022_10_01.models import DatastoreType
+
+__all__ = ["DatastoreType"]

src/promptflow-azure/promptflow/azure/operations/_artifact_utilities.py

@@ -13,7 +13,6 @@
 from azure.ai.ml._artifacts._blob_storage_helper import BlobStorageClient
 from azure.ai.ml._artifacts._gen2_storage_helper import Gen2StorageClient
 from azure.ai.ml._azure_environments import _get_storage_endpoint_from_metadata
-from azure.ai.ml._restclient.v2022_10_01.models import DatastoreType
 from azure.ai.ml._scope_dependent_operations import OperationScope
 from azure.ai.ml._utils._arm_id_utils import (
     AMLNamedArmId,
@@ -44,6 +43,7 @@
 from azure.storage.filedatalake import FileSasPermissions, generate_file_sas
 
 from promptflow._utils.logger_utils import LoggerFactory
+from promptflow.azure._utils._datastore_type import DatastoreType
 
 from ._fileshare_storeage_helper import FlowFileStorageClient
 

src/promptflow-azure/tests/sdk_cli_azure_test/unittests/test_blob_client.py

@@ -3,10 +3,16 @@
 import pytest
 
 from promptflow.azure._storage.blob.client import _datastore_cache, _get_datastore_client_key, _get_datastore_from_cache
+from promptflow.azure._utils._datastore_type import DatastoreType
 
 
 @pytest.mark.unittest
 class TestBlobClient:
+    def test_datastore_type_compat(self):
+        assert DatastoreType.AZURE_BLOB.value == "AzureBlob"
+        assert DatastoreType.AZURE_DATA_LAKE_GEN2.value == "AzureDataLakeGen2"
+        assert DatastoreType.AZURE_FILE.value == "AzureFile"
+
     def test_get_datastore_from_cache(self):
         _datastore_cache["test"] = {
             "expire_at": datetime.datetime.now() + datetime.timedelta(0, -1),  # already expire