From 371fe67b1aa8e584f0a92dac68a815d2ad8b6232 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=86=AF=E5=9F=BA=E9=AD=81?= <1412414664@qq.com>
Date: Mon, 8 Jun 2026 00:26:04 +0800
Subject: [PATCH] fix(azure): update DatastoreType import

---
 .../promptflow/azure/_storage/blob/client.py           |  2 +-
 .../promptflow/azure/_utils/_datastore_type.py         | 10 ++++++++++
 .../promptflow/azure/operations/_artifact_utilities.py |  2 +-
 .../sdk_cli_azure_test/unittests/test_blob_client.py   |  6 ++++++
 4 files changed, 18 insertions(+), 2 deletions(-)
 create mode 100644 src/promptflow-azure/promptflow/azure/_utils/_datastore_type.py

diff --git a/src/promptflow-azure/promptflow/azure/_storage/blob/client.py b/src/promptflow-azure/promptflow/azure/_storage/blob/client.py
index 75f9e2bbb4c..bccb1af33db 100644
--- a/src/promptflow-azure/promptflow/azure/_storage/blob/client.py
+++ b/src/promptflow-azure/promptflow/azure/_storage/blob/client.py
@@ -6,7 +6,6 @@
 
 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
@@ -14,6 +13,7 @@
 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()
diff --git a/src/promptflow-azure/promptflow/azure/_utils/_datastore_type.py b/src/promptflow-azure/promptflow/azure/_utils/_datastore_type.py
new file mode 100644
index 00000000000..a8d0143a541
--- /dev/null
+++ b/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"]
diff --git a/src/promptflow-azure/promptflow/azure/operations/_artifact_utilities.py b/src/promptflow-azure/promptflow/azure/operations/_artifact_utilities.py
index c4c5cf23011..f1bddd6b5b2 100644
--- a/src/promptflow-azure/promptflow/azure/operations/_artifact_utilities.py
+++ b/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
 
diff --git a/src/promptflow-azure/tests/sdk_cli_azure_test/unittests/test_blob_client.py b/src/promptflow-azure/tests/sdk_cli_azure_test/unittests/test_blob_client.py
index c39e0425ee0..5810f5b1493 100644
--- a/src/promptflow-azure/tests/sdk_cli_azure_test/unittests/test_blob_client.py
+++ b/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