Common sql bugfixes and improvements

airflow/providers/common/sql/operators/sql.py

@@ -23,7 +23,7 @@
 from packaging.version import Version
 
 from airflow.compat.functools import cached_property
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowFailException
 from airflow.hooks.base import BaseHook
 from airflow.models import BaseOperator, SkipMixin
 from airflow.providers.common.sql.hooks.sql import DbApiHook, _backported_get_hook
@@ -33,7 +33,21 @@
     from airflow.utils.context import Context
 
 
-def parse_boolean(val: str) -> str | bool:
+def _convert_to_float_if_possible(s):
+    """
+    A small helper function to convert a string to a numeric value
+    if appropriate
+
+    :param s: the string to be converted
+    """
+    try:
+        ret = float(s)
+    except (ValueError, TypeError):
+        ret = s
+    return ret
+
+
+def _parse_boolean(val: str) -> str | bool:
     """Try to parse a string into boolean.
 
     Raises ValueError if the input is not a valid true- or false-like string value.
@@ -60,6 +74,12 @@ def _get_failed_checks(checks, col=None):
     ]
 
 
+def _raise_exception(exception_string, retry_on_failure):
+    if retry_on_failure:
+        raise AirflowException(exception_string)
+    raise AirflowFailException(exception_string)
+
+
 _PROVIDERS_MATCHER = re.compile(r'airflow\.providers\.(.*)\.hooks.*')
 
 _MIN_SUPPORTED_PROVIDERS_VERSION = {
@@ -103,12 +123,14 @@ def __init__(
         conn_id: str | None = None,
         database: str | None = None,
         hook_params: dict | None = None,
+        retry_on_failure: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.conn_id = conn_id
         self.database = database
         self.hook_params = {} if hook_params is None else hook_params
+        self.retry_on_failure = retry_on_failure
 
     @cached_property
     def _hook(self):
@@ -210,12 +232,17 @@ class SQLColumnCheckOperator(BaseSQLOperator):
 
     template_fields = ("partition_clause",)
 
+    sql_check_template = """
+        SELECT '{column}' AS col_name, '{check}' AS check_type, {column}_{check} AS check_result
+        FROM (SELECT {check_statement} AS {column}_{check} FROM {table}) AS sq
+    """
+
     column_checks = {
-        "null_check": "SUM(CASE WHEN column IS NULL THEN 1 ELSE 0 END) AS column_null_check",
-        "distinct_check": "COUNT(DISTINCT(column)) AS column_distinct_check",
-        "unique_check": "COUNT(column) - COUNT(DISTINCT(column)) AS column_unique_check",
-        "min": "MIN(column) AS column_min",
-        "max": "MAX(column) AS column_max",
+        "null_check": "SUM(CASE WHEN {column} IS NULL THEN 1 ELSE 0 END)",
+        "distinct_check": "COUNT(DISTINCT({column}))",
+        "unique_check": "COUNT({column}) - COUNT(DISTINCT({column}))",
+        "min": "MIN({column})",
+        "max": "MAX({column})",
     }
 
     def __init__(
@@ -229,46 +256,58 @@ def __init__(
         **kwargs,
     ):
         super().__init__(conn_id=conn_id, database=database, **kwargs)
-        for checks in column_mapping.values():
-            for check, check_values in checks.items():
-                self._column_mapping_validation(check, check_values)
 
         self.table = table
         self.column_mapping = column_mapping
+        checks_sql = ""
+        for column, checks in self.column_mapping.items():
+            for check, check_values in checks.items():
+                self._column_mapping_validation(check, check_values)
+            checks_list = [*checks]
+            checks_sql = checks_sql + " UNION ALL ".join(
+                [
+                    self.sql_check_template.format(
+                        check_statement=self.column_checks[check].format(column=column),
+                        check=check,
+                        table=self.table,
+                        column=column,
+                    )
+                    for check in checks_list
+                ]
+            )
         self.partition_clause = partition_clause
-        # OpenLineage needs a valid SQL query with the input/output table(s) to parse
-        self.sql = f"SELECT * FROM {self.table};"
+        partition_clause_statement = f"WHERE {self.partition_clause}" if self.partition_clause else ""
+        self.sql = f"""
+            SELECT col_name, check_type, check_result FROM ({checks_sql})
+            AS check_columns {partition_clause_statement}
+        """
 
     def execute(self, context: Context):
         hook = self.get_db_hook()
         failed_tests = []
-        for column in self.column_mapping:
-            checks = [*self.column_mapping[column]]
-            checks_sql = ",".join([self.column_checks[check].replace("column", column) for check in checks])
-            partition_clause_statement = f"WHERE {self.partition_clause}" if self.partition_clause else ""
-            self.sql = f"SELECT {checks_sql} FROM {self.table} {partition_clause_statement};"
-            records = hook.get_first(self.sql)
+        records = hook.get_records(self.sql)
 
-            if not records:
-                raise AirflowException(f"The following query returned zero rows: {self.sql}")
+        if not records:
+            raise AirflowException(f"The following query returned zero rows: {self.sql}")
 
-            self.log.info("Record: %s", records)
+        self.log.info("Record: %s", records)
 
-            for idx, result in enumerate(records):
-                tolerance = self.column_mapping[column][checks[idx]].get("tolerance")
+        for row in records:
+            column, check, result = row
+            tolerance = self.column_mapping[column][check].get("tolerance")
 
-                self.column_mapping[column][checks[idx]]["result"] = result
-                self.column_mapping[column][checks[idx]]["success"] = self._get_match(
-                    self.column_mapping[column][checks[idx]], result, tolerance
-                )
+            self.column_mapping[column][check]["result"] = result
+            self.column_mapping[column][check]["success"] = self._get_match(
+                self.column_mapping[column][check], result, tolerance
+            )
 
             failed_tests.extend(_get_failed_checks(self.column_mapping[column], column))
         if failed_tests:
-            raise AirflowException(
-                f"Test failed.\nResults:\n{records!s}\n"
-                "The following tests have failed:"
-                f"\n{''.join(failed_tests)}"
-            )
+            exception_string = f"""
+                Test failed.\nResults:\n{records!s}\n
+                The following tests have failed:
+                \n{''.join(failed_tests)}"""
+            _raise_exception(exception_string, self.retry_on_failure)
 
         self.log.info("All tests have passed")
 
@@ -399,8 +438,8 @@ class SQLTableCheckOperator(BaseSQLOperator):
     template_fields = ("partition_clause",)
 
     sql_check_template = """
