# Buggy code
def Do(
    self,
    input_dict: Dict[Text, List[types.Artifact]],
    output_dict: Dict[Text, List[types.Artifact]],
    exec_properties: Dict[Text, Any],
):
    """Overrides the tfx_pusher_executor.

    Args:
      input_dict: Input dict from input key to a list of artifacts, including:
        - model_export: exported model from trainer.
        - model_blessing: model blessing path from model_validator.
      output_dict: Output dict from key to a list of artifacts, including:
        - model_push: A list of 'ModelPushPath' artifact of size one. It will
          include the model in this push execution if the model was pushed.
      exec_properties: Mostly a passthrough input dict for
        tfx.components.Pusher.executor.  custom_config.bigquery_serving_args is
        consumed by this class.  For the full set of parameters supported by
        Big Query ML, refer to https://cloud.google.com/bigquery-ml/

    Returns:
      None
    Raises:
      ValueError:
        If bigquery_serving_args is not in exec_properties.custom_config.
        If pipeline_root is not 'gs://...'
      RuntimeError: if the Big Query job failed.
    """
    self._log_startup(input_dict, output_dict, exec_properties)
    model_push = artifact_utils.get_single_instance(output_dict["model_push"])
    if not self.CheckBlessing(input_dict):
        model_push.set_int_custom_property("pushed", 0)
        return

    model_export = artifact_utils.get_single_instance(input_dict["model_export"])
    model_export_uri = model_export.uri

    custom_config = exec_properties.get("custom_config", {})
    bigquery_serving_args = custom_config.get("bigquery_serving_args", None)
    # if configuration is missing error out
    if bigquery_serving_args is None:
        raise ValueError("Big Query ML configuration was not provided")

    bq_model_uri = "`{}`.`{}`.`{}`".format(
        bigquery_serving_args["project_id"],
        bigquery_serving_args["bq_dataset_id"],
        bigquery_serving_args["model_name"],
    )

    # Deploy the model.
    model_path = path_utils.serving_model_path(model_export_uri)

    if not model_path.startswith("gs://"):
        raise ValueError("pipeline_root must be gs:// for BigQuery ML Pusher.")

    absl.logging.info(
        "Deploying the model to BigQuery ML for serving: {} from {}".format(
            bigquery_serving_args, model_path
        )
    )

    query = """
      CREATE OR REPLACE MODEL {}
      OPTIONS (model_type='tensorflow',
               model_path='{}')""".format(bq_model_uri, os.path.join(model_path, "*"))

    # TODO(zhitaoli): Refactor the executor_class_path creation into a common
    # utility function.
    executor_class_path = "%s.%s" % (self.__class__.__module__, self.__class__.__name__)
    with telemetry_utils.scoped_labels(
        {telemetry_utils.TFX_EXECUTOR: executor_class_path}
    ):
        default_query_job_config = bigquery.job.QueryJobConfig(
            labels=telemetry_utils.get_labels_dict()
        )
    client = bigquery.Client(default_query_job_config=default_query_job_config)

    try:
        query_job = client.query(query)
        query_job.result()  # Waits for the query to finish
    except Exception as e:
        raise RuntimeError("BigQuery ML Push failed: {}".format(e))

    absl.logging.info(
        "Successfully deployed model {} serving from {}".format(
            bq_model_uri, model_path
        )
    )

    # Setting the push_destination to bigquery uri
    model_push.set_int_custom_property("pushed", 1)
    model_push.set_string_custom_property("pushed_model", bq_model_uri)


# Traceback (KeyError)
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-38-f67dc69712f4> in <module>
----> 1 context.run(pusher)

~/XXX/tfx/tfx/tfx/orchestration/experimental/interactive/interactive_context.py in run_if_ipython(*args, **kwargs)
64       # __IPYTHON__ variable is set by IPython, see
65       # https://ipython.org/ipython-doc/rel-0.10.2/html/interactive/reference.html#embedding-ipython.
---> 66       return fn(*args, **kwargs)
67     else:
68       absl.logging.warning(

~/XXX/tfx/tfx/tfx/orchestration/experimental/interactive/interactive_context.py in run(self, component, enable_cache, beam_pipeline_args)
166         component, pipeline_info, driver_args, metadata_connection,
167         beam_pipeline_args, additional_pipeline_args)
--> 168     execution_id = launcher.launch().execution_id
169
170     return execution_result.ExecutionResult(

~/XXX/tfx/tfx/tfx/orchestration/launcher/base_component_launcher.py in launch(self)
203                          execution_decision.input_dict,
204                          execution_decision.output_dict,
--> 205                          execution_decision.exec_properties)
206
207     absl.logging.info('Running publisher for %s',

~/XXX/tfx/tfx/tfx/orchestration/launcher/in_process_component_launcher.py in _run_executor(self, execution_id, input_dict, output_dict, exec_properties)
65         executor_context)  # type: ignore
66
---> 67     executor.Do(input_dict, output_dict, exec_properties)

