From aa047c0d28504ea220946054746f32ad879df9fe Mon Sep 17 00:00:00 2001
From: tfx-team <tensorflow-extended-nonhuman@googlegroups.com>
Date: Thu, 25 Apr 2024 16:37:02 -0700
Subject: [PATCH] Add a property external_id to Artifact class

PiperOrigin-RevId: 628227490
---
 tfx/types/artifact.py      | 11 +++++++++++
 tfx/types/artifact_test.py | 10 ++++++++++
 2 files changed, 21 insertions(+)

diff --git a/tfx/types/artifact.py b/tfx/types/artifact.py
index 9ca5455b60..92ae830004 100644
--- a/tfx/types/artifact.py
+++ b/tfx/types/artifact.py
@@ -635,6 +635,17 @@ def producer_component(self, producer_component: str):
     """Set producer component of the artifact."""
     self._set_system_property('producer_component', producer_component)
 
+  @property
+  @doc_controls.do_not_doc_in_subclasses
+  def external_id(self) -> str:
+    """external id of the underlying artifact."""
+    return self._artifact.external_id
+
+  @external_id.setter
+  def external_id(self, external_id: str):
+    """Set external id of the underlying artifact."""
+    self._artifact.external_id = external_id
+
   # LINT.IfChange
   @property
   @doc_controls.do_not_doc_in_subclasses
diff --git a/tfx/types/artifact_test.py b/tfx/types/artifact_test.py
index 95e7ee7b50..c5713636fd 100644
--- a/tfx/types/artifact_test.py
+++ b/tfx/types/artifact_test.py
@@ -173,6 +173,7 @@ def testArtifact(self):
     self.assertEqual('MyTypeName', instance.type_name)
     self.assertEqual('', instance.state)
     self.assertFalse(instance.is_external)
+    self.assertEqual('', instance.external_id)
 
     # Default property does not have span or split_names.
     with self.assertRaisesRegex(AttributeError, "has no property 'span'"):
@@ -229,6 +230,14 @@ def testArtifact(self):
     )
     self.assertFalse(instance.get_bool_custom_property('fake_key'))
 
+    instance.mlmd_artifact.external_id = (
+        'mlmd://prod:owner/project_name:pipeline_name:type:artifact:100'
+    )
+    self.assertEqual(
+        'mlmd://prod:owner/project_name:pipeline_name:type:artifact:100',
+        instance.external_id,
+    )
+
     self.assertEqual(
         textwrap.dedent("""\
         Artifact(artifact: id: 1
@@ -272,6 +281,7 @@ def testArtifact(self):
         }
         state: DELETED
         name: "test_artifact"
+        external_id: "mlmd://prod:owner/project_name:pipeline_name:type:artifact:100"
         , artifact_type: name: "MyTypeName"
         properties {
           key: "bool1"