From b76aa54c42ab3d7c3f3707da0215056c817361bc Mon Sep 17 00:00:00 2001
From: utkarsh sharma <utkarsharma2@gmail.com>
Date: Thu, 21 Sep 2023 15:59:42 +0530
Subject: [PATCH 1/3] Respect 'soft_fail' argument when running
 DatabricksPartitionSensor

---
 .../sensors/databricks_partition.py           | 25 +++++--
 .../sensors/test_databricks_partition.py      | 66 ++++++++++++++++++-
 2 files changed, 83 insertions(+), 8 deletions(-)

diff --git a/airflow/providers/databricks/sensors/databricks_partition.py b/airflow/providers/databricks/sensors/databricks_partition.py
index 98ad0f461a9b6..7af2f78aa126f 100644
--- a/airflow/providers/databricks/sensors/databricks_partition.py
+++ b/airflow/providers/databricks/sensors/databricks_partition.py
@@ -26,7 +26,7 @@
 
 from databricks.sql.utils import ParamEscaper
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.common.sql.hooks.sql import fetch_all_handler
 from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook
 from airflow.sensors.base import BaseSensorOperator
@@ -182,7 +182,11 @@ def _generate_partition_query(
         partition_columns = self._sql_sensor(f"DESCRIBE DETAIL {table_name}")[0][7]
         self.log.debug("Partition columns: %s", partition_columns)
         if len(partition_columns) < 1:
-            raise AirflowException(f"Table {table_name} does not have partitions")
+            message = f"Table {table_name} does not have partitions"
+            if self.soft_fail:
+                raise AirflowSkipException(message)
+            raise AirflowException(message)
+
         formatted_opts = ""
         if opts:
             output_list = []
@@ -202,12 +206,16 @@ def _generate_partition_query(
                             f"""{partition_col}{self.partition_operator}{self.escaper.escape_item(partition_value)}"""
                         )
                 else:
-                    raise AirflowException(
-                        f"Column {partition_col} not part of table partitions: {partition_columns}"
-                    )
+                    message = f"Column {partition_col} not part of table partitions: {partition_columns}"
+                    if self.soft_fail:
+                        raise AirflowSkipException(message)
+                    raise AirflowException(message)
         else:
             # Raises exception if the table does not have any partitions.
-            raise AirflowException("No partitions specified to check with the sensor.")
+            message = "No partitions specified to check with the sensor."
+            if self.soft_fail:
+                raise AirflowSkipException(message)
+            raise AirflowException(message)
         formatted_opts = f"{prefix} {joiner_val.join(output_list)} {suffix}"
         self.log.debug("Formatted options: %s", formatted_opts)
 
@@ -220,4 +228,7 @@ def poke(self, context: Context) -> bool:
         if partition_result:
             return True
         else:
-            raise AirflowException(f"Specified partition(s): {self.partitions} were not found.")
+            message = f"Specified partition(s): {self.partitions} were not found."
+            if self.soft_fail:
+                raise AirflowSkipException(message)
+            raise AirflowException(message)
diff --git a/tests/providers/databricks/sensors/test_databricks_partition.py b/tests/providers/databricks/sensors/test_databricks_partition.py
index 56e77e5e59a38..5d0cce5e59fd6 100644
--- a/tests/providers/databricks/sensors/test_databricks_partition.py
+++ b/tests/providers/databricks/sensors/test_databricks_partition.py
@@ -23,7 +23,7 @@
 
 import pytest
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.models import DAG
 from airflow.providers.common.sql.hooks.sql import fetch_all_handler
 from airflow.providers.databricks.sensors.databricks_partition import DatabricksPartitionSensor
@@ -99,3 +99,67 @@ def test_unsupported_conn_type(self):
     def test_partition_sensor(self, patched_poke):
         patched_poke.return_value = True
         assert self.partition_sensor.poke({})
+
+    @pytest.mark.parametrize(
+        "soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
+    )
+    def test_fail__generate_partition_query(self, soft_fail, expected_exception):
+        self.partition_sensor.soft_fail = soft_fail
+        table_name = "test"
+        with pytest.raises(expected_exception) as e, patch(
+            "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
+            "._sql_sensor"
+        ) as _sql_sensor:
+            _sql_sensor.return_value = [[[], [], [], [], [], [], [], []]]
+            self.partition_sensor._generate_partition_query(
+                prefix="", suffix="", joiner_val="", table_name=table_name
+            )
+        assert e.match(f"Table {table_name} does not have partitions")
+
+    @pytest.mark.parametrize(
+        "soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
+    )
+    def test_fail__generate_partition_query_with_partition_col_mismatch(self, soft_fail, expected_exception):
+        self.partition_sensor.soft_fail = soft_fail
+        partition_col = "non_existent_col"
+        partition_columns = ["col1", "col2"]
+        with pytest.raises(expected_exception) as e, patch(
+            "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
+            "._sql_sensor"
+        ) as _sql_sensor:
+            _sql_sensor.return_value = [[[], [], [], [], [], [], [], partition_columns]]
+            self.partition_sensor._generate_partition_query(
+                prefix="", suffix="", joiner_val="", table_name="", opts={partition_col: "1"}
+            )
+        assert e.match(f"Column {partition_col} not part of table partitions")
+
+    @pytest.mark.parametrize(
+        "soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
+    )
+    def test_fail__generate_partition_query_with_missing_opts(self, soft_fail, expected_exception):
+        self.partition_sensor.soft_fail = soft_fail
+        with pytest.raises(expected_exception) as e, patch(
+            "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
+            "._sql_sensor"
+        ) as _sql_sensor:
+            _sql_sensor.return_value = [[[], [], [], [], [], [], [], ["col1", "col2"]]]
+            self.partition_sensor._generate_partition_query(
+                prefix="", suffix="", joiner_val="", table_name=""
+            )
+        assert e.match("No partitions specified to check with the sensor.")
+
+    @pytest.mark.parametrize(
+        "soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
+    )
+    def test_fail_poke(self, soft_fail, expected_exception):
+        self.partition_sensor.soft_fail = soft_fail
+        partitions = "test"
+        self.partition_sensor.partitions = partitions
+        with pytest.raises(expected_exception) as e, patch(
+            "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
+            "._check_table_partitions"
+        ) as _check_table_partitions:
+            _check_table_partitions.return_value = False
+            self.partition_sensor.poke(context={})
+
+        assert e.match(f"Specified partition\(s\): {partitions} were not found.")

From 265832271a4215d8fb7781ee7585662246c3e006 Mon Sep 17 00:00:00 2001
From: utkarsh sharma <utkarsharma2@gmail.com>
Date: Thu, 21 Sep 2023 16:17:35 +0530
Subject: [PATCH 2/3] Add todos

---
 airflow/providers/databricks/sensors/databricks_partition.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/airflow/providers/databricks/sensors/databricks_partition.py b/airflow/providers/databricks/sensors/databricks_partition.py
index 7af2f78aa126f..1aea220078673 100644
--- a/airflow/providers/databricks/sensors/databricks_partition.py
+++ b/airflow/providers/databricks/sensors/databricks_partition.py
@@ -182,6 +182,7 @@ def _generate_partition_query(
         partition_columns = self._sql_sensor(f"DESCRIBE DETAIL {table_name}")[0][7]
         self.log.debug("Partition columns: %s", partition_columns)
         if len(partition_columns) < 1:
+            # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
             message = f"Table {table_name} does not have partitions"
             if self.soft_fail:
                 raise AirflowSkipException(message)
@@ -206,12 +207,14 @@ def _generate_partition_query(
                             f"""{partition_col}{self.partition_operator}{self.escaper.escape_item(partition_value)}"""
                         )
                 else:
+                    # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
                     message = f"Column {partition_col} not part of table partitions: {partition_columns}"
                     if self.soft_fail:
                         raise AirflowSkipException(message)
                     raise AirflowException(message)
         else:
             # Raises exception if the table does not have any partitions.
+            # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
             message = "No partitions specified to check with the sensor."
             if self.soft_fail:
                 raise AirflowSkipException(message)
@@ -228,6 +231,7 @@ def poke(self, context: Context) -> bool:
         if partition_result:
             return True
         else:
+            # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
             message = f"Specified partition(s): {self.partitions} were not found."
             if self.soft_fail:
                 raise AirflowSkipException(message)

From fa421a617ff2d62ea2f5ccd6bd4559ea98bc208a Mon Sep 17 00:00:00 2001
From: utkarsh sharma <utkarsharma2@gmail.com>
Date: Thu, 21 Sep 2023 16:22:48 +0530
Subject: [PATCH 3/3] Move assert e.match as part of pytest.raises

---
 .../sensors/test_databricks_partition.py      | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/tests/providers/databricks/sensors/test_databricks_partition.py b/tests/providers/databricks/sensors/test_databricks_partition.py
index 5d0cce5e59fd6..839d36d4e9e67 100644
--- a/tests/providers/databricks/sensors/test_databricks_partition.py
+++ b/tests/providers/databricks/sensors/test_databricks_partition.py
@@ -106,7 +106,7 @@ def test_partition_sensor(self, patched_poke):
     def test_fail__generate_partition_query(self, soft_fail, expected_exception):
         self.partition_sensor.soft_fail = soft_fail
         table_name = "test"
-        with pytest.raises(expected_exception) as e, patch(
+        with pytest.raises(expected_exception, match=f"Table {table_name} does not have partitions"), patch(
             "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
             "._sql_sensor"
         ) as _sql_sensor:
@@ -114,7 +114,6 @@ def test_fail__generate_partition_query(self, soft_fail, expected_exception):
             self.partition_sensor._generate_partition_query(
                 prefix="", suffix="", joiner_val="", table_name=table_name
             )
-        assert e.match(f"Table {table_name} does not have partitions")
 
     @pytest.mark.parametrize(
         "soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
@@ -123,7 +122,9 @@ def test_fail__generate_partition_query_with_partition_col_mismatch(self, soft_f
         self.partition_sensor.soft_fail = soft_fail
         partition_col = "non_existent_col"
         partition_columns = ["col1", "col2"]
-        with pytest.raises(expected_exception) as e, patch(
+        with pytest.raises(
+            expected_exception, match=f"Column {partition_col} not part of table partitions"
+        ), patch(
             "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
             "._sql_sensor"
         ) as _sql_sensor:
@@ -131,14 +132,15 @@ def test_fail__generate_partition_query_with_partition_col_mismatch(self, soft_f
             self.partition_sensor._generate_partition_query(
                 prefix="", suffix="", joiner_val="", table_name="", opts={partition_col: "1"}
             )
-        assert e.match(f"Column {partition_col} not part of table partitions")
 
     @pytest.mark.parametrize(
         "soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
     )
     def test_fail__generate_partition_query_with_missing_opts(self, soft_fail, expected_exception):
         self.partition_sensor.soft_fail = soft_fail
-        with pytest.raises(expected_exception) as e, patch(
+        with pytest.raises(
+            expected_exception, match="No partitions specified to check with the sensor."
+        ), patch(
             "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
             "._sql_sensor"
         ) as _sql_sensor:
@@ -146,7 +148,6 @@ def test_fail__generate_partition_query_with_missing_opts(self, soft_fail, expec
             self.partition_sensor._generate_partition_query(
                 prefix="", suffix="", joiner_val="", table_name=""
             )
-        assert e.match("No partitions specified to check with the sensor.")
 
     @pytest.mark.parametrize(
         "soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
@@ -155,11 +156,11 @@ def test_fail_poke(self, soft_fail, expected_exception):
         self.partition_sensor.soft_fail = soft_fail
         partitions = "test"
         self.partition_sensor.partitions = partitions
-        with pytest.raises(expected_exception) as e, patch(
+        with pytest.raises(
+            expected_exception, match=f"Specified partition\(s\): {partitions} were not found."
+        ), patch(
             "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
             "._check_table_partitions"
         ) as _check_table_partitions:
             _check_table_partitions.return_value = False
             self.partition_sensor.poke(context={})
-
-        assert e.match(f"Specified partition\(s\): {partitions} were not found.")
