feat: Support ResourceName with Version. (#1609)

* feat: Support ResourceName with Version.

* remove extra line

* add unit tests for resouce_name version support

google/cloud/aiplatform/metadata/resource.py

@@ -539,7 +539,7 @@ def _extract_metadata_store_id(resource_name, resource_noun) -> str:
         pattern = re.compile(
             r"^projects\/(?P<project>[\w-]+)\/locations\/(?P<location>[\w-]+)\/metadataStores\/(?P<store>[\w-]+)\/"
             + resource_noun
-            + r"\/(?P<id>[\w-]+)$"
+            + r"\/(?P<id>[\w-]+)(?P<version>@[\w-]+)?$"
         )
         match = pattern.match(resource_name)
         if not match:

google/cloud/aiplatform/metadata/schema/utils.py

@@ -143,12 +143,15 @@ def to_dict(self):
         return results
 
 
-def create_uri_from_resource_name(resource_name: str) -> bool:
+def create_uri_from_resource_name(resource_name: str) -> str:
     """Construct the service URI for a given resource_name.
     Args:
         resource_name (str):
-            The name of the Vertex resource, in a form of
+            The name of the Vertex resource, in one of the forms:
             projects/{project}/locations/{location}/{resource_type}/{resource_id}
+            projects/{project}/locations/{location}/{resource_type}/{resource_id}@{version}
+            projects/{project}/locations/{location}/metadataStores/{store_id}/{resource_type}/{resource_id}
+            projects/{project}/locations/{location}/metadataStores/{store_id}/{resource_type}/{resource_id}@{version}
     Returns:
         The resource URI in the form of:
         https://{service-endpoint}/v1/{resource_name},
@@ -159,11 +162,11 @@ def create_uri_from_resource_name(resource_name: str) -> bool:
     """
     # TODO: support nested resource names such as models/123/evaluations/456
     match_results = re.match(
-        r"^projects\/[A-Za-z0-9-]*\/locations\/([A-Za-z0-9-]*)(\/metadataStores\/[A-Za-z0-9-]*)?(\/[A-Za-z0-9-]*\/[A-Za-z0-9-]*)+$",
+        r"^projects\/(?P<project>[\w-]+)\/locations\/(?P<location>[\w-]+)(\/metadataStores\/(?P<store>[\w-]+))?\/[\w-]+\/(?P<id>[\w-]+)(?P<version>@[\w-]+)?$",
         resource_name,
     )
     if not match_results:
         raise ValueError(f"Invalid resource_name format for {resource_name}.")
 
-    location = match_results.group(1)
+    location = match_results["location"]
     return f"https://{location}-aiplatform.googleapis.com/v1/{resource_name}"

tests/unit/aiplatform/test_metadata_resources.py

@@ -892,6 +892,26 @@ def test_init_artifact_with_id(self, get_artifact_mock):
             name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY
         )
 
+    def test_extract_metadata_store_id_with_valid_resource_name(self):
+        aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+        assert "store_id" == artifact.Artifact._extract_metadata_store_id(
+            resource_name="projects/project/locations/location/metadataStores/store_id/resource_type/resource_id",
+            resource_noun="resource_type",
+        )
+        assert "store_id" == artifact.Artifact._extract_metadata_store_id(
+            resource_name="projects/project/locations/location/metadataStores/store_id/resource_type/resource_id@version",
+            resource_noun="resource_type",
+        )
+
+    def test_extract_metadata_store_id_with_invalid_resource_name(self):
+        invalid_resouce_name = (
+            "projects/project/locations/location/resource_type/resource_id/"
+        )
+        with pytest.raises(ValueError):
+            artifact.Artifact._extract_metadata_store_id(
+                resource_name=invalid_resouce_name, resource_noun="resource_type"
+            )
+
     def test_get_or_create_artifact(
         self, get_artifact_for_get_or_create_mock, create_artifact_mock
     ):

tests/unit/aiplatform/test_metadata_schema.py

@@ -1093,6 +1093,26 @@ def test_predict_schemata_to_dict_method_returns_correct_schema(self):
 
         assert json.dumps(predict_schema_ta.to_dict()) == json.dumps(expected_results)
 
+    def test_create_uri_from_resource_name_for_valid_resouce_names(self):
+        valid_resouce_names = [
+            "projects/project/locations/location/resource_type/resource_id",
+            "projects/project/locations/location/resource_type/resource_id@version",
+            "projects/project/locations/location/metadataStores/store_id/resource_type/resource_id",
+            "projects/project/locations/location/metadataStores/store_id/resource_type/resource_id@version",
+        ]
+        for resouce_name in valid_resouce_names:
+            uri = utils.create_uri_from_resource_name(resource_name=resouce_name)
+            assert (
+                uri == f"https://location-aiplatform.googleapis.com/v1/{resouce_name}"
+            )
+
+    def test_create_uri_from_resource_name_for_invalid_resouce_names(self):
+        invalid_resouce_name = (
+            "projects/project/locations/location/resource_type/resource_id/"
+        )
+        with pytest.raises(ValueError):
+            utils.create_uri_from_resource_name(resource_name=invalid_resouce_name)
+
     def test_container_spec_to_dict_method_returns_correct_schema(self):
         container_spec = utils.ContainerSpec(
             image_uri="gcr.io/some_container_image_uri",