import airflow.configuration as config
from airflow.contrib.operators.kubernetes_pod_operator import KubernetesPodOperator as AirflowKubernetesPodOperator
class KubernetesPodOperator(AirflowKubernetesPodOperator):
def execute(self, context):
environment = config.conf.get('webserver', 'web_server_name')
ti = context['ti']
dag_id = ti.dag_id
task_id = ti.task_id
run_id = context['run_id']
try_number = str(ti._try_number)
labels = {
'ENVIRONMENT' : environment,
'DAG_ID' : dag_id,
'TASK_ID' : task_id,
'RUN_ID' : run_id,
'TRY_NUMBER' : try_number,
}
self.env_vars.update(labels)
super().execute(context)