~/XXX/tfx/tfx/tfx/extensions/google_cloud_big_query_ml/pusher/executor.py in Do(self, input_dict, output_dict, exec_properties)
59     """
60     self._log_startup(input_dict, output_dict, exec_properties)
---> 61     model_push = artifact_utils.get_single_instance(output_dict['model_push'])
62     if not self.CheckBlessing(input_dict):
63       model_push.set_int_custom_property('pushed', 0)

KeyError: 'model_push'

# Fixed code
def Do(
    self,
    input_dict: Dict[Text, List[types.Artifact]],
    output_dict: Dict[Text, List[types.Artifact]],
    exec_properties: Dict[Text, Any],
):
    """Overrides the tfx_pusher_executor.

    Args:
      input_dict: Input dict from input key to a list of artifacts, including:
        - model_export: exported model from trainer.
        - model_blessing: model blessing path from model_validator.
      output_dict: Output dict from key to a list of artifacts, including:
        - model_push: A list of 'ModelPushPath' artifact of size one. It will
          include the model in this push execution if the model was pushed.
      exec_properties: Mostly a passthrough input dict for
        tfx.components.Pusher.executor.  custom_config.bigquery_serving_args is
        consumed by this class.  For the full set of parameters supported by
        Big Query ML, refer to https://cloud.google.com/bigquery-ml/

    Returns:
      None
    Raises:
      ValueError:
        If bigquery_serving_args is not in exec_properties.custom_config.
        If pipeline_root is not 'gs://...'
      RuntimeError: if the Big Query job failed.
    """
    self._log_startup(input_dict, output_dict, exec_properties)
    model_push = artifact_utils.get_single_instance(
        output_dict[tfx_pusher_executor.PUSHED_MODEL_KEY]
    )
    if not self.CheckBlessing(input_dict):
        self._MarkNotPushed(model_push)
        return

    model_export = artifact_utils.get_single_instance(
        input_dict[tfx_pusher_executor.MODEL_KEY]
    )
    model_export_uri = model_export.uri

    custom_config = exec_properties.get(_CUSTOM_CONFIG_KEY, {})
    bigquery_serving_args = custom_config.get(SERVING_ARGS_KEY)
    # if configuration is missing error out
    if bigquery_serving_args is None:
        raise ValueError("Big Query ML configuration was not provided")

    bq_model_uri = ".".join(
        [
            bigquery_serving_args[_PROJECT_ID_KEY],
            bigquery_serving_args[_BQ_DATASET_ID_KEY],
            bigquery_serving_args[_MODEL_NAME_KEY],
        ]
    )

    # Deploy the model.
    io_utils.copy_dir(
        src=path_utils.serving_model_path(model_export_uri), dst=model_push.uri
    )
    model_path = model_push.uri
    if not model_path.startswith(_GCS_PREFIX):
        raise ValueError("pipeline_root must be gs:// for BigQuery ML Pusher.")

    logging.info(
        "Deploying the model to BigQuery ML for serving: %s from %s",
        bigquery_serving_args,
        model_path,
    )

    query = _BQML_CREATE_OR_REPLACE_MODEL_QUERY_TEMPLATE.format(
        model_uri=bq_model_uri, model_path=model_path
    )

    # TODO(zhitaoli): Refactor the executor_class_path creation into a common
    # utility function.
    executor_class_path = "%s.%s" % (self.__class__.__module__, self.__class__.__name__)
    with telemetry_utils.scoped_labels(
        {telemetry_utils.TFX_EXECUTOR: executor_class_path}
    ):
        default_query_job_config = bigquery.job.QueryJobConfig(
            labels=telemetry_utils.get_labels_dict()
        )
    client = bigquery.Client(default_query_job_config=default_query_job_config)

    try:
        query_job = client.query(query)
        query_job.result()  # Waits for the query to finish
    except Exception as e:
        raise RuntimeError("BigQuery ML Push failed: {}".format(e))

    logging.info(
        "Successfully deployed model %s serving from %s", bq_model_uri, model_path
    )

    # Setting the push_destination to bigquery uri
    self._MarkPushed(model_push, pushed_destination=bq_model_uri)