From b364cd1c298b310d198ed5ab9034b3f42b42935c Mon Sep 17 00:00:00 2001
From: dhruvesh09 <30351549+dhruvesh09@users.noreply.github.com>
Date: Fri, 9 Jul 2021 12:00:56 -0700
Subject: [PATCH] [r0.30.0 cherry-pick] commit_id=[2736197] (#4011)

---
 tfx/orchestration/kubeflow/v2/step_builder.py                | 5 +++--
 .../v2/testdata/expected_full_taxi_pipeline_job.json         | 2 +-
 .../expected_latest_blessed_model_resolver_executor.pbtxt    | 2 +-
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/tfx/orchestration/kubeflow/v2/step_builder.py b/tfx/orchestration/kubeflow/v2/step_builder.py
index f2641ec082..b56fa0bebf 100644
--- a/tfx/orchestration/kubeflow/v2/step_builder.py
+++ b/tfx/orchestration/kubeflow/v2/step_builder.py
@@ -586,8 +586,9 @@ def _build_resolver_for_latest_blessed_model(
     artifact_queries = {}
     query_filter = (
         'schema_title="{type}" AND '
-        'state={state} AND name={{$.inputs.artifacts["{input_key}"]'
-        '.metadata.{property_key}.number_value}}').format(
+        'state={state} AND '
+        'name="{{{{$.inputs.artifacts[\'{input_key}\']'
+        '.metadata[\'{property_key}\']}}}}"').format(
             type=compiler_utils.get_artifact_title(standard_artifacts.Model),
             state=metadata_store_pb2.Artifact.State.Name(
                 metadata_store_pb2.Artifact.LIVE),
diff --git a/tfx/orchestration/kubeflow/v2/testdata/expected_full_taxi_pipeline_job.json b/tfx/orchestration/kubeflow/v2/testdata/expected_full_taxi_pipeline_job.json
index c1c89bb1c4..33c9ec0627 100644
--- a/tfx/orchestration/kubeflow/v2/testdata/expected_full_taxi_pipeline_job.json
+++ b/tfx/orchestration/kubeflow/v2/testdata/expected_full_taxi_pipeline_job.json
@@ -57,7 +57,7 @@
           "resolver": {
             "outputArtifactQueries": {
               "model": {
-                "filter": "schema_title=\"tfx.Model\" AND state=LIVE AND name={$.inputs.artifacts[\"input\"].metadata.current_model_id.number_value}"
+                "filter": "schema_title=\"tfx.Model\" AND state=LIVE AND name=\"{{$.inputs.artifacts['input'].metadata['current_model_id']}}\""
               }
             }
           }
diff --git a/tfx/orchestration/kubeflow/v2/testdata/expected_latest_blessed_model_resolver_executor.pbtxt b/tfx/orchestration/kubeflow/v2/testdata/expected_latest_blessed_model_resolver_executor.pbtxt
index 5323039ac7..77bde09f0a 100644
--- a/tfx/orchestration/kubeflow/v2/testdata/expected_latest_blessed_model_resolver_executor.pbtxt
+++ b/tfx/orchestration/kubeflow/v2/testdata/expected_latest_blessed_model_resolver_executor.pbtxt
@@ -21,7 +21,7 @@ executors {
       output_artifact_queries {
         key: "model"
         value {
-          filter: "schema_title=\"tfx.Model\" AND state=LIVE AND name={$.inputs.artifacts[\"input\"].metadata.current_model_id.number_value}"
+          filter: "schema_title=\"tfx.Model\" AND state=LIVE AND name=\"{{$.inputs.artifacts['input'].metadata['current_model_id']}}\""
         }
       }
     }