text: # Copyright 2019 Google LLC. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Definition for Airflow component for TFX."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from __future__ import unicode_literals
  
  import functools
  from airflow import models
  from airflow.operators import python_operator
  
  from typing import Any, Dict, Text
  
  from ml_metadata.proto import metadata_store_pb2
  from tfx.components.base import base_component
  from tfx.orchestration import component_launcher
  from tfx.orchestration import data_types
  
  
  def _airflow_component_launcher(
      component: base_component.BaseComponent,
      pipeline_info: data_types.PipelineInfo, driver_args: data_types.DriverArgs,
      metadata_connection_config: metadata_store_pb2.ConnectionConfig,
      additional_pipeline_args: Dict[Text, Any], **kwargs) -> None:
    """Helper function to launch TFX component execution.
  
    This helper function will be called with Airflow env objects which contains
    run_id that we need to pass into TFX ComponentLauncher.
  
    Args:
      component: TFX BaseComponent instance. This instance holds all inputs and
        outputs placeholders as well as component properties.
      pipeline_info: a data_types.PipelineInfo instance that holds pipeline
        properties
      driver_args: component specific args for driver.
      metadata_connection_config: configuration for how to connect to metadata.
      additional_pipeline_args: a dict of additional pipeline args. Currently
        supporting following keys: beam_pipeline_args.
      **kwargs: Context arguments that will be passed in by Airflow, including:
        - ti: TaskInstance object from which we can get run_id of the running
              pipeline.
        For more details, please refer to the code:
        https://github.com/apache/airflow/blob/master/airflow/operators/python_operator.py
    """
    # Populate run id from Airflow task instance.
    pipeline_info.run_id = kwargs['ti'].get_dagrun().run_id
    launcher = component_launcher.ComponentLauncher(
        component=component,
        pipeline_info=pipeline_info,
        driver_args=driver_args,
        metadata_connection_config=metadata_connection_config,
        additional_pipeline_args=additional_pipeline_args)
    launcher.launch()
  
  
  class AirflowComponent(python_operator.PythonOperator):
    """Airflow-specific TFX Component.
  
    This class wrap a component run into its own PythonOperator in Airflow.
    """
  
    def __init__(self, parent_dag: models.DAG,
                 component: base_component.BaseComponent,
                 pipeline_info: data_types.PipelineInfo, enable_cache: bool,
                 metadata_connection_config: metadata_store_pb2.ConnectionConfig,
                 additional_pipeline_args: Dict[Text, Any]):
      """Constructs an Airflow implementation of TFX component.
  
      Args:
        parent_dag: An AirflowPipeline instance as the pipeline DAG.
        component: An instance of base_component.BaseComponent that holds all
          properties of a logical component.
        pipeline_info: An instance of data_types.PipelineInfo that holds pipeline
          properties.
        enable_cache: Whether or not cache is enabled for this component run.
        metadata_connection_config: A config proto for metadata connection.
        additional_pipeline_args: Additional pipeline args.
      """
      # Prepare parameters to create TFX worker.
      driver_args = data_types.DriverArgs(enable_cache=enable_cache)
  
      super(AirflowComponent, self).__init__(
          task_id=component.component_id,
          provide_context=True,
          python_callable=functools.partial(
              _airflow_component_launcher,
              component=component,
              pipeline_info=pipeline_info,
              driver_args=driver_args,
              metadata_connection_config=metadata_connection_config,
              additional_pipeline_args=additional_pipeline_args),
          dag=parent_dag)
  