feat: Add Metadata SDK support and samples for get method (#1516)

* feat: Add get() method to Metadata Resource base class.

* Add unit tests for artifact, execution, and context

* Add samples for get execution and get artifact

* fix lint issues

* Fix unit tests.

Author: SinaChavoshi

google/cloud/aiplatform/metadata/resource.py

@@ -189,6 +189,50 @@ def get_or_create(
             )
         return resource
 
+    @classmethod
+    def get(
+        cls,
+        resource_id: str,
+        metadata_store_id: str = "default",
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> "_Resource":
+        """Retrieves a Metadata resource.
+
+        Args:
+            resource_id (str):
+                Required. The <resource_id> portion of the resource name with the format:
+                projects/123/locations/us-central1/metadataStores/<metadata_store_id>/<resource_noun>/<resource_id>.
+            metadata_store_id (str):
+                The <metadata_store_id> portion of the resource name with
+                the format:
+                projects/123/locations/us-central1/metadataStores/<metadata_store_id>/<resource_noun>/<resource_id>
+                If not provided, the MetadataStore's ID will be set to "default".
+            project (str):
+                Project used to retrieve or create this resource. Overrides project set in
+                aiplatform.init.
+            location (str):
+                Location used to retrieve or create this resource. Overrides location set in
+                aiplatform.init.
+            credentials (auth_credentials.Credentials):
+                Custom credentials used to retrieve or create this resource. Overrides
+                credentials set in aiplatform.init.
+
+        Returns:
+            resource (_Resource):
+                Instantiated representation of the managed Metadata resource or None if no resouce was found.
+
+        """
+        resource = cls._get(
+            resource_name=resource_id,
+            metadata_store_id=metadata_store_id,
+            project=project,
+            location=location,
+            credentials=credentials,
+        )
+        return resource
+
     def sync_resource(self):
         """Syncs local resource with the resource in metadata store."""
         self._gca_resource = getattr(self.api_client, self._getter_method)(

samples/model-builder/conftest.py

@@ -576,6 +576,13 @@ def mock_get_execution(mock_execution):
         yield mock_get_execution
 
 
+@pytest.fixture
+def mock_execution_get(mock_execution):
+    with patch.object(aiplatform.Execution, "get") as mock_execution_get:
+        mock_execution_get.return_value = mock_execution
+        yield mock_execution_get
+
+
 @pytest.fixture
 def mock_create_execution(mock_execution):
     with patch.object(aiplatform.Execution, "create") as mock_create_execution:
@@ -590,6 +597,13 @@ def mock_get_artifact(mock_artifact):
         yield mock_get_artifact
 
 
+@pytest.fixture
+def mock_artifact_get(mock_artifact):
+    with patch.object(aiplatform.Artifact, "get") as mock_artifact_get:
+        mock_artifact_get.return_value = mock_artifact
+        yield mock_artifact_get
+
+
 @pytest.fixture
 def mock_pipeline_job_create(mock_pipeline_job):
     with patch.object(aiplatform, "PipelineJob") as mock_pipeline_job_create:

samples/model-builder/experiment_tracking/get_artifact_sample.py

@@ -17,12 +17,12 @@
 
 #  [START aiplatform_sdk_get_artifact_sample]
 def get_artifact_sample(
-    uri: str,
+    artifact_id: str,
     project: str,
     location: str,
 ):
-    artifact = aiplatform.Artifact.get_with_uri(
-        uri=uri, project=project, location=location
+    artifact = aiplatform.Artifact.get(
+        resource_id=artifact_id, project=project, location=location
     )
 
     return artifact

samples/model-builder/experiment_tracking/get_artifact_sample_test.py

@@ -17,15 +17,15 @@
 import test_constants
 
 
-def test_get_artifact_sample(mock_artifact, mock_get_with_uri):
+def test_get_artifact_sample(mock_artifact, mock_artifact_get):
     artifact = get_artifact_sample.get_artifact_sample(
-        uri=test_constants.MODEL_ARTIFACT_URI,
+        artifact_id=test_constants.RESOURCE_ID,
         project=test_constants.PROJECT,
         location=test_constants.LOCATION,
     )
 
-    mock_get_with_uri.assert_called_with(
-        uri=test_constants.MODEL_ARTIFACT_URI,
+    mock_artifact_get.assert_called_with(
+        resource_id=test_constants.RESOURCE_ID,
         project=test_constants.PROJECT,
         location=test_constants.LOCATION,
     )

samples/model-builder/experiment_tracking/get_artifact_with_uri_sample.py

@@ -0,0 +1,31 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_get_artifact_with_uri_sample]
+def get_artifact_with_uri_sample(
+    uri: str,
+    project: str,
+    location: str,
+):
+    artifact = aiplatform.Artifact.get_with_uri(
+        uri=uri, project=project, location=location
+    )
+
+    return artifact
+
+
+#  [END aiplatform_sdk_get_artifact_with_uri_sample]

samples/model-builder/experiment_tracking/get_artifact_with_uri_sample_test.py

@@ -0,0 +1,33 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import get_artifact_with_uri_sample
+
+import test_constants
+
+
+def test_get_artifact_with_uri_sample(mock_artifact, mock_get_with_uri):
+    artifact = get_artifact_with_uri_sample.get_artifact_with_uri_sample(
+        uri=test_constants.MODEL_ARTIFACT_URI,
+        project=test_constants.PROJECT,
+        location=test_constants.LOCATION,
+    )
+
+    mock_get_with_uri.assert_called_with(
+        uri=test_constants.MODEL_ARTIFACT_URI,
+        project=test_constants.PROJECT,
+        location=test_constants.LOCATION,
+    )
+
+    assert artifact is mock_artifact

samples/model-builder/experiment_tracking/get_execution_sample.py

@@ -0,0 +1,31 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_get_execution_sample]
+def get_execution_sample(
+    execution_id: str,
+    project: str,
+    location: str,
+):
+    execution = aiplatform.Execution.get(
+        resource_id=execution_id, project=project, location=location
+    )
+
+    return execution
+
+
+#  [END aiplatform_sdk_get_execution_sample]

samples/model-builder/experiment_tracking/get_execution_sample_test.py

@@ -0,0 +1,33 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import get_artifact_sample
+
+import test_constants
+
+
+def test_get_artifact_sample(mock_artifact, mock_artifact_get):
+    artifact = get_artifact_sample.get_artifact_sample(
+        artifact_id=test_constants.RESOURCE_ID,
+        project=test_constants.PROJECT,
+        location=test_constants.LOCATION,
+    )
+
+    mock_artifact_get.assert_called_with(
+        resource_id=test_constants.RESOURCE_ID,
+        project=test_constants.PROJECT,
+        location=test_constants.LOCATION,
+    )
+
+    assert artifact is mock_artifact

tests/unit/aiplatform/test_metadata_resources.py

@@ -419,6 +419,28 @@ def test_get_or_create_context(
         expected_context.name = _TEST_CONTEXT_NAME
         assert my_context._gca_resource == expected_context
 
+    def test_get_context(self, get_context_mock):
+        aiplatform.init(project=_TEST_PROJECT)
+
+        my_context = context.Context.get(
+            resource_id=_TEST_CONTEXT_ID,
+            metadata_store_id=_TEST_METADATA_STORE,
+        )
+
+        expected_context = GapicContext(
+            schema_title=_TEST_SCHEMA_TITLE,
+            schema_version=_TEST_SCHEMA_VERSION,
+            display_name=_TEST_DISPLAY_NAME,
+            description=_TEST_DESCRIPTION,
+            metadata=_TEST_METADATA,
+        )
+        get_context_mock.assert_called_once_with(
+            name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY
+        )
+
+        expected_context.name = _TEST_CONTEXT_NAME
+        assert my_context._gca_resource == expected_context
+
     @pytest.mark.usefixtures("get_context_mock")
     @pytest.mark.usefixtures("create_context_mock")
     def test_update_context(self, update_context_mock):
@@ -633,6 +655,28 @@ def test_get_or_create_execution(
         expected_execution.name = _TEST_EXECUTION_NAME
         assert my_execution._gca_resource == expected_execution
 
+    def test_get_execution(self, get_execution_mock):
+        aiplatform.init(project=_TEST_PROJECT)
+
+        my_execution = execution.Execution.get(
+            resource_id=_TEST_EXECUTION_ID,
+            metadata_store_id=_TEST_METADATA_STORE,
+        )
+
+        expected_execution = GapicExecution(
+            schema_title=_TEST_SCHEMA_TITLE,
+            schema_version=_TEST_SCHEMA_VERSION,
+            display_name=_TEST_DISPLAY_NAME,
+            description=_TEST_DESCRIPTION,
+            metadata=_TEST_METADATA,
+        )
+        get_execution_mock.assert_called_once_with(
+            name=_TEST_EXECUTION_NAME, retry=base._DEFAULT_RETRY
+        )
+
+        expected_execution.name = _TEST_EXECUTION_NAME
+        assert my_execution._gca_resource == expected_execution
+
     @pytest.mark.usefixtures("get_execution_mock")
     @pytest.mark.usefixtures("create_execution_mock")
     def test_update_execution(self, update_execution_mock):
@@ -883,6 +927,28 @@ def test_get_or_create_artifact(
         expected_artifact.name = _TEST_ARTIFACT_NAME
         assert my_artifact._gca_resource == expected_artifact
 
+    def test_get_artifact(self, get_artifact_mock):
+        aiplatform.init(project=_TEST_PROJECT)
+
+        my_artifact = artifact.Artifact.get(
+            resource_id=_TEST_ARTIFACT_ID,
+            metadata_store_id=_TEST_METADATA_STORE,
+        )
+
+        expected_artifact = GapicArtifact(
+            schema_title=_TEST_SCHEMA_TITLE,
+            schema_version=_TEST_SCHEMA_VERSION,
+            display_name=_TEST_DISPLAY_NAME,
+            description=_TEST_DESCRIPTION,
+            metadata=_TEST_METADATA,
+        )
+        get_artifact_mock.assert_called_once_with(
+            name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY
+        )
+
+        expected_artifact.name = _TEST_ARTIFACT_NAME
+        assert my_artifact._gca_resource == expected_artifact
+
     @pytest.mark.usefixtures("get_artifact_mock")
     @pytest.mark.usefixtures("create_artifact_mock")
     def test_update_artifact(self, update_artifact_mock):