From 737dc2bacd89e3c9af06bb408abe5ed9304a1397 Mon Sep 17 00:00:00 2001 From: sina chavoshi Date: Wed, 24 Aug 2022 15:40:32 -0700 Subject: [PATCH] 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 | 2 +- .../cloud/aiplatform/metadata/schema/utils.py | 11 ++++++---- .../aiplatform/test_metadata_resources.py | 20 +++++++++++++++++++ tests/unit/aiplatform/test_metadata_schema.py | 20 +++++++++++++++++++ 4 files changed, 48 insertions(+), 5 deletions(-) diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py index 00133ba789..84fefc2131 100644 --- a/google/cloud/aiplatform/metadata/resource.py +++ b/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[\w-]+)\/locations\/(?P[\w-]+)\/metadataStores\/(?P[\w-]+)\/" + resource_noun - + r"\/(?P[\w-]+)$" + + r"\/(?P[\w-]+)(?P@[\w-]+)?$" ) match = pattern.match(resource_name) if not match: diff --git a/google/cloud/aiplatform/metadata/schema/utils.py b/google/cloud/aiplatform/metadata/schema/utils.py index c742824b76..1b4a5e4f6c 100644 --- a/google/cloud/aiplatform/metadata/schema/utils.py +++ b/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[\w-]+)\/locations\/(?P[\w-]+)(\/metadataStores\/(?P[\w-]+))?\/[\w-]+\/(?P[\w-]+)(?P@[\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}" diff --git a/tests/unit/aiplatform/test_metadata_resources.py b/tests/unit/aiplatform/test_metadata_resources.py index 247657e08f..733eb1dd3f 100644 --- a/tests/unit/aiplatform/test_metadata_resources.py +++ b/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 ): diff --git a/tests/unit/aiplatform/test_metadata_schema.py b/tests/unit/aiplatform/test_metadata_schema.py index 65de1304ea..0003838968 100644 --- a/tests/unit/aiplatform/test_metadata_schema.py +++ b/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",