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 index fefa6a9d6eef..ced22578d307 100644 --- 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 index 43b019438a86..84feb7df6298 100644 --- 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 index 83f79aa20ee9..9518ec1c3e22 100644 --- 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,