Httphook: rertry on AirflowException

dags/library/cern_publication_records.py

@@ -6,7 +6,8 @@
 from airflow.providers.postgres.operators.postgres import PostgresOperator
 from common.utils import get_total_results_count
 from executor_config import kubernetes_executor_config
-from library.utils import get_url
+from library.utils import get_endpoint
+from requests.exceptions import HTTPError
 from tenacity import retry_if_exception_type, stop_after_attempt
 
 
@@ -19,7 +20,7 @@ def library_cern_publication_records_dag():
     @task(multiple_outputs=True, executor_config=kubernetes_executor_config)
     def generate_params(key, **kwargs):
         year = kwargs["params"].get("year")
-        url = get_url(key, year)
+        url = get_endpoint(key, year)
         return {
             "endpoint": url,
             "type_of_query": key,
@@ -32,7 +33,7 @@ def fetch_count(parameters):
             endpoint=parameters["endpoint"],
             _retry_args={
                 "stop": stop_after_attempt(3),
-                "retry": retry_if_exception_type(Exception),
+                "retry": retry_if_exception_type(HTTPError),
             },
         )
         count = get_total_results_count(response.text)

dags/library/utils.py

@@ -1,4 +1,4 @@
-def get_url(key, year):
+def get_endpoint(key, year):
     url = {
         "publications_total_count": r"search?wl=0&ln=en&cc=Published+"
         + r"Articles&sc=1&p=%28affiliation%3ACERN+or+595%3A%27For+annual+report%27%29+and+"

dags/open_access/gold_open_access_mechanisms.py

@@ -7,6 +7,7 @@
 from airflow.providers.postgres.operators.postgres import PostgresOperator
 from common.utils import get_total_results_count
 from executor_config import kubernetes_executor_config
+from requests.exceptions import ProxyError
 from tenacity import retry_if_exception_type, stop_after_attempt
 
 
@@ -42,7 +43,7 @@ def fetch_count(parameters):
             endpoint=parameters["endpoint"],
             _retry_args={
                 "stop": stop_after_attempt(3),
-                "retry": retry_if_exception_type(Exception),
+                "retry": retry_if_exception_type(ProxyError),
             },
         )
         count = get_total_results_count(response.text)

dags/open_access/open_access.py

@@ -4,6 +4,7 @@
 import open_access.utils as utils
 import pendulum
 from airflow.decorators import dag, task
+from airflow.exceptions import AirflowException
 from airflow.providers.http.hooks.http import HttpHook
 from airflow.providers.postgres.operators.postgres import PostgresOperator
 from common.utils import get_total_results_count
@@ -43,7 +44,7 @@ def fetch_count(parameters):
             endpoint=parameters["endpoint"],
             _retry_args={
                 "stop": stop_after_attempt(3),
-                "retry": retry_if_exception_type(Exception),
+                "retry": retry_if_exception_type(AirflowException),
             },
         )
         type_of_query = parameters["type_of_query"]

dags/open_access/utils.py

@@ -6,6 +6,7 @@
     get_golden_access_records_ids,
     get_green_access_records_ids,
 )
+from requests.exceptions import ProxyError
 from tenacity import retry_if_exception_type, stop_after_attempt
 
 
@@ -20,7 +21,7 @@ def get_count_http_hook(total, url, record_extractor):
             endpoint=full_url,
             _retry_args={
                 "stop": stop_after_attempt(3),
-                "retry": retry_if_exception_type(Exception),
+                "retry": retry_if_exception_type(ProxyError),
             },
         )
         records_ids_count = records_ids_count + len(record_extractor(response.text))