From 298d19e55e781377f4135876af3387e5be573b13 Mon Sep 17 00:00:00 2001 From: Jordi Escrich Date: Tue, 16 Jan 2024 13:39:30 +0100 Subject: [PATCH 1/4] Accept 202 status code response --- airflow/providers/apache/druid/hooks/druid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/apache/druid/hooks/druid.py b/airflow/providers/apache/druid/hooks/druid.py index 11c415e42ca60..aa0b99415b426 100644 --- a/airflow/providers/apache/druid/hooks/druid.py +++ b/airflow/providers/apache/druid/hooks/druid.py @@ -107,7 +107,7 @@ def submit_indexing_job( req_index = requests.post(url, data=json_index_spec, headers=self.header, auth=self.get_auth()) code = req_index.status_code - if code != 200: + if code not in [200, 202]: self.log.error("Error submitting the Druid job to %s (%s) %s", url, code, req_index.content) raise AirflowException(f"Did not get 200 when submitting the Druid job to {url}") From f460b44a6978e2b9dfd55e9a1fc2aebdf5a3c0d6 Mon Sep 17 00:00:00 2001 From: Jordi Escrich Date: Tue, 16 Jan 2024 13:52:03 +0100 Subject: [PATCH 2/4] Fix status message --- airflow/providers/apache/druid/hooks/druid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/apache/druid/hooks/druid.py b/airflow/providers/apache/druid/hooks/druid.py index aa0b99415b426..99c12055e3c73 100644 --- a/airflow/providers/apache/druid/hooks/druid.py +++ b/airflow/providers/apache/druid/hooks/druid.py @@ -109,7 +109,7 @@ def submit_indexing_job( code = req_index.status_code if code not in [200, 202]: self.log.error("Error submitting the Druid job to %s (%s) %s", url, code, req_index.content) - raise AirflowException(f"Did not get 200 when submitting the Druid job to {url}") + raise AirflowException(f"Did not get 200 or 202 when submitting the Druid job to {url}") req_json = req_index.json() # Wait until the job is completed From 64b5075e1ac450761c59fefb59b13bce21fb860b Mon Sep 17 00:00:00 2001 From: Jordi Escrich Date: Wed, 17 Jan 2024 15:14:35 +0100 Subject: [PATCH 3/4] Addresses [comment](https://github.com/apache/airflow/pull/36813#discussion_r1454463089) --- airflow/providers/apache/druid/hooks/druid.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/providers/apache/druid/hooks/druid.py b/airflow/providers/apache/druid/hooks/druid.py index 99c12055e3c73..85bb1167fc91c 100644 --- a/airflow/providers/apache/druid/hooks/druid.py +++ b/airflow/providers/apache/druid/hooks/druid.py @@ -107,7 +107,8 @@ def submit_indexing_job( req_index = requests.post(url, data=json_index_spec, headers=self.header, auth=self.get_auth()) code = req_index.status_code - if code not in [200, 202]: + not_accepted = not (200 <= code < 300) + if not_accepted: self.log.error("Error submitting the Druid job to %s (%s) %s", url, code, req_index.content) raise AirflowException(f"Did not get 200 or 202 when submitting the Druid job to {url}") From a203c24b7b0fe3de9df60a04dba2b5f48f76f896 Mon Sep 17 00:00:00 2001 From: Jordi Escrich Date: Mon, 22 Jan 2024 13:18:16 +0100 Subject: [PATCH 4/4] Make tests cover actual behaviour --- tests/providers/apache/druid/hooks/test_druid.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/providers/apache/druid/hooks/test_druid.py b/tests/providers/apache/druid/hooks/test_druid.py index 221371eac6cec..5a389cb710c65 100644 --- a/tests/providers/apache/druid/hooks/test_druid.py +++ b/tests/providers/apache/druid/hooks/test_druid.py @@ -65,6 +65,7 @@ def test_submit_gone_wrong(self, requests_mock): def test_submit_ok(self, requests_mock): task_post = requests_mock.post( "http://druid-overlord:8081/druid/indexer/v1/task", + status_code=200, text='{"task":"9f8a7359-77d4-4612-b0cd-cc2f6a3c28de"}', ) status_check = requests_mock.get( @@ -81,6 +82,7 @@ def test_submit_ok(self, requests_mock): def test_submit_sql_based_ingestion_ok(self, requests_mock): task_post = requests_mock.post( "http://druid-overlord:8081/druid/v2/sql/task", + status_code=202, text='{"taskId":"9f8a7359-77d4-4612-b0cd-cc2f6a3c28de"}', ) status_check = requests_mock.get(