Fix issue where invalid artifact proto could be passed to MLMD.

This fixes #2598.

PiperOrigin-RevId: 336965037

RELEASE.md

@@ -27,6 +27,8 @@
     tfx/examples.
 *   Reuse Examples artifact type introduced in TFX 0.23 to allow older release
     jobs running together with TFX 0.23+ release.
+*   Fixed issues where custom property access of a missing property created an
+    invalid MLMD Artifact protobuf message.
 
 ### For pipeline authors
 

tfx/types/artifact.py

@@ -468,10 +468,14 @@ def has_custom_property(self, key: Text) -> bool:
 
   def get_string_custom_property(self, key: Text) -> Text:
     """Get a custom property of string type."""
+    if key not in self._artifact.custom_properties:
+      return ''
     return self._artifact.custom_properties[key].string_value
 
   def get_int_custom_property(self, key: Text) -> int:
     """Get a custom property of int type."""
+    if key not in self._artifact.custom_properties:
+      return 0
     return self._artifact.custom_properties[key].int_value
 
   def copy_from(self, other: 'Artifact'):

tfx/types/artifact_test.py

@@ -279,6 +279,9 @@ def testArtifactProperties(self):
     self.assertEqual(my_artifact.int2, 222)
     self.assertEqual(my_artifact.string1, '111')
     self.assertEqual(my_artifact.string2, '222')
+    self.assertEqual(my_artifact.get_string_custom_property('invalid'), '')
+    self.assertEqual(my_artifact.get_int_custom_property('invalid'), 0)
+    self.assertNotIn('invalid', my_artifact._artifact.custom_properties)
 
     with self.assertRaisesRegexp(
         AttributeError, "Cannot set unknown property 'invalid' on artifact"):