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."""
  
  import functools
  from typing import Any, Dict, List, Type
  
  import airflow
  from airflow.operators import python
  from ml_metadata.proto import metadata_store_pb2
  from tfx.dsl.components.base import base_node
  from tfx.orchestration import data_types, metadata
  from tfx.orchestration.config import base_component_config
  from tfx.orchestration.launcher import base_component_launcher
  from tfx.utils import telemetry_utils
  
  
  def _airflow_component_launcher(
      component: base_node.BaseNode,
      component_launcher_class: Type[
          base_component_launcher.BaseComponentLauncher
      ],
      pipeline_info: data_types.PipelineInfo,
      driver_args: data_types.DriverArgs,
      metadata_connection_config: metadata_store_pb2.ConnectionConfig,
      beam_pipeline_args: List[str],
      additional_pipeline_args: Dict[str, Any],
      component_config: base_component_config.BaseComponentConfig,
      exec_properties: Dict[str, Any],
      **kwargs: Any
  ) -> 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.
        component_launcher_class: The class of the launcher to launch the component.
        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.
        beam_pipeline_args: Pipeline arguments for Beam powered Components.
        additional_pipeline_args: A dict of additional pipeline args.
        component_config: Component config to launch the component.
        exec_properties: Execution properties from the ComponentSpec.
        **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
      """
      component.exec_properties.update(exec_properties)
  
      # Populate run id from Airflow task instance.
      pipeline_info.run_id = kwargs["ti"].get_dagrun().run_id
      launcher = component_launcher_class.create(
          component=component,
          pipeline_info=pipeline_info,
          driver_args=driver_args,
          metadata_connection=metadata.Metadata(metadata_connection_config),
          beam_pipeline_args=beam_pipeline_args,
          additional_pipeline_args=additional_pipeline_args,
          component_config=component_config,
      )
      with telemetry_utils.scoped_labels(
          {telemetry_utils.LABEL_TFX_RUNNER: "airflow"}
      ):
          launcher.launch()
  
  
  class AirflowComponent(python.PythonOperator):
      """Airflow-specific TFX Component.
      This class wrap a component run into its own PythonOperator in Airflow.
      """
  
      def __init__(
          self,
          *,
          parent_dag: airflow.DAG,
          component: base_node.BaseNode,
          component_launcher_class: Type[
              base_component_launcher.BaseComponentLauncher
          ],
          pipeline_info: data_types.PipelineInfo,
          enable_cache: bool,
          metadata_connection_config: metadata_store_pb2.ConnectionConfig,
          beam_pipeline_args: List[str],
          additional_pipeline_args: Dict[str, Any],
          component_config: base_component_config.BaseComponentConfig
      ):
          """Constructs an Airflow implementation of TFX component.
  
          Args:
            parent_dag: An AirflowPipeline instance as the pipeline DAG.
            component: An instance of base_node.BaseNode that holds all
              properties of a logical component.
            component_launcher_class: The class of the launcher to launch the
              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.
            beam_pipeline_args: Pipeline arguments for Beam powered Components.
            additional_pipeline_args: Additional pipeline args.
            component_config: Component config to launch the component.
          """
          # Prepare parameters to create TFX worker.
          driver_args = data_types.DriverArgs(enable_cache=enable_cache)
  
          exec_properties = component.exec_properties
  
          super().__init__(
              task_id=component.id,
              provide_context=True,
              python_callable=functools.partial(
                  _airflow_component_launcher,
                  component=component,
                  component_launcher_class=component_launcher_class,
                  pipeline_info=pipeline_info,
                  driver_args=driver_args,
                  metadata_connection_config=metadata_connection_config,
                  beam_pipeline_args=beam_pipeline_args,
                  additional_pipeline_args=additional_pipeline_args,
                  component_config=component_config,
              ),
              # op_kwargs is a templated field for PythonOperator, which means Airflow
              # will inspect the dictionary and resolve any templated fields.
              op_kwargs={"exec_properties": exec_properties},
              dag=parent_dag,
          )
  