Bug 3323988: Regex fix and indices correction for model download (#36317)

* Bug 3323988: Regex fix and indices correction for model download

* fixing test case

Author: kshitij-microsoft

sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_constants.py

@@ -41,7 +41,7 @@
 GEN2_STORAGE_CLIENT_NAME = "Gen2StorageClient"
 DEFAULT_CONNECTION_TIMEOUT = 14400
 STORAGE_URI_REGEX = (
-    r"(https:\/\/([a-zA-Z0-9@:%_\\\-+~#?&=]+)[a-zA-Z0-9@:%._\\\-+~#?&=]+\.?)\/([a-zA-Z0-9@:%._\\\-+~#?&=]+)\/(.*)"
+    r"(https:\/\/([a-zA-Z0-9@:%_\\\-+~#?&=]+)[a-zA-Z0-9@:%._\\\-+~#?&=]+\.?)\/([a-zA-Z0-9@:%._\\\-+~#?&=]+)\/?(.*)"
 )
 
 WORKSPACE_MANAGED_DATASTORE_WITH_SLASH = "azureml://datastores/workspacemanageddatastore/"

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_storage_utils.py

@@ -191,7 +191,10 @@ def get_ds_name_and_path_prefix(asset_uri: str, registry_name: Optional[str] = N
     if registry_name:
         try:
             split_paths = re.findall(STORAGE_URI_REGEX, asset_uri)
-            path_prefix = split_paths[0][3]
+            if split_paths[0][3] == "":
+                path_prefix = split_paths[0][2]
+            else:
+                path_prefix = split_paths[0][3]
         except Exception as e:
             msg = "Registry asset URI could not be parsed."
             raise MlException(message=msg, no_personal_data_message=msg) from e

sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_storage_utils.py

@@ -1,20 +1,15 @@
-import pytest
-
-from azure.ai.ml._utils._storage_utils import get_ds_name_and_path_prefix
 from pathlib import Path
 from unittest.mock import Mock, patch
 
+import pytest
+
 from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationScope
+from azure.ai.ml._utils._storage_utils import get_ds_name_and_path_prefix
 from azure.ai.ml.entities._assets import Code, Data, Environment, Model
 from azure.ai.ml.entities._assets._artifacts.artifact import ArtifactStorageInfo
-from azure.ai.ml.operations import (
-    DataOperations,
-    DatastoreOperations,
-    EnvironmentOperations,
-    ModelOperations,
-)
-from azure.ai.ml.operations._code_operations import CodeOperations
 from azure.ai.ml.exceptions import ErrorTarget
+from azure.ai.ml.operations import DataOperations, DatastoreOperations, EnvironmentOperations, ModelOperations
+from azure.ai.ml.operations._code_operations import CodeOperations
 
 
 @pytest.fixture
@@ -125,7 +120,7 @@ def test_storage_uri_to_prefix(
     def test_storage_uri_to_prefix_malformed(
         self,
     ) -> None:
-        reg_uri_bad = "https://ccccccccddd4512d.blob.core.windows.net/5823bbb4-bb28-497c-b9f2-1ff3a0778b10"
+        reg_uri_bad = "https://ccccccccddd4512d.blob.core.windows.net"
         workspace_uri_bad = "azureml://subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/000000000000000/workspaces/some_test_3/datastores/workspaceblobstore/path/LocalUpload/26960525964086056a7301dd061fb9be/lightgbm_mlflow_model"
 
         with pytest.raises(Exception) as e: