How to access dag_run configs inside a kubernetes pod operator

Viewed 1157

I'm riding the struggle bus pretty hardcore right now with this on. Trying to access the dag_run.conf['myValue'] via the KubernetesPodOperator in airflow 1.10.9 (GCP Composer). I'm able to access it in a pythonOperator and write to xcom. In addition we can read from XCom in the podOperator but the dag_run configs are not working, any help would appreciated! See below for DAG

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow_common.airflow_decorators import django_decorator

from airflow_common.operator_defaults import KubernetesOperator, PythonOperator

DAG_ID = "example_trigger_target_dag"
dag = DAG(
    dag_id=DAG_ID,
    default_args={"owner": "airflow"},
    start_date=days_ago(2),
    schedule_interval=None,
    tags=['example'],
    max_active_runs=145,
)

def run_this_func(**kwargs):
    org= kwargs["dag_run"].conf["org"]

    # write to xcom locally from kwargs
    kwargs['ti'].xcom_push(key='org', value=org)

run_this = PythonOperator(
    task_id="run_this",
    python_callable=run_this_func,
    provide_context=True,
    dag=dag
    )

commands = ['python', '-c']
args = [" org = {{ dag_run.conf['org'] }}"
        , "print('org dag_run is: ', org)"
        , "import time"
        , "time.sleep(15)"
]

# wrapper over the KuberentesPodOperator to set company defaults
k = KubernetesOperator(
    dag,
    name="test-kubernetes",
    cmds=commands,
    arguments=args,
    xcom_push=True,
    }
).get_operator()


@django_decorator
def xcom_read(**kwargs):
    org = kwargs['ti'].xcom_pull(task_ids='run_this', key='org')
    print("got org_id: ", org)


read_xcom = PythonOperator(
    task_id='test-read-xcom',
    python_callable=xcom_read,
    op_kwargs={},
    provide_context=True,
    dag=dag)

run_this >> k >> read_xcom

The above code should work, and not all is required for what im testing. In args, i'm using the {{ }} template to get the dag_run value but its not actually being set to be printed to the backend.

Update:

class KubernetesOperator:
    def __init__(self, dag, arguments, name, image_ver="stable",
     cmds=None, env_vars=None, xcom_push=False):
        
        if cmds is None:
            cmds = ["python", "-c"]

        self.dag = dag
        self.arguments = arguments
        self.name = name
        self.image_ver = image_ver
        self.cmds = cmds
        self.env_vars = env_vars
        self.xcom_push = xcom_push

    def get_operator(self):
        
        # apply variables and return
        return KubernetesPodOperator(
            namespace='default',
            image="<image goes here>",
            cmds=self.cmds,
            arguments=self.arguments,
            startup_timeout_seconds=300,
            labels={"airflow": "{ntl}".format(ntl=self.name)},
            task_id=self.name,
            name=self.name,
            get_logs=True,
            image_pull_policy='Always',
            dag=self.dag,
            env_vars=self.env_vars,
            xcom_push=self.xcom_push)
    
0 Answers
Related