[ML] Enable SDK data and datastore get calls to use new error message…

… format (#26788) * enable SDK data and datastore get calls to use new error message format
Azure · Oct 17, 2022 · 9c473c8 · 9c473c8
1 parent 60ef9f0
commit 9c473c8
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 50 deletions.
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py
@@ -121,37 +121,39 @@ def get(self, name: str, version: Optional[str] = None, label: Optional[str] = N
         :return: Data asset object.
         :rtype: ~azure.ai.ml.entities.Data
         """
-        if version and label:
-            msg = "Cannot specify both version and label."
-            raise ValidationException(
-                message=msg,
-                target=ErrorTarget.DATA,
-                no_personal_data_message=msg,
-                error_category=ErrorCategory.USER_ERROR,
-                error_type=ValidationErrorType.INVALID_VALUE,
-            )
-
-        if label:
-            return _resolve_label_to_asset(self, name, label)
+        try:
+            if version and label:
+                msg = "Cannot specify both version and label."
+                raise ValidationException(
+                    message=msg,
+                    target=ErrorTarget.DATA,
+                    no_personal_data_message=msg,
+                    error_category=ErrorCategory.USER_ERROR,
+                    error_type=ValidationErrorType.INVALID_VALUE,
+                )
 
-        if not version:
-            msg = "Must provide either version or label."
-            raise ValidationException(
-                message=msg,
-                target=ErrorTarget.DATA,
-                no_personal_data_message=msg,
-                error_category=ErrorCategory.USER_ERROR,
-                error_type=ValidationErrorType.MISSING_FIELD,
+            if label:
+                return _resolve_label_to_asset(self, name, label)
+
+            if not version:
+                msg = "Must provide either version or label."
+                raise ValidationException(
+                    message=msg,
+                    target=ErrorTarget.DATA,
+                    no_personal_data_message=msg,
+                    error_category=ErrorCategory.USER_ERROR,
+                    error_type=ValidationErrorType.MISSING_FIELD,
+                )
+            data_version_resource = self._operation.get(
+                resource_group_name=self._resource_group_name,
+                workspace_name=self._workspace_name,
+                name=name,
+                version=version,
+                **self._init_kwargs,
             )
-        data_version_resource = self._operation.get(
-            resource_group_name=self._resource_group_name,
-            workspace_name=self._workspace_name,
-            name=name,
-            version=version,
-            **self._init_kwargs,
-        )
-
-        return Data._from_rest_object(data_version_resource)
+            return Data._from_rest_object(data_version_resource)
+        except (ValidationException, SchemaValidationError) as ex:
+            log_and_raise_error(ex)
 
     # @monitor_with_activity(logger, "Data.CreateOrUpdate", ActivityType.PUBLICAPI)
     def create_or_update(self, data: Data) -> Data:

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_datastore_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_datastore_operations.py
@@ -101,16 +101,18 @@ def get(self, name: str, *, include_secrets: bool = False) -> Datastore:
         :return: Datastore with the specified name.
         :rtype: Datastore
         """
-
-        datastore_resource = self._operation.get(
-            name=name,
-            resource_group_name=self._operation_scope.resource_group_name,
-            workspace_name=self._workspace_name,
-            **self._init_kwargs
-        )
-        if include_secrets:
-            self._fetch_and_populate_secret(datastore_resource)
-        return Datastore._from_rest_object(datastore_resource)
+        try:
+            datastore_resource = self._operation.get(
+                name=name,
+                resource_group_name=self._operation_scope.resource_group_name,
+                workspace_name=self._workspace_name,
+                **self._init_kwargs
+            )
+            if include_secrets:
+                self._fetch_and_populate_secret(datastore_resource)
+            return Datastore._from_rest_object(datastore_resource)
+        except (ValidationException, SchemaValidationError) as ex:
+            log_and_raise_error(ex)
 
     def _fetch_and_populate_secret(self, datastore_resource: DatastoreData) -> None:
         if datastore_resource.name and not isinstance(
@@ -128,16 +130,18 @@ def get_default(self, *, include_secrets: bool = False) -> Datastore:
         :return: The default datastore.
         :rtype: Datastore
         """
-
-        datastore_resource = self._operation.list(
-            resource_group_name=self._operation_scope.resource_group_name,
-            workspace_name=self._workspace_name,
-            is_default=True,
-            **self._init_kwargs
-        ).next()
-        if include_secrets:
-            self._fetch_and_populate_secret(datastore_resource)
-        return Datastore._from_rest_object(datastore_resource)
+        try:
+            datastore_resource = self._operation.list(
+                resource_group_name=self._operation_scope.resource_group_name,
+                workspace_name=self._workspace_name,
+                is_default=True,
+                **self._init_kwargs
+            ).next()
+            if include_secrets:
+                self._fetch_and_populate_secret(datastore_resource)
+            return Datastore._from_rest_object(datastore_resource)
+        except (ValidationException, SchemaValidationError) as ex:
+            log_and_raise_error(ex)
 
     # @monitor_with_activity(logger, "Datastore.CreateOrUpdate", ActivityType.PUBLICAPI)
     def create_or_update(self, datastore: Datastore) -> Datastore:

diff --git a/sdk/ml/azure-ai-ml/tests/dataset/unittests/test_data_operations.py b/sdk/ml/azure-ai-ml/tests/dataset/unittests/test_data_operations.py
@@ -91,8 +91,9 @@ def test_get_with_version(self, mock_data_operations: DataOperations) -> None:
 
     def test_get_no_version(self, mock_data_operations: DataOperations) -> None:
         name = "random_name"
-        with pytest.raises(Exception):
+        with pytest.raises(Exception) as ex:
             mock_data_operations.get(name=name)
+        assert "At least one required parameter is missing" in str(ex.value)
 
     def test_create_with_spec_file(
         self,