[airflow] branch main updated: Make presto and trino compatible with airflow 2.1 (#23061)

[potiuk]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 5164cdbe98 Make presto and trino compatible with airflow 2.1 (#23061)
5164cdbe98 is described below

commit 5164cdbe98ad63754d969b4b300a7a0167565e33
Author: Ping Zhang <pin...@umich.edu>
AuthorDate: Mon Apr 18 22:29:49 2022 -0700

    Make presto and trino compatible with airflow 2.1 (#23061)
---
 airflow/providers/presto/hooks/presto.py | 5 ++++-
 airflow/providers/trino/hooks/trino.py   | 5 ++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/presto/hooks/presto.py 
b/airflow/providers/presto/hooks/presto.py
index 97f6ec566e..95ecf86b52 100644
--- a/airflow/providers/presto/hooks/presto.py
+++ b/airflow/providers/presto/hooks/presto.py
@@ -47,11 +47,14 @@ def generate_presto_client_info() -> str:
         )
         for format_map in AIRFLOW_VAR_NAME_FORMAT_MAPPING.values()
     }
+    # try_number isn't available in context for airflow < 2.2.5
+    # https://github.com/apache/airflow/issues/23059
+    try_number = context_var.get('try_number', '')
     task_info = {
         'dag_id': context_var['dag_id'],
         'task_id': context_var['task_id'],
         'execution_date': context_var['execution_date'],
-        'try_number': context_var['try_number'],
+        'try_number': try_number,
         'dag_run_id': context_var['dag_run_id'],
         'dag_owner': context_var['dag_owner'],
     }
diff --git a/airflow/providers/trino/hooks/trino.py 
b/airflow/providers/trino/hooks/trino.py
index ee97eeac1a..5e48ea2338 100644
--- a/airflow/providers/trino/hooks/trino.py
+++ b/airflow/providers/trino/hooks/trino.py
@@ -47,11 +47,14 @@ def generate_trino_client_info() -> str:
         )
         for format_map in AIRFLOW_VAR_NAME_FORMAT_MAPPING.values()
     }
+    # try_number isn't available in context for airflow < 2.2.5
+    # https://github.com/apache/airflow/issues/23059
+    try_number = context_var.get('try_number', '')
     task_info = {
         'dag_id': context_var['dag_id'],
         'task_id': context_var['task_id'],
         'execution_date': context_var['execution_date'],
-        'try_number': context_var['try_number'],
+        'try_number': try_number,
         'dag_run_id': context_var['dag_run_id'],
         'dag_owner': context_var['dag_owner'],
     }