diff --git a/catalog/dags/common/ingestion_server.py b/catalog/dags/common/ingestion_server.py index 54e0c3622e6..43a628fb12e 100644 --- a/catalog/dags/common/ingestion_server.py +++ b/catalog/dags/common/ingestion_server.py @@ -4,7 +4,6 @@ from urllib.parse import urlparse from airflow.exceptions import AirflowSkipException -from airflow.models import Variable from airflow.providers.http.operators.http import SimpleHttpOperator from airflow.providers.http.sensors.http import HttpSensor from requests import Response @@ -160,13 +159,12 @@ def api_health_check( retry. The task is set to retry with exponential backoff, such that the retry delay doubles between each attempt. """ - access_token = Variable.get("API_ACCESS_TOKEN", "not_set") return HttpSensor( task_id="api_health_check", http_conn_id=API_CONN_ID, endpoint=f"{media_type}", request_params={"internal__index": f"{media_type}-{index_suffix}"}, - headers={"Authorization": f"Bearer {access_token}"}, + headers={"Authorization": "Bearer {{ var.value.API_ACCESS_TOKEN }}"}, method="GET", response_check=response_check_api_health_check, mode="reschedule", diff --git a/catalog/tests/dags/common/test_ingestion_server.py b/catalog/tests/dags/common/test_ingestion_server.py index b75b6526962..aca42ce4c9a 100644 --- a/catalog/tests/dags/common/test_ingestion_server.py +++ b/catalog/tests/dags/common/test_ingestion_server.py @@ -33,12 +33,10 @@ def clean_db(): @pytest.fixture() def healthcheck_dag(): - index_suffix = "my-test-suffix" - # Create a DAG that just has an api_health_check task with DAG(dag_id=TEST_DAG_ID, schedule=None, start_date=TEST_START_DATE) as dag: ingestion_server.api_health_check( - media_type="image", index_suffix=index_suffix, timeout=timedelta(days=1) + media_type="image", index_suffix="my_test_suffix", timeout=timedelta(days=1) ) return dag