Add Operator to Retrieve Task State

shipyard_airflow/dags/samples/airflow_task_state_operators.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+#
+# 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.
+"""
+### Airflow Task State
+"""
+import airflow
+from airflow import DAG
+from airflow.operators import TaskStateOperator
+from airflow.operators.bash_operator import BashOperator
+from datetime import datetime, timedelta
+
+default_args = {
+    'owner': 'airflow',
+    'depends_on_past': False,
+    'start_date': airflow.utils.dates.days_ago(2),
+    'email': ['airflow@example.com'],
+    'email_on_failure': False,
+    'email_on_retry': False,
+    'retries': 1,
+    'retry_delay': timedelta(minutes=1),
+}
+
+dag = DAG('airflow_task_state', default_args=default_args, schedule_interval=None)
+
+# Get Task State
+t1 = TaskStateOperator(
+    task_id='airflow_task_state',
+    airflow_dag_id='airflow_cli',
+    airflow_task_id='airflow_task_state',
+    airflow_execution_date='2017-06-25T21:27:52.809436',
+    dag=dag)
+
+# Use XCOM to Retrieve Task State
+t2 = BashOperator(
+    task_id='pull',
+    bash_command="echo {{ ti.xcom_pull(task_ids='airflow_task_state', key='task_state') }}",
+    xcom_push=True,
+    dag=dag)
+
+t2.set_upstream(t1)
+

shipyard_airflow/plugins/airflow_task_state_operators.py

@@ -0,0 +1,81 @@
+# -*- coding: utf-8 -*-
+#
+# 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.
+
+import logging
+import subprocess
+import sys
+import os
+import shlex
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.plugins_manager import AirflowPlugin
+from airflow.utils.decorators import apply_defaults
+
+class TaskStateOperator(BaseOperator):
+    """
+    Retrieve Task State
+    :airflow_dag_id: Dag ID
+    :airflow_task_id: Task ID
+    :airflow_execution_date: Task Execution Date
+    """
+    @apply_defaults
+    def __init__(self,
+                 airflow_command=None,
+                 airflow_dag_id=None,
+                 airflow_task_id=None,
+                 airflow_execution_date=None,
+                 *args, **kwargs):
+
+        super(TaskStateOperator, self).__init__(*args, **kwargs)
+        self.airflow_dag_id = airflow_dag_id
+        self.airflow_task_id = airflow_task_id
+        self.airflow_execution_date = airflow_execution_date
+        self.airflow_command = "%s %s %s %s" % ("airflow task_state", airflow_dag_id, airflow_task_id, airflow_execution_date)
+
+    def execute(self, context):
+
+        logging.info("Running Airflow Command: " + self.airflow_command)
+
+        # Execute Airflow CLI Command
+        airflow_cli = subprocess.Popen(shlex.split(self.airflow_command), stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+
+        # Logs Output
+        # Filter out logging messages from standard output and keep only the relevant information
+        line = ''
+        for line in iter(airflow_cli.stdout.readline, b''):
+            line = line.strip()
+
+            if line.startswith( '[' ):
+                pass
+            else:
+                logging.info(line)
+                task_state = line
+
+        # Wait for child process to terminate. Set and return returncode attribute.
+        airflow_cli.wait()
+
+        # Raise Execptions if Task State Command Fails
+        if airflow_cli.returncode:
+            raise AirflowException("Failed to Retrieve Task State")
+
+        # Return XCOM State
+        task_instance = context['task_instance']
+        task_instance.xcom_push('task_state', task_state)
+
+
+class TaskStatePlugin(AirflowPlugin):
+    name = "task_state_plugin"
+    operators = [TaskStateOperator]
+