From 102fe934f6eef92baddf6b1d1dca609d3608bf45 Mon Sep 17 00:00:00 2001 From: Geoffroy Dimur Date: Tue, 29 Aug 2023 14:36:59 +0200 Subject: [PATCH 1/3] Fix : Add "QUEUED" to RUN_LIFE_CYCLE_STATES following deployement of queue feature on databricks jobs --- airflow/providers/databricks/hooks/databricks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index 2a7a8b91ff6bd..82e79d223935d 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -55,7 +55,7 @@ WORKSPACE_GET_STATUS_ENDPOINT = ("GET", "api/2.0/workspace/get-status") -RUN_LIFE_CYCLE_STATES = ["PENDING", "RUNNING", "TERMINATING", "TERMINATED", "SKIPPED", "INTERNAL_ERROR"] +RUN_LIFE_CYCLE_STATES = ["PENDING", "RUNNING", "TERMINATING", "TERMINATED", "SKIPPED", "INTERNAL_ERROR", "QUEUED"] SPARK_VERSIONS_ENDPOINT = ("GET", "api/2.0/clusters/spark-versions") From e9941078f6993b3c8784883138cd205bc8709333 Mon Sep 17 00:00:00 2001 From: Geoffroy Dimur Date: Tue, 29 Aug 2023 15:18:15 +0200 Subject: [PATCH 2/3] Lint --- airflow/providers/databricks/hooks/databricks.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py index 82e79d223935d..e8b124b99fb18 100644 --- a/airflow/providers/databricks/hooks/databricks.py +++ b/airflow/providers/databricks/hooks/databricks.py @@ -55,7 +55,15 @@ WORKSPACE_GET_STATUS_ENDPOINT = ("GET", "api/2.0/workspace/get-status") -RUN_LIFE_CYCLE_STATES = ["PENDING", "RUNNING", "TERMINATING", "TERMINATED", "SKIPPED", "INTERNAL_ERROR", "QUEUED"] +RUN_LIFE_CYCLE_STATES = [ + "PENDING", + "RUNNING", + "TERMINATING", + "TERMINATED", + "SKIPPED", + "INTERNAL_ERROR", + "QUEUED", +] SPARK_VERSIONS_ENDPOINT = ("GET", "api/2.0/clusters/spark-versions") From c53d456337e7e11836f5be7eaabb76f8a7a811f0 Mon Sep 17 00:00:00 2001 From: Geoffroy Dimur Date: Tue, 29 Aug 2023 16:16:58 +0200 Subject: [PATCH 3/3] Extend databricks tests --- tests/providers/databricks/hooks/test_databricks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py index 8644d95cd9c4c..52f224a04f5e2 100644 --- a/tests/providers/databricks/hooks/test_databricks.py +++ b/tests/providers/databricks/hooks/test_databricks.py @@ -940,7 +940,7 @@ def test_is_terminal_true(self): assert run_state.is_terminal def test_is_terminal_false(self): - non_terminal_states = ["PENDING", "RUNNING", "TERMINATING"] + non_terminal_states = ["PENDING", "RUNNING", "TERMINATING", "QUEUED"] for state in non_terminal_states: run_state = RunState(state, "", "") assert not run_state.is_terminal