Support for data share() (Azure#29894)

* first

* data

* ggg

* remove the testcase that is not valid anymore

* import order

Author: banibrata-de

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py

@@ -7,8 +7,10 @@
 import os
 from pathlib import Path
 from typing import Dict, List, Optional, Union, Iterable
+from contextlib import contextmanager
 
 from marshmallow.exceptions import ValidationError as SchemaValidationError
+from azure.ai.ml._utils._registry_utils import get_registry_client
 
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml.entities import PipelineJob, PipelineJobSettings
@@ -278,11 +280,16 @@ def create_or_update(self, data: Data) -> Data:
                             target=ErrorTarget.DATA,
                             error_category=ErrorCategory.USER_ERROR,
                         )
-                    data = data._to_rest_object()
+                    data_res_obj = data._to_rest_object()
                     result = self._service_client.resource_management_asset_reference.begin_import_method(
-                        resource_group_name=self._resource_group_name, registry_name=self._registry_name, body=data
-                    )
-                    return result
+                        resource_group_name=self._resource_group_name,
+                        registry_name=self._registry_name,
+                        body=data_res_obj,
+                    ).result()
+
+                    if not result:
+                        data_res_obj = self._get(name=data.name, version=data.version)
+                        return Data._from_rest_object(data_res_obj)
 
                 sas_uri = get_sas_uri_for_registry_asset(
                     service_client=self._service_client,
@@ -544,19 +551,24 @@ def _get_latest_version(self, name: str) -> Data:
         )
         return self.get(name, version=latest_version)
 
-    # pylint: disable=no-self-use
-    def _prepare_to_copy(
-        self, data: Data, name: Optional[str] = None, version: Optional[str] = None
-    ) -> WorkspaceAssetReference:
-        """Returns WorkspaceAssetReference to copy a registered data to registry given the asset id.
+    @monitor_with_activity(logger, "data.Share", ActivityType.PUBLICAPI)
+    def share(self, name, version, *, share_with_name, share_with_version, registry_name) -> Data:
+        """Share a data asset from workspace to registry.
 
-        :param data: Registered data
-        :type data: Data
-        :param name: Destination name
+        :param name: Name of data asset.
         :type name: str
-        :param version: Destination version
+        :param version: Version of data asset.
         :type version: str
+        :param share_with_name: Name of data asset to share with.
+        :type share_with_name: str
+        :param share_with_version: Version of data asset to share with.
+        :type share_with_version: str
+        :param registry_name: Name of the destination registry.
+        :type registry_name: str
+        :return: Data asset object.
+        :rtype: ~azure.ai.ml.entities.Data
         """
+
         #  Get workspace info to get workspace GUID
         workspace = self._service_client.workspaces.get(
             resource_group_name=self._resource_group_name, workspace_name=self._workspace_name
@@ -569,16 +581,47 @@ def _prepare_to_copy(
             workspace_location,
             workspace_guid,
             AzureMLResourceType.DATA,
-            data.name,
-            data.version,
+            name,
+            version,
         )
 
-        return WorkspaceAssetReference(
-            name=name if name else data.name,
-            version=version if version else data.version,
+        data_ref = WorkspaceAssetReference(
+            name=share_with_name if share_with_name else name,
+            version=share_with_version if share_with_version else version,
             asset_id=asset_id,
         )
 
+        with self._set_registry_client(registry_name):
+            return self.create_or_update(data_ref)
+
+    @contextmanager
+    def _set_registry_client(self, registry_name: str) -> None:
+        """Sets the registry client for the data operations.
+
+        :param registry_name: Name of the registry.
+        :type registry_name: str
+        """
+        rg_ = self._operation_scope._resource_group_name
+        sub_ = self._operation_scope._subscription_id
+        registry_ = self._operation_scope.registry_name
+        client_ = self._service_client
+        data_versions_operation_ = self._operation
+
+        try:
+            _client, _rg, _sub = get_registry_client(self._service_client._config.credential, registry_name)
+            self._operation_scope.registry_name = registry_name
+            self._operation_scope._resource_group_name = _rg
+            self._operation_scope._subscription_id = _sub
+            self._service_client = _client
+            self._operation = _client.data_versions
+            yield
+        finally:
+            self._operation_scope.registry_name = registry_
+            self._operation_scope._resource_group_name = rg_
+            self._operation_scope._subscription_id = sub_
+            self._service_client = client_
+            self._operation = data_versions_operation_
+
 
 def _assert_local_path_matches_asset_type(
     local_path: str,

sdk/ml/azure-ai-ml/tests/dataset/unittests/test_data_operations.py

@@ -563,40 +563,3 @@ def test_create_with_datastore(
                 show_progress=True,
                 ignore_file=None,
             )
-
-    def test_promote_data_from_workspace(
-        self, mock_data_operations_in_registry: DataOperations, mock_data_operations: DataOperations, tmp_path: Path
-    ) -> None:
-        data_asset_name = f"data_random_string"
-        p = tmp_path / "data_full.yml"
-        data_path = tmp_path / "data.pkl"
-        data_path.write_text("hello world")
-        p.write_text(
-            f"""
-    name: {data_asset_name}
-    path: ./data.pkl
-    version: 3"""
-        )
-
-        with patch(
-            "azure.ai.ml._artifacts._artifact_utilities._upload_to_datastore",
-            return_value=ArtifactStorageInfo(
-                name=data_asset_name,
-                version="3",
-                relative_path="path",
-                datastore_arm_id="/subscriptions/mock/resourceGroups/mock/providers/Microsoft.MachineLearningServices/workspaces/mock/datastores/datastore_id",
-                container_name="containerName",
-            ),
-        ) as mock_upload, patch(
-            "azure.ai.ml.operations._data_operations.Data._from_rest_object",
-            return_value=Data(),
-        ):
-            data = load_data(source=p)
-            data_to_promote = mock_data_operations._prepare_to_copy(data, "new_name", "new_version")
-            assert data_to_promote.name == "new_name"
-            assert data_to_promote.version == "new_version"
-            mock_data_operations_in_registry._operation.get.side_effect = Mock(
-                side_effect=ResourceNotFoundError("Test")
-            )
-            mock_data_operations_in_registry.create_or_update(data_to_promote)
-            mock_data_operations_in_registry._service_client.resource_management_asset_reference.begin_import_method.assert_called_once()