-        SELECT '_check_name' AS check_name, MIN(_check_name) AS check_result
-        FROM (SELECT CASE WHEN check_statement THEN 1 ELSE 0 END AS _check_name FROM table) AS sq
+        SELECT {check_name} AS check_name, MIN({check_name}) AS check_result
+        FROM (SELECT CASE WHEN {check_statement} THEN 1 ELSE 0 END AS {check_name} FROM {table}) AS sq
     """
 
     def __init__(
@@ -418,16 +457,11 @@ def __init__(
         self.table = table
         self.checks = checks
         self.partition_clause = partition_clause
-        # OpenLineage needs a valid SQL query with the input/output table(s) to parse
-        self.sql = f"SELECT * FROM {self.table};"
-
-    def execute(self, context: Context):
-        hook = self.get_db_hook()
         checks_sql = " UNION ALL ".join(
             [
-                self.sql_check_template.replace("check_statement", value["check_statement"])
-                .replace("_check_name", check_name)
-                .replace("table", self.table)
+                self.sql_check_template.format(
+                    check_statement=value["check_statement"], check_name=check_name, table=self.table
+                )
                 for check_name, value in self.checks.items()
             ]
         )
@@ -437,6 +471,8 @@ def execute(self, context: Context):
             AS check_table {partition_clause_statement}
         """
 
+    def execute(self, context: Context):
+        hook = self.get_db_hook()
         records = hook.get_records(self.sql)
 
         if not records:
@@ -446,15 +482,16 @@ def execute(self, context: Context):
 
         for row in records:
             check, result = row
-            self.checks[check]["success"] = parse_boolean(str(result))
+            self.checks[check]["success"] = _parse_boolean(str(result))
 
         failed_tests = _get_failed_checks(self.checks)
         if failed_tests:
-            raise AirflowException(
-                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n"
-                "The following tests have failed:"
-                f"\n{', '.join(failed_tests)}"
-            )
+            exception_string = f"""
+                Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n
+                The following tests have failed:
+                \n{', '.join(failed_tests)}
+            """
+            _raise_exception(exception_string, self.retry_on_failure)
 
         self.log.info("All tests have passed")
 
@@ -514,7 +551,9 @@ def execute(self, context: Context):
         if not records:
             raise AirflowException("The query returned None")
         elif not all(bool(r) for r in records):
-            raise AirflowException(f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}")
+            _raise_exception(
+                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}", self.retry_on_failure
+            )
 
         self.log.info("Success.")
 
@@ -591,7 +630,7 @@ def execute(self, context: Context):
             tests = []
 
         if not all(tests):
-            raise AirflowException(error_msg)
+            _raise_exception(error_msg, self.retry_on_failure)
 
     def _to_float(self, records):
         return [float(record) for record in records]
@@ -743,7 +782,9 @@ def execute(self, context: Context):
                     ratios[k],
                     self.metrics_thresholds[k],
                 )
-            raise AirflowException(f"The following tests have failed:\n {', '.join(sorted(failed_tests))}")
+            _raise_exception(
+                f"The following tests have failed:\n {', '.join(sorted(failed_tests))}", self.retry_on_failure
+            )
 
         self.log.info("All tests have passed")
 
@@ -820,7 +861,7 @@ def execute(self, context: Context):
                 f'Result: {result} is not within thresholds '
                 f'{meta_data.get("min_threshold")} and {meta_data.get("max_threshold")}'
             )
-            raise AirflowException(error_msg)
+            _raise_exception(error_msg, self.retry_on_failure)
 
         self.log.info("Test %s Successful.", self.task_id)
 
@@ -903,7 +944,7 @@ def execute(self, context: Context):
                     follow_branch = self.follow_task_ids_if_true
             elif isinstance(query_result, str):
                 # return result is not Boolean, try to convert from String to Boolean
-                if parse_boolean(query_result):
+                if _parse_boolean(query_result):
                     follow_branch = self.follow_task_ids_if_true
             elif isinstance(query_result, int):
                 if bool(query_result):
@@ -921,17 +962,3 @@ def execute(self, context: Context):
             )
 
         self.skip_all_except(context["ti"], follow_branch)
-
-
-def _convert_to_float_if_possible(s):
-    """
-    A small helper function to convert a string to a numeric value
-    if appropriate
-
-    :param s: the string to be converted
-    """
-    try:
-        ret = float(s)
-    except (ValueError, TypeError):
-        ret = s
-    return ret