fix: keep the caller's project, location and credentials when metadata schemas reload

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

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/<metadata_store_id>/contexts/<resource_id>.
+            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

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/<metadata_store_id>/executions/<resource_id>.
+            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
 

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,
+        )