From 4f7aad89f6ed19188ae23a3246a84b17751de7ff Mon Sep 17 00:00:00 2001 From: abhishekacharya200 Date: Sun, 1 Nov 2026 00:45:08 +1100 Subject: [PATCH] fix: keep the caller's project, location and credentials when metadata schemas reload create and start_execution passed these to the service, then reloaded the new resource without them, so the reload fell back to the values from aiplatform.init. --- .../metadata/schema/base_artifact.py | 8 +- .../metadata/schema/base_context.py | 32 ++++++- .../metadata/schema/base_execution.py | 35 +++++++- tests/unit/aiplatform/test_metadata_schema.py | 84 +++++++++++++++++++ 4 files changed, 153 insertions(+), 6 deletions(-) diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py index 316653b511..f2140265d0 100644 --- a/google/cloud/aiplatform/metadata/schema/base_artifact.py +++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py @@ -197,7 +197,13 @@ def create( ) # Reinstantiate this class using the newly created resource. - self._init_with_resource_name(artifact_name=new_artifact_instance.resource_name) + self._init_with_resource_name( + artifact_name=new_artifact_instance.resource_name, + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, + ) return self @classmethod diff --git a/google/cloud/aiplatform/metadata/schema/base_context.py b/google/cloud/aiplatform/metadata/schema/base_context.py index d22e3ec2b6..798410879c 100644 --- a/google/cloud/aiplatform/metadata/schema/base_context.py +++ b/google/cloud/aiplatform/metadata/schema/base_context.py @@ -89,19 +89,41 @@ def _init_with_resource_name( self, *, context_name: str, + metadata_store_id: str = "default", + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, ): """Initializes the Artifact instance using an existing resource. Args: context_name (str): Context name with the following format, this is globally unique in a metadataStore: projects/123/locations/us-central1/metadataStores//contexts/. + metadata_store_id (str): + Optional. MetadataStore to retrieve the Context from. If not set, metadata_store_id is set to "default". + If context_name is a fully-qualified resource, its metadata_store_id overrides this one. + project (str): + Optional. Project to retrieve the Context from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve the Context from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Context. Overrides + credentials set in aiplatform.init. """ # Add User Agent Header for metrics tracking if one is not specified # If one is already specified this call was initiated by a sub class. if not base_constants.USER_AGENT_SDK_COMMAND: base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_context.BaseContextSchema._init_with_resource_name" - super(BaseContextSchema, self).__init__(resource_name=context_name) + super(BaseContextSchema, self).__init__( + resource_name=context_name, + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, + ) def create( self, @@ -156,7 +178,13 @@ def create( ) # Reinstantiate this class using the newly created resource. - self._init_with_resource_name(context_name=new_context.resource_name) + self._init_with_resource_name( + context_name=new_context.resource_name, + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, + ) return self @classmethod diff --git a/google/cloud/aiplatform/metadata/schema/base_execution.py b/google/cloud/aiplatform/metadata/schema/base_execution.py index c514077c65..c4ad7f474a 100644 --- a/google/cloud/aiplatform/metadata/schema/base_execution.py +++ b/google/cloud/aiplatform/metadata/schema/base_execution.py @@ -95,19 +95,41 @@ def _init_with_resource_name( self, *, execution_name: str, + metadata_store_id: str = "default", + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, ): """Initializes the Execution instance using an existing resource. Args: execution_name (str): The Execution name with the following format, this is globally unique in a metadataStore. projects/123/locations/us-central1/metadataStores//executions/. + metadata_store_id (str): + Optional. MetadataStore to retrieve the Execution from. If not set, metadata_store_id is set to "default". + If execution_name is a fully-qualified resource, its metadata_store_id overrides this one. + project (str): + Optional. Project to retrieve the Execution from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve the Execution from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Execution. Overrides + credentials set in aiplatform.init. """ # Add User Agent Header for metrics tracking if one is not specified # If one is already specified this call was initiated by a sub class. if not base_constants.USER_AGENT_SDK_COMMAND: base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_execution.BaseExecutionSchema._init_with_resource_name" - super(BaseExecutionSchema, self).__init__(execution_name=execution_name) + super(BaseExecutionSchema, self).__init__( + execution_name=execution_name, + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, + ) def create( self, @@ -164,7 +186,11 @@ def create( ) # Reinstantiate this class using the newly created resource. self._init_with_resource_name( - execution_name=new_execution_instance.resource_name + execution_name=new_execution_instance.resource_name, + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, ) return self @@ -304,7 +330,10 @@ def start_execution( # Reinstantiate this class using the newly created resource. self._init_with_resource_name( - execution_name=new_execution_instance.resource_name + execution_name=new_execution_instance.resource_name, + project=project, + location=location, + credentials=credentials, ) return self diff --git a/tests/unit/aiplatform/test_metadata_schema.py b/tests/unit/aiplatform/test_metadata_schema.py index 8251e91332..17eb0dad11 100644 --- a/tests/unit/aiplatform/test_metadata_schema.py +++ b/tests/unit/aiplatform/test_metadata_schema.py @@ -1585,3 +1585,87 @@ class TestExecution(base_execution.BaseExecutionSchema): assert kwargs["execution"].display_name == _TEST_DISPLAY_NAME assert kwargs["execution"].description == _TEST_DESCRIPTION assert kwargs["execution"].metadata == _TEST_UPDATED_METADATA + + +class TestMetadataSchemaReloadUsesCallerSettings: + """The reload after create must use the project, location and credentials given.""" + + SETTINGS = { + "project": "other-project", + "location": "europe-west4", + "credentials": mock.sentinel.credentials, + } + + def test_artifact_create_reloads_with_caller_settings( + self, base_artifact_init_with_resouce_name_mock + ): + class TestArtifact(base_artifact.BaseArtifactSchema): + schema_title = _TEST_SCHEMA_TITLE + + with mock.patch( + "google.cloud.aiplatform.metadata.artifact.Artifact.create", + return_value=mock.Mock(resource_name=_TEST_ARTIFACT_NAME), + ): + TestArtifact().create(metadata_store_id="store", **self.SETTINGS) + + base_artifact_init_with_resouce_name_mock.assert_called_once_with( + artifact_name=_TEST_ARTIFACT_NAME, + metadata_store_id="store", + **self.SETTINGS, + ) + + def test_execution_create_reloads_with_caller_settings( + self, base_execution_init_with_resouce_name_mock + ): + class TestExecution(base_execution.BaseExecutionSchema): + schema_title = _TEST_SCHEMA_TITLE + + with mock.patch( + "google.cloud.aiplatform.metadata.execution.Execution.create", + return_value=mock.Mock(resource_name=_TEST_EXECUTION_NAME), + ): + TestExecution().create(metadata_store_id="store", **self.SETTINGS) + + base_execution_init_with_resouce_name_mock.assert_called_once_with( + execution_name=_TEST_EXECUTION_NAME, + metadata_store_id="store", + **self.SETTINGS, + ) + + def test_start_execution_reloads_with_caller_settings( + self, base_execution_init_with_resouce_name_mock + ): + class TestExecution(base_execution.BaseExecutionSchema): + schema_title = _TEST_SCHEMA_TITLE + + with mock.patch( + "google.cloud.aiplatform.metadata.metadata." + "_ExperimentTracker.start_execution", + return_value=mock.Mock(resource_name=_TEST_EXECUTION_NAME), + ): + TestExecution(metadata=_TEST_UPDATED_METADATA).start_execution( + **self.SETTINGS + ) + + base_execution_init_with_resouce_name_mock.assert_called_once_with( + execution_name=_TEST_EXECUTION_NAME, + **self.SETTINGS, + ) + + def test_context_create_reloads_with_caller_settings( + self, base_context_init_with_resouce_name_mock + ): + class TestContext(base_context.BaseContextSchema): + schema_title = _TEST_SCHEMA_TITLE + + with mock.patch( + "google.cloud.aiplatform.metadata.context.Context.create", + return_value=mock.Mock(resource_name=_TEST_CONTEXT_NAME), + ): + TestContext().create(metadata_store_id="store", **self.SETTINGS) + + base_context_init_with_resouce_name_mock.assert_called_once_with( + context_name=_TEST_CONTEXT_NAME, + metadata_store_id="store", + **self.SETTINGS, + )