From 711d6bc1e0fd62c5fbcde2ae7e9311b5042d128e Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 27 Sep 2022 10:44:58 -0400
Subject: [PATCH 01/63] Fix BigQueryTableCheckOperator test

Signed-off-by: Benji Lampel <benjamin@astronomer.io>
---
 .../providers/google/cloud/bigquery/example_bigquery_queries.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py b/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py
index 52637d37ce22d..d7375b21784f8 100644
--- a/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py
+++ b/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py
@@ -223,7 +223,7 @@
         table_check = BigQueryTableCheckOperator(
             task_id="table_check",
             table=f"{DATASET}.{TABLE_1}",
-            checks={"row_count_check": {"check_statement": {"COUNT(*) = 4"}}},
+            checks={"row_count_check": {"check_statement": "COUNT(*) = 4"}},
         )
         # [END howto_operator_bigquery_table_check]
 

From 3c97433eb2712d5088d8efe08ae02702dc654d14 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 27 Sep 2022 10:50:43 -0400
Subject: [PATCH 02/63] Remove job_id generation in table/col check operators

The job_id is automatically generated by hook.insert_job()
if an empty string is passed, so job_id generation in the
operator is removed in favor of the existing code.

Signed-off-by: Benji Lampel <benjamin@astronomer.io>
---
 .../providers/google/cloud/operators/bigquery.py | 16 ++--------------
 1 file changed, 2 insertions(+), 14 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 33e51833e59a6..e7162e8c72c68 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -607,13 +607,7 @@ def execute(self, context=None):
             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};"
 
-            job_id = hook.generate_job_id(
-                dag_id=self.dag_id,
-                task_id=self.task_id,
-                logical_date=context["logical_date"],
-                configuration=self.configuration,
-            )
-            job = self._submit_job(hook, job_id=job_id)
+            job = self._submit_job(hook, job_id="")
             context["ti"].xcom_push(key="job_id", value=job.job_id)
             records = list(job.result().to_dataframe().values.flatten())
 
@@ -727,13 +721,7 @@ def execute(self, context=None):
         self.sql = f"SELECT check_name, check_result FROM ({checks_sql}) "
         f"AS check_table {partition_clause_statement};"
 
-        job_id = hook.generate_job_id(
-            dag_id=self.dag_id,
-            task_id=self.task_id,
-            logical_date=context["logical_date"],
-            configuration=self.configuration,
-        )
-        job = self._submit_job(hook, job_id=job_id)
+        job = self._submit_job(hook, job_id="")
         context["ti"].xcom_push(key="job_id", value=job.job_id)
         records = job.result().to_dataframe()
 

From 11a26fd5789b2d5680a4aeedd3bf7ac2f4855e4d Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 14:27:27 -0400
Subject: [PATCH 03/63] Rework SQL query building

SQL query building is moved to the init() method of the column
and table check operators to lessen the amount of duplicate code
in the child operator. It also has the added effect of, ideally,
passing a more complete query to OpenLineage.

In doing the above, the column check operator had to be reworked and
now matches the logic of the table check operator in terms of
returning multiple rows and only sending one query to the database.
---
 airflow/providers/common/sql/operators/sql.py | 84 +++++++++++--------
 .../google/cloud/operators/bigquery.py        | 45 ++++------
 2 files changed, 65 insertions(+), 64 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 536b72fbff8e8..66e1c4c80bde0 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -210,12 +210,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_statment} 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,38 +234,50 @@ 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:
@@ -399,8 +416,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 +435,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 +449,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:
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index e7162e8c72c68..8226357b548b0 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -601,28 +601,27 @@ def execute(self, context=None):
         """Perform checks on the given columns."""
         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};"
 
-            job = self._submit_job(hook, job_id="")
-            context["ti"].xcom_push(key="job_id", value=job.job_id)
-            records = list(job.result().to_dataframe().values.flatten())
+        job = self._submit_job(hook, job_id="")
+        context["ti"].xcom_push(key="job_id", value=job.job_id)
+        records = job.result().to_dataframe()
 
-            if not records:
-                raise AirflowException(f"The following query returned zero rows: {self.sql}")
+        if records.empty:
+            raise AirflowException(f"The following query returned zero rows: {self.sql}")
 
-            self.log.info("Record: %s", records)
+        records.columns = records.columns.str.lower()
+        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.iterrows():
+            column = row[1].get("col_name")
+            check = row[1].get("check_type")
+            result = row[1].get("check_result")
+            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:
@@ -709,18 +708,6 @@ def _submit_job(
     def execute(self, context=None):
         """Execute the given checks on the table."""
         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)
-                for check_name, value in self.checks.items()
-            ]
-        )
-        partition_clause_statement = f"WHERE {self.partition_clause}" if self.partition_clause else ""
-        self.sql = f"SELECT check_name, check_result FROM ({checks_sql}) "
-        f"AS check_table {partition_clause_statement};"
-
         job = self._submit_job(hook, job_id="")
         context["ti"].xcom_push(key="job_id", value=job.job_id)
         records = job.result().to_dataframe()

From de7fafb24439026ea4966ee770a46684687c2a1b Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 14:52:13 -0400
Subject: [PATCH 04/63] Remove self.sql overwrite parameter in BigQuery check
 operators.

---
 airflow/providers/google/cloud/operators/bigquery.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 8226357b548b0..1db64ba338eef 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -578,8 +578,6 @@ def __init__(
         self.location = location
         self.impersonation_chain = impersonation_chain
         self.labels = labels
-        # OpenLineage needs a valid SQL query with the input/output table(s) to parse
-        self.sql = ""
 
     def _submit_job(
         self,
@@ -686,8 +684,6 @@ def __init__(
         self.location = location
         self.impersonation_chain = impersonation_chain
         self.labels = labels
-        # OpenLineage needs a valid SQL query with the input/output table(s) to parse
-        self.sql = ""
 
     def _submit_job(
         self,

From b315f0b9f4ce05ab904a50d084b459f79c37103f Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 15:25:39 -0400
Subject: [PATCH 05/63] Add option to fail check operators without retries

Adds a new parameter, retry_on_failure, and a new function to
determine if operators should retry or not on test failure.
---
 airflow/providers/common/sql/operators/sql.py | 43 ++++++++++++-------
 .../google/cloud/operators/bigquery.py        | 26 ++++++-----
 2 files changed, 43 insertions(+), 26 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 66e1c4c80bde0..63950889c6a79 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/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
@@ -60,6 +60,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 +109,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):
@@ -281,11 +289,11 @@ def execute(self, context: Context):
 
             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")
 
@@ -464,11 +472,12 @@ def execute(self, context: Context):
 
         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")
 
@@ -528,7 +537,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.")
 
@@ -605,7 +616,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]
@@ -757,7 +768,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")
 
@@ -834,7 +847,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)
 
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 1db64ba338eef..73879530be57a 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -39,6 +39,7 @@
     SQLTableCheckOperator,
     SQLValueCheckOperator,
     _get_failed_checks,
+    _raise_exception,
     parse_boolean,
 )
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
@@ -264,7 +265,9 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
         if not records:
             raise AirflowException("The query returned empty results")
         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("Record: %s", event["records"])
         self.log.info("Success.")
 
@@ -623,11 +626,11 @@ def execute(self, context=None):
 
             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")
 
@@ -721,11 +724,12 @@ def execute(self, context=None):
 
         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")
 

From a560b459034a757ff038985c3f6dfd4c4f364be0 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 16:03:29 -0400
Subject: [PATCH 06/63] Rename and reorder private functions, fix typo

---
 airflow/providers/common/sql/operators/sql.py | 36 +++++++++----------
 .../google/cloud/operators/bigquery.py        |  4 +--
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 63950889c6a79..846bbe3bd3703 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -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.
@@ -220,7 +234,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
 
     sql_check_template = """
         SELECT '{column}' AS col_name, '{check}' AS check_type, {column}_{check} AS check_result
-        FROM (SELECT {check_statment} AS {column}_{check} FROM {table}) AS sq
+        FROM (SELECT {check_statement} AS {column}_{check} FROM {table}) AS sq
     """
 
     column_checks = {
@@ -468,7 +482,7 @@ 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:
@@ -930,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):
@@ -948,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
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 73879530be57a..e1c21fe329b03 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -39,8 +39,8 @@
     SQLTableCheckOperator,
     SQLValueCheckOperator,
     _get_failed_checks,
+    _parse_boolean,
     _raise_exception,
-    parse_boolean,
 )
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
@@ -720,7 +720,7 @@ def execute(self, context=None):
         for row in records.iterrows():
             check = row[1].get("check_name")
             result = row[1].get("check_result")
-            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:

From 8c987fd047cf01f1ee839a37190b46bfdad34270 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 09:56:33 -0400
Subject: [PATCH 07/63] Update airflow/providers/common/sql/operators/sql.py

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>
---
 airflow/providers/common/sql/operators/sql.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 846bbe3bd3703..908accf71e963 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -74,7 +74,7 @@ def _get_failed_checks(checks, col=None):
     ]
 
 
-def _raise_exception(exception_string, retry_on_failure):
+def _raise_exception(exception_string: str, retry_on_failure: bool) -> NoReturn:
     if retry_on_failure:
         raise AirflowException(exception_string)
     raise AirflowFailException(exception_string)

From 4364a3987ca8bd3ff2ba3fe9b50697439549d566 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 10:35:42 -0400
Subject: [PATCH 08/63] Small updates to helper functions

---
 airflow/providers/common/sql/operators/sql.py | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 908accf71e963..842dd3158309d 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING, Any, Iterable, Mapping, Sequence, SupportsAbs
+from typing import TYPE_CHECKING, Any, Iterable, Mapping, NoReturn, Sequence, SupportsAbs
 
 from packaging.version import Version
 
@@ -33,18 +33,11 @@
     from airflow.utils.context import Context
 
 
-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
-    """
+def _convert_to_float_if_possible(s: str) -> float | str:
     try:
-        ret = float(s)
+        return float(s)
     except (ValueError, TypeError):
-        ret = s
-    return ret
+        return s
 
 
 def _parse_boolean(val: str) -> str | bool:
@@ -60,7 +53,7 @@ def _parse_boolean(val: str) -> str | bool:
     raise ValueError(f"{val!r} is not a boolean-like string value")
 
 
-def _get_failed_checks(checks, col=None):
+def _get_failed_checks(checks: dict[str, Any], col: str | None = None):
     if col:
         return [
             f"Column: {col}\nCheck: {check},\nCheck Values: {check_values}\n"

From 30839fbe8f2f90fbb7878875f3b78ca9a5ef245c Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 10:42:40 -0400
Subject: [PATCH 09/63] Use _raise_exception when query returns 0 rows

---
 airflow/providers/common/sql/operators/sql.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 842dd3158309d..2605a94e7af9c 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -281,7 +281,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            raise AirflowException(f"The following query returned zero rows: {self.sql}")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         self.log.info("Record: %s", records)
 
@@ -469,7 +469,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            raise AirflowException(f"The following query returned zero rows: {self.sql}")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         self.log.info("Record:\n%s", records)
 
@@ -542,7 +542,7 @@ def execute(self, context: Context):
 
         self.log.info("Record: %s", records)
         if not records:
-            raise AirflowException("The query returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
         elif not all(bool(r) for r in records):
             _raise_exception(
                 f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}", self.retry_on_failure
@@ -594,7 +594,7 @@ def execute(self, context: Context):
         records = self.get_db_hook().get_first(self.sql)
 
         if not records:
-            raise AirflowException("The query returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         pass_value_conv = _convert_to_float_if_possible(self.pass_value)
         is_numeric_value_check = isinstance(pass_value_conv, float)
@@ -695,7 +695,7 @@ def __init__(
         if ratio_formula not in self.ratio_formulas:
             msg_template = "Invalid diff_method: {diff_method}. Supported diff methods are: {diff_methods}"
 
-            raise AirflowException(
+            raise AirflowFailException(
                 msg_template.format(diff_method=ratio_formula, diff_methods=self.ratio_formulas)
             )
         self.ratio_formula = ratio_formula
@@ -720,9 +720,9 @@ def execute(self, context: Context):
         row1 = hook.get_first(self.sql1)
 
         if not row2:
-            raise AirflowException(f"The query {self.sql2} returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql2}", self.retry_on_failure)
         if not row1:
-            raise AirflowException(f"The query {self.sql1} returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql1}", self.retry_on_failure)
 
         current = dict(zip(self.metrics_sorted, row1))
         reference = dict(zip(self.metrics_sorted, row2))
@@ -820,6 +820,8 @@ def __init__(
     def execute(self, context: Context):
         hook = self.get_db_hook()
         result = hook.get_first(self.sql)[0]
+        if not result:
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         if isinstance(self.min_threshold, float):
             lower_bound = self.min_threshold

From a787d3a357b6ba8eb11af97cd613558eb172afea Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 14:47:46 -0400
Subject: [PATCH 10/63] Update tests and operator code

Updates tests to reflect changes in operator code, and fixed
bugs in operators as well. Mainly moving the code to check for
failed tests into the column and table check operators as it works
slightly differently for each and doesn't make much sense as a
top-level function.
---
 airflow/providers/common/sql/operators/sql.py | 31 ++++-----
 .../common/sql/operators/test_sql.py          | 64 ++++++++++++++-----
 2 files changed, 61 insertions(+), 34 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 2605a94e7af9c..b8d4b05679918 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -53,20 +53,6 @@ def _parse_boolean(val: str) -> str | bool:
     raise ValueError(f"{val!r} is not a boolean-like string value")
 
 
-def _get_failed_checks(checks: dict[str, Any], col: str | None = None):
-    if col:
-        return [
-            f"Column: {col}\nCheck: {check},\nCheck Values: {check_values}\n"
-            for check, check_values in checks.items()
-            if not check_values["success"]
-        ]
-    return [
-        f"\tCheck: {check},\n\tCheck Values: {check_values}\n"
-        for check, check_values in checks.items()
-        if not check_values["success"]
-    ]
-
-
 def _raise_exception(exception_string: str, retry_on_failure: bool) -> NoReturn:
     if retry_on_failure:
         raise AirflowException(exception_string)
@@ -276,8 +262,8 @@ def __init__(
         """
 
     def execute(self, context: Context):
-        hook = self.get_db_hook()
         failed_tests = []
+        hook = self.get_db_hook()
         records = hook.get_records(self.sql)
 
         if not records:
@@ -294,7 +280,14 @@ def execute(self, context: Context):
                 self.column_mapping[column][check], result, tolerance
             )
 
-            failed_tests.extend(_get_failed_checks(self.column_mapping[column], column))
+        for col, checks in self.column_mapping.items():
+            failed_tests.extend(
+                [
+                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+                    for check, check_values in checks.items()
+                    if not check_values["success"]
+                ]
+            )
         if failed_tests:
             exception_string = f"""
                 Test failed.\nResults:\n{records!s}\n
@@ -477,7 +470,11 @@ def execute(self, context: Context):
             check, result = row
             self.checks[check]["success"] = _parse_boolean(str(result))
 
-        failed_tests = _get_failed_checks(self.checks)
+        failed_tests = [
+            f"\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for check, check_values in self.checks.items()
+            if not check_values["success"]
+        ]
         if failed_tests:
             exception_string = f"""
                 Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 373b1b99aed2b..1c118f4d733d4 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -45,9 +45,6 @@
 
 
 class MockHook:
-    def get_first(self):
-        return
-
     def get_records(self):
         return
 
@@ -70,15 +67,15 @@ class TestColumnCheckOperator:
 
     invalid_column_mapping = {"Y": {"invalid_check_name": {"expectation": 5}}}
 
-    def _construct_operator(self, monkeypatch, column_mapping, return_vals):
-        def get_first_return(*arg):
-            return return_vals
+    def _construct_operator(self, monkeypatch, column_mapping, records):
+        def get_records(*arg):
+            return records
 
         operator = SQLColumnCheckOperator(
             task_id="test_task", table="test_table", column_mapping=column_mapping
         )
         monkeypatch.setattr(operator, "get_db_hook", _get_mock_db_hook)
-        monkeypatch.setattr(MockHook, "get_first", get_first_return)
+        monkeypatch.setattr(MockHook, "get_records", get_records)
         return operator
 
     def test_check_not_in_column_checks(self, monkeypatch):
@@ -86,29 +83,62 @@ def test_check_not_in_column_checks(self, monkeypatch):
             self._construct_operator(monkeypatch, self.invalid_column_mapping, ())
 
     def test_pass_all_checks_exact_check(self, monkeypatch):
-        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 10, 10, 1, 19))
+        records = [
+            ('X', 'null_check', 0),
+            ('X', 'distinct_check', 10),
+            ('X', 'unique_check', 10),
+            ('X', 'min', 1),
+            ('X', 'max', 19),
+        ]
+        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
 
     def test_max_less_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
-            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 10, 10, 1, 21))
+            records = [
+                ('X', 'null_check', 1),
+                ('X', 'distinct_check', 10),
+                ('X', 'unique_check', 10),
+                ('X', 'min', 1),
+                ('X', 'max', 21),
+            ]
+            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
             operator.execute(context=MagicMock())
             assert operator.column_mapping["X"]["max"]["success"] is False
 
     def test_max_greater_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
-            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 10, 10, 1, 9))
+            records = [
+                ('X', 'null_check', 1),
+                ('X', 'distinct_check', 10),
+                ('X', 'unique_check', 10),
+                ('X', 'min', 1),
+                ('X', 'max', 9),
+            ]
+            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
             operator.execute(context=MagicMock())
             assert operator.column_mapping["X"]["max"]["success"] is False
 
     def test_pass_all_checks_inexact_check(self, monkeypatch):
-        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 9, 12, 0, 15))
+        records = [
+            ('X', 'null_check', 0),
+            ('X', 'distinct_check', 9),
+            ('X', 'unique_check', 12),
+            ('X', 'min', 0),
+            ('X', 'max', 15),
+        ]
+        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
 
     def test_fail_all_checks_check(self, monkeypatch):
-        operator = operator = self._construct_operator(
-            monkeypatch, self.valid_column_mapping, (1, 12, 11, -1, 20)
-        )
+        records = [
+            ('X', 'null_check', 1),
+            ('X', 'distinct_check', 12),
+            ('X', 'unique_check', 11),
+            ('X', 'min', -1),
+            ('X', 'max', 20),
+        ]
+        operator = operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         with pytest.raises(AirflowException):
             operator.execute(context=MagicMock())
 
@@ -120,9 +150,9 @@ class TestTableCheckOperator:
         "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
     }
 
-    def _construct_operator(self, monkeypatch, checks, return_df):
+    def _construct_operator(self, monkeypatch, checks, records):
         def get_records(*arg):
-            return return_df
+            return records
 
         operator = SQLTableCheckOperator(task_id="test_task", table="test_table", checks=checks)
         monkeypatch.setattr(operator, "get_db_hook", _get_mock_db_hook)
@@ -263,7 +293,7 @@ def setUp(self):
     def test_execute_no_records(self, mock_get_db_hook):
         mock_get_db_hook.return_value.get_first.return_value = []
 
-        with pytest.raises(AirflowException, match=r"The query returned None"):
+        with pytest.raises(AirflowException, match=r"The following query returned zero rows: sql"):
             self._operator.execute({})
 
     @mock.patch.object(SQLCheckOperator, "get_db_hook")

From 9005a4b502e923b831ff0d7519bfe105b3b2596b Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 16:40:25 -0400
Subject: [PATCH 11/63] Fix _failed_checks() to match update in parent
 operators

---
 .../providers/google/cloud/operators/bigquery.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index e1c21fe329b03..82749bb57edf6 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -38,7 +38,6 @@
     SQLIntervalCheckOperator,
     SQLTableCheckOperator,
     SQLValueCheckOperator,
-    _get_failed_checks,
     _parse_boolean,
     _raise_exception,
 )
@@ -624,7 +623,14 @@ def execute(self, context=None):
                 self.column_mapping[column][check], result, tolerance
             )
 
-            failed_tests.extend(_get_failed_checks(self.column_mapping[column], column))
+        for col, checks in self.column_mapping.items():
+            failed_tests.extend(
+                [
+                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+                    for check, check_values in checks.items()
+                    if not check_values["success"]
+                ]
+            )
         if failed_tests:
             exception_string = f"""
                 Test failed.\nResults:\n{records!s}\n
@@ -722,7 +728,11 @@ def execute(self, context=None):
             result = row[1].get("check_result")
             self.checks[check]["success"] = _parse_boolean(str(result))
 
-        failed_tests = _get_failed_checks(self.checks)
+        failed_tests = [
+            f"\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for check, check_values in self.checks.items()
+            if not check_values["success"]
+        ]
         if failed_tests:
             exception_string = f"""
                 Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n

From c00c934c2a877bff4aec12d0d580c8f46a5731c6 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Fri, 30 Sep 2022 09:33:39 -0400
Subject: [PATCH 12/63] Insert quotes around check_name table op's sql template
 to fix sql error

---
 airflow/providers/common/sql/operators/sql.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index b8d4b05679918..8365c86bb57ae 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -424,7 +424,7 @@ class SQLTableCheckOperator(BaseSQLOperator):
     template_fields = ("partition_clause",)
 
     sql_check_template = """
-        SELECT {check_name} AS check_name, MIN({check_name}) AS check_result
+        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
     """
 

From b2eb10abf0a96885584219b1fb33af2a4130cf0d Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 27 Sep 2022 10:44:58 -0400
Subject: [PATCH 13/63] Fix BigQueryTableCheckOperator test

Signed-off-by: Benji Lampel <benjamin@astronomer.io>
---
 .../providers/google/cloud/bigquery/example_bigquery_queries.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py b/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py
index 52637d37ce22d..d7375b21784f8 100644
--- a/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py
+++ b/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py
@@ -223,7 +223,7 @@
         table_check = BigQueryTableCheckOperator(
             task_id="table_check",
             table=f"{DATASET}.{TABLE_1}",
-            checks={"row_count_check": {"check_statement": {"COUNT(*) = 4"}}},
+            checks={"row_count_check": {"check_statement": "COUNT(*) = 4"}},
         )
         # [END howto_operator_bigquery_table_check]
 

From e13c59c0f444098af5af609c317918bf19a309f2 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 27 Sep 2022 10:50:43 -0400
Subject: [PATCH 14/63] Remove job_id generation in table/col check operators

The job_id is automatically generated by hook.insert_job()
if an empty string is passed, so job_id generation in the
operator is removed in favor of the existing code.

Signed-off-by: Benji Lampel <benjamin@astronomer.io>
---
 .../providers/google/cloud/operators/bigquery.py | 16 ++--------------
 1 file changed, 2 insertions(+), 14 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 33e51833e59a6..e7162e8c72c68 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -607,13 +607,7 @@ def execute(self, context=None):
             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};"
 
-            job_id = hook.generate_job_id(
-                dag_id=self.dag_id,
-                task_id=self.task_id,
-                logical_date=context["logical_date"],
-                configuration=self.configuration,
-            )
-            job = self._submit_job(hook, job_id=job_id)
+            job = self._submit_job(hook, job_id="")
             context["ti"].xcom_push(key="job_id", value=job.job_id)
             records = list(job.result().to_dataframe().values.flatten())
 
@@ -727,13 +721,7 @@ def execute(self, context=None):
         self.sql = f"SELECT check_name, check_result FROM ({checks_sql}) "
         f"AS check_table {partition_clause_statement};"
 
-        job_id = hook.generate_job_id(
-            dag_id=self.dag_id,
-            task_id=self.task_id,
-            logical_date=context["logical_date"],
-            configuration=self.configuration,
-        )
-        job = self._submit_job(hook, job_id=job_id)
+        job = self._submit_job(hook, job_id="")
         context["ti"].xcom_push(key="job_id", value=job.job_id)
         records = job.result().to_dataframe()
 

From 708ad7a1ba5956593673a5f24b7636f76d8c826b Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 14:27:27 -0400
Subject: [PATCH 15/63] Rework SQL query building

SQL query building is moved to the init() method of the column
and table check operators to lessen the amount of duplicate code
in the child operator. It also has the added effect of, ideally,
passing a more complete query to OpenLineage.

In doing the above, the column check operator had to be reworked and
now matches the logic of the table check operator in terms of
returning multiple rows and only sending one query to the database.
---
 airflow/providers/common/sql/operators/sql.py | 84 +++++++++++--------
 .../google/cloud/operators/bigquery.py        | 45 ++++------
 2 files changed, 65 insertions(+), 64 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 536b72fbff8e8..66e1c4c80bde0 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -210,12 +210,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_statment} 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,38 +234,50 @@ 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:
@@ -399,8 +416,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 +435,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 +449,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:
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index e7162e8c72c68..8226357b548b0 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -601,28 +601,27 @@ def execute(self, context=None):
         """Perform checks on the given columns."""
         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};"
 
-            job = self._submit_job(hook, job_id="")
-            context["ti"].xcom_push(key="job_id", value=job.job_id)
-            records = list(job.result().to_dataframe().values.flatten())
+        job = self._submit_job(hook, job_id="")
+        context["ti"].xcom_push(key="job_id", value=job.job_id)
+        records = job.result().to_dataframe()
 
-            if not records:
-                raise AirflowException(f"The following query returned zero rows: {self.sql}")
+        if records.empty:
+            raise AirflowException(f"The following query returned zero rows: {self.sql}")
 
-            self.log.info("Record: %s", records)
+        records.columns = records.columns.str.lower()
+        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.iterrows():
+            column = row[1].get("col_name")
+            check = row[1].get("check_type")
+            result = row[1].get("check_result")
+            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:
@@ -709,18 +708,6 @@ def _submit_job(
     def execute(self, context=None):
         """Execute the given checks on the table."""
         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)
-                for check_name, value in self.checks.items()
-            ]
-        )
-        partition_clause_statement = f"WHERE {self.partition_clause}" if self.partition_clause else ""
-        self.sql = f"SELECT check_name, check_result FROM ({checks_sql}) "
-        f"AS check_table {partition_clause_statement};"
-
         job = self._submit_job(hook, job_id="")
         context["ti"].xcom_push(key="job_id", value=job.job_id)
         records = job.result().to_dataframe()

From 62c01bf6d810b0bd7273636bdaa89058b9329251 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 14:52:13 -0400
Subject: [PATCH 16/63] Remove self.sql overwrite parameter in BigQuery check
 operators.

---
 airflow/providers/google/cloud/operators/bigquery.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 8226357b548b0..1db64ba338eef 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -578,8 +578,6 @@ def __init__(
         self.location = location
         self.impersonation_chain = impersonation_chain
         self.labels = labels
-        # OpenLineage needs a valid SQL query with the input/output table(s) to parse
-        self.sql = ""
 
     def _submit_job(
         self,
@@ -686,8 +684,6 @@ def __init__(
         self.location = location
         self.impersonation_chain = impersonation_chain
         self.labels = labels
-        # OpenLineage needs a valid SQL query with the input/output table(s) to parse
-        self.sql = ""
 
     def _submit_job(
         self,

From f66bfe4aa2c0483da4674b604bcb7079781af60c Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 15:25:39 -0400
Subject: [PATCH 17/63] Add option to fail check operators without retries

Adds a new parameter, retry_on_failure, and a new function to
determine if operators should retry or not on test failure.
---
 airflow/providers/common/sql/operators/sql.py | 43 ++++++++++++-------
 .../google/cloud/operators/bigquery.py        | 26 ++++++-----
 2 files changed, 43 insertions(+), 26 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 66e1c4c80bde0..63950889c6a79 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/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
@@ -60,6 +60,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 +109,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):
@@ -281,11 +289,11 @@ def execute(self, context: Context):
 
             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")
 
@@ -464,11 +472,12 @@ def execute(self, context: Context):
 
         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")
 
@@ -528,7 +537,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.")
 
@@ -605,7 +616,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]
@@ -757,7 +768,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")
 
@@ -834,7 +847,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)
 
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 1db64ba338eef..73879530be57a 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -39,6 +39,7 @@
     SQLTableCheckOperator,
     SQLValueCheckOperator,
     _get_failed_checks,
+    _raise_exception,
     parse_boolean,
 )
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
@@ -264,7 +265,9 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
         if not records:
             raise AirflowException("The query returned empty results")
         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("Record: %s", event["records"])
         self.log.info("Success.")
 
@@ -623,11 +626,11 @@ def execute(self, context=None):
 
             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")
 
@@ -721,11 +724,12 @@ def execute(self, context=None):
 
         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")
 

From ef7c2a17e301ac7ebfca4f5031e22640ba1a0921 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 16:03:29 -0400
Subject: [PATCH 18/63] Rename and reorder private functions, fix typo

---
 airflow/providers/common/sql/operators/sql.py | 36 +++++++++----------
 .../google/cloud/operators/bigquery.py        |  4 +--
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 63950889c6a79..846bbe3bd3703 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -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.
@@ -220,7 +234,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
 
     sql_check_template = """
         SELECT '{column}' AS col_name, '{check}' AS check_type, {column}_{check} AS check_result
-        FROM (SELECT {check_statment} AS {column}_{check} FROM {table}) AS sq
+        FROM (SELECT {check_statement} AS {column}_{check} FROM {table}) AS sq
     """
 
     column_checks = {
@@ -468,7 +482,7 @@ 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:
@@ -930,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):
@@ -948,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
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 73879530be57a..e1c21fe329b03 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -39,8 +39,8 @@
     SQLTableCheckOperator,
     SQLValueCheckOperator,
     _get_failed_checks,
+    _parse_boolean,
     _raise_exception,
-    parse_boolean,
 )
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
@@ -720,7 +720,7 @@ def execute(self, context=None):
         for row in records.iterrows():
             check = row[1].get("check_name")
             result = row[1].get("check_result")
-            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:

From 586dc39da2148aeb4506405584a254312302c2d7 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 09:56:33 -0400
Subject: [PATCH 19/63] Update airflow/providers/common/sql/operators/sql.py

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>
---
 airflow/providers/common/sql/operators/sql.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 846bbe3bd3703..908accf71e963 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -74,7 +74,7 @@ def _get_failed_checks(checks, col=None):
     ]
 
 
-def _raise_exception(exception_string, retry_on_failure):
+def _raise_exception(exception_string: str, retry_on_failure: bool) -> NoReturn:
     if retry_on_failure:
         raise AirflowException(exception_string)
     raise AirflowFailException(exception_string)

From 419b167fc576718ec57cad58e9b5d157c0bcef83 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 10:35:42 -0400
Subject: [PATCH 20/63] Small updates to helper functions

---
 airflow/providers/common/sql/operators/sql.py | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 908accf71e963..842dd3158309d 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING, Any, Iterable, Mapping, Sequence, SupportsAbs
+from typing import TYPE_CHECKING, Any, Iterable, Mapping, NoReturn, Sequence, SupportsAbs
 
 from packaging.version import Version
 
@@ -33,18 +33,11 @@
     from airflow.utils.context import Context
 
 
-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
-    """
+def _convert_to_float_if_possible(s: str) -> float | str:
     try:
-        ret = float(s)
+        return float(s)
     except (ValueError, TypeError):
-        ret = s
-    return ret
+        return s
 
 
 def _parse_boolean(val: str) -> str | bool:
@@ -60,7 +53,7 @@ def _parse_boolean(val: str) -> str | bool:
     raise ValueError(f"{val!r} is not a boolean-like string value")
 
 
-def _get_failed_checks(checks, col=None):
+def _get_failed_checks(checks: dict[str, Any], col: str | None = None):
     if col:
         return [
             f"Column: {col}\nCheck: {check},\nCheck Values: {check_values}\n"

From 55f9bb818dad0b338dc89d0298c51896c78e1f2c Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 10:42:40 -0400
Subject: [PATCH 21/63] Use _raise_exception when query returns 0 rows

---
 airflow/providers/common/sql/operators/sql.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 842dd3158309d..2605a94e7af9c 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -281,7 +281,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            raise AirflowException(f"The following query returned zero rows: {self.sql}")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         self.log.info("Record: %s", records)
 
@@ -469,7 +469,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            raise AirflowException(f"The following query returned zero rows: {self.sql}")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         self.log.info("Record:\n%s", records)
 
@@ -542,7 +542,7 @@ def execute(self, context: Context):
 
         self.log.info("Record: %s", records)
         if not records:
-            raise AirflowException("The query returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
         elif not all(bool(r) for r in records):
             _raise_exception(
                 f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}", self.retry_on_failure
@@ -594,7 +594,7 @@ def execute(self, context: Context):
         records = self.get_db_hook().get_first(self.sql)
 
         if not records:
-            raise AirflowException("The query returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         pass_value_conv = _convert_to_float_if_possible(self.pass_value)
         is_numeric_value_check = isinstance(pass_value_conv, float)
@@ -695,7 +695,7 @@ def __init__(
         if ratio_formula not in self.ratio_formulas:
             msg_template = "Invalid diff_method: {diff_method}. Supported diff methods are: {diff_methods}"
 
-            raise AirflowException(
+            raise AirflowFailException(
                 msg_template.format(diff_method=ratio_formula, diff_methods=self.ratio_formulas)
             )
         self.ratio_formula = ratio_formula
@@ -720,9 +720,9 @@ def execute(self, context: Context):
         row1 = hook.get_first(self.sql1)
 
         if not row2:
-            raise AirflowException(f"The query {self.sql2} returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql2}", self.retry_on_failure)
         if not row1:
-            raise AirflowException(f"The query {self.sql1} returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql1}", self.retry_on_failure)
 
         current = dict(zip(self.metrics_sorted, row1))
         reference = dict(zip(self.metrics_sorted, row2))
@@ -820,6 +820,8 @@ def __init__(
     def execute(self, context: Context):
         hook = self.get_db_hook()
         result = hook.get_first(self.sql)[0]
+        if not result:
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         if isinstance(self.min_threshold, float):
             lower_bound = self.min_threshold

From e8ce8794cf62c224c57289f016c8c83563723a43 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 14:47:46 -0400
Subject: [PATCH 22/63] Update tests and operator code

Updates tests to reflect changes in operator code, and fixed
bugs in operators as well. Mainly moving the code to check for
failed tests into the column and table check operators as it works
slightly differently for each and doesn't make much sense as a
top-level function.
---
 airflow/providers/common/sql/operators/sql.py | 31 ++++-----
 .../common/sql/operators/test_sql.py          | 64 ++++++++++++++-----
 2 files changed, 61 insertions(+), 34 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 2605a94e7af9c..b8d4b05679918 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -53,20 +53,6 @@ def _parse_boolean(val: str) -> str | bool:
     raise ValueError(f"{val!r} is not a boolean-like string value")
 
 
-def _get_failed_checks(checks: dict[str, Any], col: str | None = None):
-    if col:
-        return [
-            f"Column: {col}\nCheck: {check},\nCheck Values: {check_values}\n"
-            for check, check_values in checks.items()
-            if not check_values["success"]
-        ]
-    return [
-        f"\tCheck: {check},\n\tCheck Values: {check_values}\n"
-        for check, check_values in checks.items()
-        if not check_values["success"]
-    ]
-
-
 def _raise_exception(exception_string: str, retry_on_failure: bool) -> NoReturn:
     if retry_on_failure:
         raise AirflowException(exception_string)
@@ -276,8 +262,8 @@ def __init__(
         """
 
     def execute(self, context: Context):
-        hook = self.get_db_hook()
         failed_tests = []
+        hook = self.get_db_hook()
         records = hook.get_records(self.sql)
 
         if not records:
@@ -294,7 +280,14 @@ def execute(self, context: Context):
                 self.column_mapping[column][check], result, tolerance
             )
 
-            failed_tests.extend(_get_failed_checks(self.column_mapping[column], column))
+        for col, checks in self.column_mapping.items():
+            failed_tests.extend(
+                [
+                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+                    for check, check_values in checks.items()
+                    if not check_values["success"]
+                ]
+            )
         if failed_tests:
             exception_string = f"""
                 Test failed.\nResults:\n{records!s}\n
@@ -477,7 +470,11 @@ def execute(self, context: Context):
             check, result = row
             self.checks[check]["success"] = _parse_boolean(str(result))
 
-        failed_tests = _get_failed_checks(self.checks)
+        failed_tests = [
+            f"\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for check, check_values in self.checks.items()
+            if not check_values["success"]
+        ]
         if failed_tests:
             exception_string = f"""
                 Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 373b1b99aed2b..1c118f4d733d4 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -45,9 +45,6 @@
 
 
 class MockHook:
-    def get_first(self):
-        return
-
     def get_records(self):
         return
 
@@ -70,15 +67,15 @@ class TestColumnCheckOperator:
 
     invalid_column_mapping = {"Y": {"invalid_check_name": {"expectation": 5}}}
 
-    def _construct_operator(self, monkeypatch, column_mapping, return_vals):
-        def get_first_return(*arg):
-            return return_vals
+    def _construct_operator(self, monkeypatch, column_mapping, records):
+        def get_records(*arg):
+            return records
 
         operator = SQLColumnCheckOperator(
             task_id="test_task", table="test_table", column_mapping=column_mapping
         )
         monkeypatch.setattr(operator, "get_db_hook", _get_mock_db_hook)
-        monkeypatch.setattr(MockHook, "get_first", get_first_return)
+        monkeypatch.setattr(MockHook, "get_records", get_records)
         return operator
 
     def test_check_not_in_column_checks(self, monkeypatch):
@@ -86,29 +83,62 @@ def test_check_not_in_column_checks(self, monkeypatch):
             self._construct_operator(monkeypatch, self.invalid_column_mapping, ())
 
     def test_pass_all_checks_exact_check(self, monkeypatch):
-        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 10, 10, 1, 19))
+        records = [
+            ('X', 'null_check', 0),
+            ('X', 'distinct_check', 10),
+            ('X', 'unique_check', 10),
+            ('X', 'min', 1),
+            ('X', 'max', 19),
+        ]
+        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
 
     def test_max_less_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
-            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 10, 10, 1, 21))
+            records = [
+                ('X', 'null_check', 1),
+                ('X', 'distinct_check', 10),
+                ('X', 'unique_check', 10),
+                ('X', 'min', 1),
+                ('X', 'max', 21),
+            ]
+            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
             operator.execute(context=MagicMock())
             assert operator.column_mapping["X"]["max"]["success"] is False
 
     def test_max_greater_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
-            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 10, 10, 1, 9))
+            records = [
+                ('X', 'null_check', 1),
+                ('X', 'distinct_check', 10),
+                ('X', 'unique_check', 10),
+                ('X', 'min', 1),
+                ('X', 'max', 9),
+            ]
+            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
             operator.execute(context=MagicMock())
             assert operator.column_mapping["X"]["max"]["success"] is False
 
     def test_pass_all_checks_inexact_check(self, monkeypatch):
-        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 9, 12, 0, 15))
+        records = [
+            ('X', 'null_check', 0),
+            ('X', 'distinct_check', 9),
+            ('X', 'unique_check', 12),
+            ('X', 'min', 0),
+            ('X', 'max', 15),
+        ]
+        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
 
     def test_fail_all_checks_check(self, monkeypatch):
-        operator = operator = self._construct_operator(
-            monkeypatch, self.valid_column_mapping, (1, 12, 11, -1, 20)
-        )
+        records = [
+            ('X', 'null_check', 1),
+            ('X', 'distinct_check', 12),
+            ('X', 'unique_check', 11),
+            ('X', 'min', -1),
+            ('X', 'max', 20),
+        ]
+        operator = operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         with pytest.raises(AirflowException):
             operator.execute(context=MagicMock())
 
@@ -120,9 +150,9 @@ class TestTableCheckOperator:
         "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
     }
 
-    def _construct_operator(self, monkeypatch, checks, return_df):
+    def _construct_operator(self, monkeypatch, checks, records):
         def get_records(*arg):
-            return return_df
+            return records
 
         operator = SQLTableCheckOperator(task_id="test_task", table="test_table", checks=checks)
         monkeypatch.setattr(operator, "get_db_hook", _get_mock_db_hook)
@@ -263,7 +293,7 @@ def setUp(self):
     def test_execute_no_records(self, mock_get_db_hook):
         mock_get_db_hook.return_value.get_first.return_value = []
 
-        with pytest.raises(AirflowException, match=r"The query returned None"):
+        with pytest.raises(AirflowException, match=r"The following query returned zero rows: sql"):
             self._operator.execute({})
 
     @mock.patch.object(SQLCheckOperator, "get_db_hook")

From 3e8dcd9fa366bbb55b96670b1a89f55ae252ff2c Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 16:40:25 -0400
Subject: [PATCH 23/63] Fix _failed_checks() to match update in parent
 operators

---
 .../providers/google/cloud/operators/bigquery.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index e1c21fe329b03..82749bb57edf6 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -38,7 +38,6 @@
     SQLIntervalCheckOperator,
     SQLTableCheckOperator,
     SQLValueCheckOperator,
-    _get_failed_checks,
     _parse_boolean,
     _raise_exception,
 )
@@ -624,7 +623,14 @@ def execute(self, context=None):
                 self.column_mapping[column][check], result, tolerance
             )
 
-            failed_tests.extend(_get_failed_checks(self.column_mapping[column], column))
+        for col, checks in self.column_mapping.items():
+            failed_tests.extend(
+                [
+                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+                    for check, check_values in checks.items()
+                    if not check_values["success"]
+                ]
+            )
         if failed_tests:
             exception_string = f"""
                 Test failed.\nResults:\n{records!s}\n
@@ -722,7 +728,11 @@ def execute(self, context=None):
             result = row[1].get("check_result")
             self.checks[check]["success"] = _parse_boolean(str(result))
 
-        failed_tests = _get_failed_checks(self.checks)
+        failed_tests = [
+            f"\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for check, check_values in self.checks.items()
+            if not check_values["success"]
+        ]
         if failed_tests:
             exception_string = f"""
                 Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n

From 0ce4b6f6a34661ed6d9852f3864450557e93973d Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Fri, 30 Sep 2022 09:33:39 -0400
Subject: [PATCH 24/63] Insert quotes around check_name table op's sql template
 to fix sql error

---
 airflow/providers/common/sql/operators/sql.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index b8d4b05679918..8365c86bb57ae 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -424,7 +424,7 @@ class SQLTableCheckOperator(BaseSQLOperator):
     template_fields = ("partition_clause",)
 
     sql_check_template = """
-        SELECT {check_name} AS check_name, MIN({check_name}) AS check_result
+        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
     """
 

From bf54c24a6e9fea636920716f26cc4ca4979f60d8 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 4 Oct 2022 11:27:07 -0400
Subject: [PATCH 25/63] Remove unnecessary list comprehension

---
 airflow/providers/common/sql/operators/sql.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 8365c86bb57ae..442ae1c833db5 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -242,7 +242,6 @@ def __init__(
         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(
@@ -251,7 +250,7 @@ def __init__(
                         table=self.table,
                         column=column,
                     )
-                    for check in checks_list
+                    for check in checks
                 ]
             )
         self.partition_clause = partition_clause

From c3e64eafa72f0e5d760ac1abe2cc08fac58bec93 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 6 Oct 2022 13:32:56 -0400
Subject: [PATCH 26/63] Added assertions in existing column and table tests

---
 tests/providers/common/sql/operators/test_sql.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 1c118f4d733d4..8fb7bf2c05e22 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -92,6 +92,10 @@ def test_pass_all_checks_exact_check(self, monkeypatch):
         ]
         operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
+        assert [
+            operator.column_mapping["X"][check]["success"] is True
+            for check in [*operator.column_mapping["X"]]
+        ]
 
     def test_max_less_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
@@ -129,6 +133,10 @@ def test_pass_all_checks_inexact_check(self, monkeypatch):
         ]
         operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
+        assert [
+            operator.column_mapping["X"][check]["success"] is True
+            for check in [*operator.column_mapping["X"]]
+        ]
 
     def test_fail_all_checks_check(self, monkeypatch):
         records = [
@@ -200,6 +208,7 @@ def test_pass_all_checks_check(self, monkeypatch):
         records = [('row_count_check', 1), ('column_sum_check', 'y')]
         operator = self._construct_operator(monkeypatch, self.checks, records)
         operator.execute(context=MagicMock())
+        assert [operator.checks[check]["success"] is True for check in operator.checks.keys()]
 
     def test_fail_all_checks_check(self, monkeypatch):
         records = [('row_count_check', 0), ('column_sum_check', 'n')]

From b1262e66228aa50c14c63f63a4047c5806cd4f59 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 6 Oct 2022 16:30:00 -0400
Subject: [PATCH 27/63] Add new tests for TableCheckOperator

---
 .../common/sql/operators/test_sql.py          | 46 +++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 8fb7bf2c05e22..fa0bd523d0601 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -158,6 +158,30 @@ class TestTableCheckOperator:
         "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
     }
 
+    correct_generate_sql_query_no_partitions = """
+    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
+    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    UNION ALL
+    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
+    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table ) AS sq
+    """
+
+    correct_generate_sql_query_with_partition = """
+    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
+    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10) AS sq
+    UNION ALL
+    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
+    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
+    """
+
+    correct_generate_sql_query_with_partition_and_where = """
+    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
+    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10 AND id = 100) AS sq
+    UNION ALL
+    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
+    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
+    """
+
     def _construct_operator(self, monkeypatch, checks, records):
         def get_records(*arg):
             return records
@@ -216,6 +240,28 @@ def test_fail_all_checks_check(self, monkeypatch):
         with pytest.raises(AirflowException):
             operator.execute(context=MagicMock())
 
+    def test_generate_sql_query_no_partitions(self, monkeypatch):
+        operator = self._construct_operator(monkeypatch, self.checks, ())
+        assert (
+            operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_no_partitions.lstrip()
+        )
+
+    def test_generate_sql_query_with_partitions(self, monkeypatch):
+        operator = self._construct_operator(monkeypatch, self.checks, ())
+        operator.partition_clause = "col_a > 10"
+        assert (
+            operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_partition.lstrip()
+        )
+
+    def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
+        operator = self._construct_operator(monkeypatch, self.checks, ())
+        operator.partition_clause = "col_a > 10"
+        self.checks["row_count_check"]["where"] = "id = 100"
+        assert (
+            operator._generate_sql_query().lstrip()
+            == self.correct_generate_sql_query_with_partition_and_where.lstrip()
+        )
+
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)
 INTERVAL = datetime.timedelta(hours=12)

From 1a5233576ac87241a2ca2d8405c46456e44d4e78 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Fri, 7 Oct 2022 10:26:08 -0400
Subject: [PATCH 28/63] Update operator logic and tests

Adds "where" option in checks dictionaries for column and table
operators, which may be renamed. This allows for check-level
partitioning, whereas the partition_clause param will always be
for all checks. New tests are added for this addition.
---
 airflow/providers/common/sql/operators/sql.py | 92 ++++++++++++-------
 .../common/sql/operators/test_sql.py          | 47 +++++++---
 2 files changed, 91 insertions(+), 48 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 442ae1c833db5..1f117025c2f4c 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -238,27 +238,16 @@ def __init__(
 
         self.table = table
         self.column_mapping = column_mapping
-        checks_sql = ""
+        self.partition_clause = partition_clause
+
+        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_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
-                ]
-            )
-        self.partition_clause = partition_clause
-        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}
-        """
+            checks_sql.append(self._generate_sql_query(column, checks))
+        checks_sql = " UNION ALL ".join(checks_sql)
+
+        self.sql = f"SELECT col_name, check_type, check_result FROM ({checks_sql}) AS check_columns"
 
     def execute(self, context: Context):
         failed_tests = []
@@ -296,6 +285,31 @@ def execute(self, context: Context):
 
         self.log.info("All tests have passed")
 
+    def _generate_sql_query(self, column, checks):
+        def _generate_partition_clause(check):
+            if self.partition_clause and "where" not in checks[check]:
+                return " WHERE " + self.partition_clause
+            elif not self.partition_clause and "where" in checks[check]:
+                return " WHERE " + checks[check]["where"]
+            elif self.partition_clause and "where" in checks[check]:
+                return " WHERE " + self.partition_clause + " AND " + checks[check]["where"]
+            else:
+                return ""
+
+        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,
+                    partition_clause=_generate_partition_clause(check),
+                )
+                for check in checks
+            ]
+        )
+        return checks_sql
+
     def _get_match(self, check_values, record, tolerance=None) -> bool:
         match_boolean = True
         if "geq_to" in check_values:
@@ -423,8 +437,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} {partition_clause}) AS sq
     """
 
     def __init__(
@@ -442,19 +456,7 @@ def __init__(
         self.table = table
         self.checks = checks
         self.partition_clause = partition_clause
-        checks_sql = " UNION ALL ".join(
-            [
-                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()
-            ]
-        )
-        partition_clause_statement = f"WHERE {self.partition_clause}" if self.partition_clause else ""
-        self.sql = f"""
-            SELECT check_name, check_result FROM ({checks_sql})
-            AS check_table {partition_clause_statement}
-        """
+        self.sql = f"SELECT check_name, check_result FROM ({self._generate_sql_query()}) AS check_table"
 
     def execute(self, context: Context):
         hook = self.get_db_hook()
@@ -484,6 +486,30 @@ def execute(self, context: Context):
 
         self.log.info("All tests have passed")
 
+    def _generate_sql_query(self):
+        def _generate_partition_clause(check_name):
+            if self.partition_clause and "where" not in self.checks[check_name]:
+                return "WHERE " + self.partition_clause
+            elif not self.partition_clause and "where" in self.checks[check_name]:
+                return "WHERE " + self.checks[check_name]["where"]
+            elif self.partition_clause and "where" in self.checks[check_name]:
+                return "WHERE " + self.partition_clause + " AND " + self.checks[check_name]["where"]
+            else:
+                return ""
+
+        checks_sql = "UNION ALL".join(
+            [
+                self.sql_check_template.format(
+                    check_statement=value["check_statement"],
+                    check_name=check_name,
+                    table=self.table,
+                    partition_clause=_generate_partition_clause(check_name),
+                )
+                for check_name, value in self.checks.items()
+            ]
+        )
+        return checks_sql
+
 
 class SQLCheckOperator(BaseSQLOperator):
     """
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index fa0bd523d0601..9ade60c989dcb 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -153,34 +153,44 @@ def test_fail_all_checks_check(self, monkeypatch):
 
 class TestTableCheckOperator:
 
+    count_check = "COUNT(*) == 1000"
+    sum_check = "col_a + col_b < col_c"
     checks = {
-        "row_count_check": {"check_statement": "COUNT(*) == 1000"},
-        "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
+        "row_count_check": {"check_statement": f"{count_check}"},
+        "column_sum_check": {"check_statement": f"{sum_check}"},
     }
 
-    correct_generate_sql_query_no_partitions = """
+    correct_generate_sql_query_no_partitions = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table ) AS sq
     """
 
-    correct_generate_sql_query_with_partition = """
+    correct_generate_sql_query_with_partition = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
-    """
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
+    """  # noqa 501
 
-    correct_generate_sql_query_with_partition_and_where = """
+    correct_generate_sql_query_with_partition_and_where = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10 AND id = 100) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10 AND id = 100) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
-    """
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
+    """  # noqa 501
+
+    correct_generate_sql_query_with_where = f"""
+    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    UNION ALL
+    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE id = 100) AS sq
+    """  # noqa 501
 
     def _construct_operator(self, monkeypatch, checks, records):
         def get_records(*arg):
@@ -254,14 +264,21 @@ def test_generate_sql_query_with_partitions(self, monkeypatch):
         )
 
     def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
-        operator = self._construct_operator(monkeypatch, self.checks, ())
+        checks = self.checks
+        checks["row_count_check"]["where"] = "id = 100"
+        operator = self._construct_operator(monkeypatch, checks, ())
         operator.partition_clause = "col_a > 10"
-        self.checks["row_count_check"]["where"] = "id = 100"
         assert (
             operator._generate_sql_query().lstrip()
             == self.correct_generate_sql_query_with_partition_and_where.lstrip()
         )
 
+    def test_generate_sql_query_with_check_partition(self, monkeypatch):
+        checks = self.checks
+        checks["column_sum_check"]["where"] = "id = 100"
+        operator = self._construct_operator(monkeypatch, checks, ())
+        assert operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_where.lstrip()
+
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)
 INTERVAL = datetime.timedelta(hours=12)

From 46ca885577d1a5da23aba840ccb1b42a19e62927 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Fri, 7 Oct 2022 11:35:16 -0400
Subject: [PATCH 29/63] Add testing for column check operator and line edits

Cleans up operator and adds testing for new generator function.
---
 airflow/providers/common/sql/operators/sql.py |  19 +--
 .../common/sql/operators/test_sql.py          | 119 +++++++++++++++---
 2 files changed, 112 insertions(+), 26 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 1f117025c2f4c..01bb1e6bb0f53 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -213,7 +213,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
 
     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
+        FROM (SELECT {check_statement} AS {column}_{check} FROM {table} {partition_clause}) AS sq
     """
 
     column_checks = {
@@ -240,12 +240,12 @@ def __init__(
         self.column_mapping = column_mapping
         self.partition_clause = partition_clause
 
-        checks_sql = []
+        checks_sql_list = []
         for column, checks in self.column_mapping.items():
             for check, check_values in checks.items():
                 self._column_mapping_validation(check, check_values)
-            checks_sql.append(self._generate_sql_query(column, checks))
-        checks_sql = " UNION ALL ".join(checks_sql)
+            checks_sql_list.append(self._generate_sql_query(column, checks))
+        checks_sql = "UNION ALL".join(checks_sql_list)
 
         self.sql = f"SELECT col_name, check_type, check_result FROM ({checks_sql}) AS check_columns"
 
@@ -288,15 +288,15 @@ def execute(self, context: Context):
     def _generate_sql_query(self, column, checks):
         def _generate_partition_clause(check):
             if self.partition_clause and "where" not in checks[check]:
-                return " WHERE " + self.partition_clause
+                return "WHERE " + self.partition_clause
             elif not self.partition_clause and "where" in checks[check]:
-                return " WHERE " + checks[check]["where"]
+                return "WHERE " + checks[check]["where"]
             elif self.partition_clause and "where" in checks[check]:
-                return " WHERE " + self.partition_clause + " AND " + checks[check]["where"]
+                return "WHERE " + self.partition_clause + " AND " + checks[check]["where"]
             else:
                 return ""
 
-        checks_sql = " UNION ALL ".join(
+        checks_sql = "UNION ALL".join(
             [
                 self.sql_check_template.format(
                     check_statement=self.column_checks[check].format(column=column),
@@ -438,7 +438,8 @@ class SQLTableCheckOperator(BaseSQLOperator):
 
     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} {partition_clause}) AS sq
+    FROM (SELECT CASE WHEN {check_statement} THEN 1 ELSE 0 END AS {check_name}
+          FROM {table} {partition_clause}) AS sq
     """
 
     def __init__(
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 9ade60c989dcb..584de8ab95a72 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -65,8 +65,47 @@ class TestColumnCheckOperator:
         }
     }
 
+    short_valid_column_mapping = {
+        "X": {
+            "null_check": {"equal_to": 0},
+            "distinct_check": {"equal_to": 10, "tolerance": 0.1},
+        }
+    }
+
     invalid_column_mapping = {"Y": {"invalid_check_name": {"expectation": 5}}}
 
+    correct_generate_sql_query_no_partitions = """
+        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
+        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table ) AS sq
+    UNION ALL
+        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
+        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table ) AS sq
+    """
+
+    correct_generate_sql_query_with_partition = """
+        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
+        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table WHERE Y > 1) AS sq
+    UNION ALL
+        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
+        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table WHERE Y > 1) AS sq
+    """  # noqa 501
+
+    correct_generate_sql_query_with_partition_and_where = """
+        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
+        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table WHERE Y > 1 AND Z < 100) AS sq
+    UNION ALL
+        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
+        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table WHERE Y > 1) AS sq
+    """  # noqa 501
+
+    correct_generate_sql_query_with_where = """
+        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
+        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table ) AS sq
+    UNION ALL
+        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
+        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table WHERE Z < 100) AS sq
+    """  # 501
+
     def _construct_operator(self, monkeypatch, column_mapping, records):
         def get_records(*arg):
             return records
@@ -150,6 +189,44 @@ def test_fail_all_checks_check(self, monkeypatch):
         with pytest.raises(AirflowException):
             operator.execute(context=MagicMock())
 
+    def test_generate_sql_query_no_partitions(self, monkeypatch):
+        checks = self.short_valid_column_mapping["X"]
+        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
+        assert (
+            operator._generate_sql_query("X", checks).lstrip()
+            == self.correct_generate_sql_query_no_partitions.lstrip()
+        )
+
+    def test_generate_sql_query_with_partitions(self, monkeypatch):
+        checks = self.short_valid_column_mapping["X"]
+        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
+        operator.partition_clause = "Y > 1"
+        assert (
+            operator._generate_sql_query("X", checks).lstrip()
+            == self.correct_generate_sql_query_with_partition.lstrip()
+        )
+
+    def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
+        self.short_valid_column_mapping["X"]["null_check"]["where"] = "Z < 100"
+        checks = self.short_valid_column_mapping["X"]
+        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
+        operator.partition_clause = "Y > 1"
+        assert (
+            operator._generate_sql_query("X", checks).lstrip()
+            == self.correct_generate_sql_query_with_partition_and_where.lstrip()
+        )
+        del self.short_valid_column_mapping["X"]["null_check"]["where"]
+
+    def test_generate_sql_query_with_check_partition(self, monkeypatch):
+        self.short_valid_column_mapping["X"]["distinct_check"]["where"] = "Z < 100"
+        checks = self.short_valid_column_mapping["X"]
+        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
+        assert (
+            operator._generate_sql_query("X", checks).lstrip()
+            == self.correct_generate_sql_query_with_where.lstrip()
+        )
+        del self.short_valid_column_mapping["X"]["distinct_check"]["where"]
+
 
 class TestTableCheckOperator:
 
@@ -162,35 +239,43 @@ class TestTableCheckOperator:
 
     correct_generate_sql_query_no_partitions = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
+          FROM test_table ) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
+          FROM test_table ) AS sq
     """
 
     correct_generate_sql_query_with_partition = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
+          FROM test_table WHERE col_a > 10) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
-    """  # noqa 501
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
+          FROM test_table WHERE col_a > 10) AS sq
+    """
 
     correct_generate_sql_query_with_partition_and_where = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10 AND id = 100) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
+          FROM test_table WHERE col_a > 10 AND id = 100) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
-    """  # noqa 501
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
+          FROM test_table WHERE col_a > 10) AS sq
+    """
 
     correct_generate_sql_query_with_where = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
+          FROM test_table ) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE id = 100) AS sq
-    """  # noqa 501
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
+          FROM test_table WHERE id = 100) AS sq
+    """
 
     def _construct_operator(self, monkeypatch, checks, records):
         def get_records(*arg):
@@ -264,20 +349,20 @@ def test_generate_sql_query_with_partitions(self, monkeypatch):
         )
 
     def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
-        checks = self.checks
-        checks["row_count_check"]["where"] = "id = 100"
-        operator = self._construct_operator(monkeypatch, checks, ())
+        self.checks["row_count_check"]["where"] = "id = 100"
+        operator = self._construct_operator(monkeypatch, self.checks, ())
         operator.partition_clause = "col_a > 10"
         assert (
             operator._generate_sql_query().lstrip()
             == self.correct_generate_sql_query_with_partition_and_where.lstrip()
         )
+        del self.checks["row_count_check"]["where"]
 
     def test_generate_sql_query_with_check_partition(self, monkeypatch):
-        checks = self.checks
-        checks["column_sum_check"]["where"] = "id = 100"
-        operator = self._construct_operator(monkeypatch, checks, ())
+        self.checks["column_sum_check"]["where"] = "id = 100"
+        operator = self._construct_operator(monkeypatch, self.checks, ())
         assert operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_where.lstrip()
+        del self.checks["column_sum_check"]["where"]
 
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)

From 9014f6505fc7bc06c34d94dffe5962e502a66b9f Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 11 Oct 2022 09:51:00 -0400
Subject: [PATCH 30/63] Change name 'where' to 'partition_clause' in check
 dictionaries

---
 airflow/providers/common/sql/operators/sql.py | 31 ++++++++++---------
 .../common/sql/operators/test_sql.py          | 16 +++++-----
 2 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 01bb1e6bb0f53..7c57ab9549533 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -191,6 +191,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
                     "tolerance": 0.2,
                 },
                 "max": {"less_than": 1000, "geq_to": 10, "tolerance": 0.01},
+                "partition_clause": "foreign_key IS NOT NULL",
             }
         }
 
@@ -287,12 +288,12 @@ def execute(self, context: Context):
 
     def _generate_sql_query(self, column, checks):
         def _generate_partition_clause(check):
-            if self.partition_clause and "where" not in checks[check]:
+            if self.partition_clause and "partition_clause" not in checks[check]:
                 return "WHERE " + self.partition_clause
-            elif not self.partition_clause and "where" in checks[check]:
-                return "WHERE " + checks[check]["where"]
-            elif self.partition_clause and "where" in checks[check]:
-                return "WHERE " + self.partition_clause + " AND " + checks[check]["where"]
+            elif not self.partition_clause and "partition_clause" in checks[check]:
+                return "WHERE " + checks[check]["partition_clause"]
+            elif self.partition_clause and "partition_clause" in checks[check]:
+                return "WHERE " + self.partition_clause + " AND " + checks[check]["partition_clause"]
             else:
                 return ""
 
@@ -409,13 +410,14 @@ class SQLTableCheckOperator(BaseSQLOperator):
     Checks should be written to return a boolean result.
 
     :param table: the table to run checks on
-    :param checks: the dictionary of checks, e.g.:
+    :param checks: the dictionary of checks, where check names are followed by a dictionary containing at
+        least a check statement, and optionally a partition clause, e.g.:
 
     .. code-block:: python
 
         {
             "row_count_check": {"check_statement": "COUNT(*) = 1000"},
-            "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
+            "column_sum_check": {"check_statement": "col_a + col_b < col_c", "partition_clause": "col_a IS NOT NULL"},
         }
 
 
@@ -489,16 +491,18 @@ def execute(self, context: Context):
 
     def _generate_sql_query(self):
         def _generate_partition_clause(check_name):
-            if self.partition_clause and "where" not in self.checks[check_name]:
+            if self.partition_clause and "partition_clause" not in self.checks[check_name]:
                 return "WHERE " + self.partition_clause
-            elif not self.partition_clause and "where" in self.checks[check_name]:
-                return "WHERE " + self.checks[check_name]["where"]
-            elif self.partition_clause and "where" in self.checks[check_name]:
-                return "WHERE " + self.partition_clause + " AND " + self.checks[check_name]["where"]
+            elif not self.partition_clause and "partition_clause" in self.checks[check_name]:
+                return "WHERE " + self.checks[check_name]["partition_clause"]
+            elif self.partition_clause and "partition_clause" in self.checks[check_name]:
+                return (
+                    "WHERE " + self.partition_clause + " AND " + self.checks[check_name]["partition_clause"]
+                )
             else:
                 return ""
 
-        checks_sql = "UNION ALL".join(
+        return "UNION ALL".join(
             [
                 self.sql_check_template.format(
                     check_statement=value["check_statement"],
@@ -509,7 +513,6 @@ def _generate_partition_clause(check_name):
                 for check_name, value in self.checks.items()
             ]
         )
-        return checks_sql
 
 
 class SQLCheckOperator(BaseSQLOperator):
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 584de8ab95a72..0f0902344201f 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -207,7 +207,7 @@ def test_generate_sql_query_with_partitions(self, monkeypatch):
         )
 
     def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
-        self.short_valid_column_mapping["X"]["null_check"]["where"] = "Z < 100"
+        self.short_valid_column_mapping["X"]["null_check"]["partition_clause"] = "Z < 100"
         checks = self.short_valid_column_mapping["X"]
         operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
         operator.partition_clause = "Y > 1"
@@ -215,17 +215,17 @@ def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatc
             operator._generate_sql_query("X", checks).lstrip()
             == self.correct_generate_sql_query_with_partition_and_where.lstrip()
         )
-        del self.short_valid_column_mapping["X"]["null_check"]["where"]
+        del self.short_valid_column_mapping["X"]["null_check"]["partition_clause"]
 
     def test_generate_sql_query_with_check_partition(self, monkeypatch):
-        self.short_valid_column_mapping["X"]["distinct_check"]["where"] = "Z < 100"
+        self.short_valid_column_mapping["X"]["distinct_check"]["partition_clause"] = "Z < 100"
         checks = self.short_valid_column_mapping["X"]
         operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
         assert (
             operator._generate_sql_query("X", checks).lstrip()
             == self.correct_generate_sql_query_with_where.lstrip()
         )
-        del self.short_valid_column_mapping["X"]["distinct_check"]["where"]
+        del self.short_valid_column_mapping["X"]["distinct_check"]["partition_clause"]
 
 
 class TestTableCheckOperator:
@@ -349,20 +349,20 @@ def test_generate_sql_query_with_partitions(self, monkeypatch):
         )
 
     def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
-        self.checks["row_count_check"]["where"] = "id = 100"
+        self.checks["row_count_check"]["partition_clause"] = "id = 100"
         operator = self._construct_operator(monkeypatch, self.checks, ())
         operator.partition_clause = "col_a > 10"
         assert (
             operator._generate_sql_query().lstrip()
             == self.correct_generate_sql_query_with_partition_and_where.lstrip()
         )
-        del self.checks["row_count_check"]["where"]
+        del self.checks["row_count_check"]["partition_clause"]
 
     def test_generate_sql_query_with_check_partition(self, monkeypatch):
-        self.checks["column_sum_check"]["where"] = "id = 100"
+        self.checks["column_sum_check"]["partition_clause"] = "id = 100"
         operator = self._construct_operator(monkeypatch, self.checks, ())
         assert operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_where.lstrip()
-        del self.checks["column_sum_check"]["where"]
+        del self.checks["column_sum_check"]["partition_clause"]
 
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)

From ee5f516339bd0f14d6cc74d2d12b264d3898cce2 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 11 Oct 2022 12:35:59 -0400
Subject: [PATCH 31/63] Update docs and use f strings

---
 airflow/providers/common/sql/operators/sql.py | 17 ++++++-----
 .../operators.rst                             | 28 +++++++++++--------
 2 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 7c57ab9549533..6d2976f5b981b 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -174,6 +174,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
     - geq_to: value that results should be greater than or equal to
     - leq_to: value that results should be less than or equal to
     - tolerance: the percentage that the result may be off from the expected value
+    - partition_clause: an extra clause passed into a WHERE statement to partition data
 
     :param table: the table to run checks on
     :param column_mapping: the dictionary of columns and their associated checks, e.g.
@@ -184,6 +185,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
             "col_name": {
                 "null_check": {
                     "equal_to": 0,
+                    "partition_clause": "foreign_key IS NOT NULL",
                 },
                 "min": {
                     "greater_than": 5,
@@ -191,7 +193,6 @@ class SQLColumnCheckOperator(BaseSQLOperator):
                     "tolerance": 0.2,
                 },
                 "max": {"less_than": 1000, "geq_to": 10, "tolerance": 0.01},
-                "partition_clause": "foreign_key IS NOT NULL",
             }
         }
 
@@ -289,11 +290,11 @@ def execute(self, context: Context):
     def _generate_sql_query(self, column, checks):
         def _generate_partition_clause(check):
             if self.partition_clause and "partition_clause" not in checks[check]:
-                return "WHERE " + self.partition_clause
+                return f"WHERE {self.partition_clause}"
             elif not self.partition_clause and "partition_clause" in checks[check]:
-                return "WHERE " + checks[check]["partition_clause"]
+                return f"WHERE {checks[check]['partition_clause']}"
             elif self.partition_clause and "partition_clause" in checks[check]:
-                return "WHERE " + self.partition_clause + " AND " + checks[check]["partition_clause"]
+                return f"WHERE {self.partition_clause} AND {checks[check]['partition_clause']}"
             else:
                 return ""
 
@@ -492,13 +493,11 @@ def execute(self, context: Context):
     def _generate_sql_query(self):
         def _generate_partition_clause(check_name):
             if self.partition_clause and "partition_clause" not in self.checks[check_name]:
-                return "WHERE " + self.partition_clause
+                return f"WHERE {self.partition_clause}"
             elif not self.partition_clause and "partition_clause" in self.checks[check_name]:
-                return "WHERE " + self.checks[check_name]["partition_clause"]
+                return f"WHERE {self.checks[check_name]['partition_clause']}"
             elif self.partition_clause and "partition_clause" in self.checks[check_name]:
-                return (
-                    "WHERE " + self.partition_clause + " AND " + self.checks[check_name]["partition_clause"]
-                )
+                return f"WHERE {self.partition_clause} AND {self.checks[check_name]['partition_clause']}"
             else:
                 return ""
 
diff --git a/docs/apache-airflow-providers-common-sql/operators.rst b/docs/apache-airflow-providers-common-sql/operators.rst
index 402197fb5204f..0283ac0ecaea9 100644
--- a/docs/apache-airflow-providers-common-sql/operators.rst
+++ b/docs/apache-airflow-providers-common-sql/operators.rst
@@ -28,16 +28,14 @@ Check SQL Table Columns
 
 Use the :class:`~airflow.providers.common.sql.operators.sql.SQLColumnCheckOperator` to run data quality
 checks against columns of a given table. As well as a connection ID and table, a column_mapping
-describing the relationship between columns and tests to run must be supplied. An example column
-mapping is a set of three nested dictionaries and looks like:
+describing the relationship between columns and tests to run must be supplied. An example column mapping
+is a set of three nested dictionaries and looks like:
 
 .. code-block:: python
 
         column_mapping = {
             "col_name": {
-                "null_check": {
-                    "equal_to": 0,
-                },
+                "null_check": {"equal_to": 0, "partition_clause": "other_col LIKE 'this'"},
                 "min": {
                     "greater_than": 5,
                     "leq_to": 10,
@@ -56,8 +54,8 @@ The valid checks are:
 - min: checks the minimum value in the column
 - max: checks the maximum value in the column
 
-Each entry in the check's dictionary is either a condition for success of the check or the tolerance. The
-conditions for success are:
+Each entry in the check's dictionary is either a condition for success of the check, the tolerance,
+or a partition clause. The conditions for success are:
 
 - greater_than
 - geq_to
@@ -69,7 +67,9 @@ When specifying conditions, equal_to is not compatible with other conditions. Bo
 bound condition may be specified in the same check. The tolerance is a percentage that the result may
 be out of bounds but still considered successful.
 
-
+The partition clauses may be given at the operator level as a parameter where it partitions all checks,
+at the column level in the column mapping where it partitions all checks for that column, or at the
+check level for a column where it partitions just that check.
 
 The below example demonstrates how to instantiate the SQLColumnCheckOperator task.
 
@@ -96,14 +96,20 @@ checks argument is a set of two nested dictionaries and looks like:
                 "row_count_check": {
                     "check_statement": "COUNT(*) = 1000",
                 },
-                "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
+                "column_sum_check": {
+                    "check_statement": "col_a + col_b < col_c",
+                    "partition_clause": "col_a IS NOT NULL",
+                },
             },
         )
 
 The first set of keys are the check names, which are referenced in the templated query the operator builds.
-The dictionary key under the check name must be check_statement, with the value a SQL statement that
+A dictionary key under the check name must include check_statement and the value a SQL statement that
 resolves to a boolean (this can be any string or int that resolves to a boolean in
-airflow.operators.sql.parse_boolean).
+airflow.operators.sql.parse_boolean). The other possible key to supply is partition_clause, which is a
+check level statement that will partition the data in the table using a WHERE clause for that check.
+This statement is compatible with the parameter partition_clause, where the latter filters across all
+checks.
 
 The below example demonstrates how to instantiate the SQLTableCheckOperator task.
 

From bdccba7321eb069139c1a97f582c417684109078 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 11 Oct 2022 12:37:56 -0400
Subject: [PATCH 32/63] Edit operator docstring

---
 airflow/providers/common/sql/operators/sql.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 6d2976f5b981b..5d96723c8ee32 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -418,7 +418,8 @@ class SQLTableCheckOperator(BaseSQLOperator):
 
         {
             "row_count_check": {"check_statement": "COUNT(*) = 1000"},
-            "column_sum_check": {"check_statement": "col_a + col_b < col_c", "partition_clause": "col_a IS NOT NULL"},
+            "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
+            "third_check": {"check_statement": "MIN(col) = 1", "partition_clause": "col IS NOT NULL"},
         }
 
 

From fe8ba7071544ebc13652d1068639ad21d877c2fb Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 12 Oct 2022 09:11:38 -0400
Subject: [PATCH 33/63] Move _raise_exception to base class to simplify method

---
 airflow/providers/common/sql/operators/sql.py | 41 ++++++++-----------
 1 file changed, 18 insertions(+), 23 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 5d96723c8ee32..b1c572d71f186 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -53,12 +53,6 @@ def _parse_boolean(val: str) -> str | bool:
     raise ValueError(f"{val!r} is not a boolean-like string value")
 
 
-def _raise_exception(exception_string: str, retry_on_failure: bool) -> NoReturn:
-    if retry_on_failure:
-        raise AirflowException(exception_string)
-    raise AirflowFailException(exception_string)
-
-
 _PROVIDERS_MATCHER = re.compile(r'airflow\.providers\.(.*)\.hooks.*')
 
 _MIN_SUPPORTED_PROVIDERS_VERSION = {
@@ -162,6 +156,11 @@ def get_db_hook(self) -> DbApiHook:
         """
         return self._hook
 
+    def _raise_exception(self, exception_string: str) -> NoReturn:
+        if self.retry_on_failure:
+            raise AirflowException(exception_string)
+        raise AirflowFailException(exception_string)
+
 
 class SQLColumnCheckOperator(BaseSQLOperator):
     """
@@ -257,7 +256,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
 
         self.log.info("Record: %s", records)
 
@@ -283,7 +282,7 @@ def execute(self, context: Context):
                 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._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
 
@@ -468,7 +467,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
 
         self.log.info("Record:\n%s", records)
 
@@ -487,7 +486,7 @@ def execute(self, context: Context):
                 The following tests have failed:
                 \n{', '.join(failed_tests)}
             """
-            _raise_exception(exception_string, self.retry_on_failure)
+            self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
 
@@ -568,11 +567,9 @@ def execute(self, context: Context):
 
         self.log.info("Record: %s", records)
         if not records:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
         elif not all(bool(r) for r in records):
-            _raise_exception(
-                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}", self.retry_on_failure
-            )
+            self._raise_exception(f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}")
 
         self.log.info("Success.")
 
@@ -620,7 +617,7 @@ def execute(self, context: Context):
         records = self.get_db_hook().get_first(self.sql)
 
         if not records:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
 
         pass_value_conv = _convert_to_float_if_possible(self.pass_value)
         is_numeric_value_check = isinstance(pass_value_conv, float)
@@ -649,7 +646,7 @@ def execute(self, context: Context):
             tests = []
 
         if not all(tests):
-            _raise_exception(error_msg, self.retry_on_failure)
+            self._raise_exception(error_msg)
 
     def _to_float(self, records):
         return [float(record) for record in records]
@@ -746,9 +743,9 @@ def execute(self, context: Context):
         row1 = hook.get_first(self.sql1)
 
         if not row2:
-            _raise_exception(f"The following query returned zero rows: {self.sql2}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql2}")
         if not row1:
-            _raise_exception(f"The following query returned zero rows: {self.sql1}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql1}")
 
         current = dict(zip(self.metrics_sorted, row1))
         reference = dict(zip(self.metrics_sorted, row2))
@@ -801,9 +798,7 @@ def execute(self, context: Context):
                     ratios[k],
                     self.metrics_thresholds[k],
                 )
-            _raise_exception(
-                f"The following tests have failed:\n {', '.join(sorted(failed_tests))}", self.retry_on_failure
-            )
+            self._raise_exception(f"The following tests have failed:\n {', '.join(sorted(failed_tests))}")
 
         self.log.info("All tests have passed")
 
@@ -847,7 +842,7 @@ def execute(self, context: Context):
         hook = self.get_db_hook()
         result = hook.get_first(self.sql)[0]
         if not result:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
 
         if isinstance(self.min_threshold, float):
             lower_bound = self.min_threshold
@@ -882,7 +877,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_exception(error_msg, self.retry_on_failure)
+            self._raise_exception(error_msg)
 
         self.log.info("Test %s Successful.", self.task_id)
 

From 3e159c04e102e7c1518e88e94bff55520fff094e Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 12 Oct 2022 10:03:19 -0400
Subject: [PATCH 34/63] Updates from code review

---
 airflow/providers/common/sql/operators/sql.py | 69 ++++++++-----------
 1 file changed, 30 insertions(+), 39 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index b1c572d71f186..ebac975060396 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -251,7 +251,6 @@ def __init__(
         self.sql = f"SELECT col_name, check_type, check_result FROM ({checks_sql}) AS check_columns"
 
     def execute(self, context: Context):
-        failed_tests = []
         hook = self.get_db_hook()
         records = hook.get_records(self.sql)
 
@@ -260,8 +259,7 @@ def execute(self, context: Context):
 
         self.log.info("Record: %s", records)
 
-        for row in records:
-            column, check, result = row
+        for column, check, result in records:
             tolerance = self.column_mapping[column][check].get("tolerance")
 
             self.column_mapping[column][check]["result"] = result
@@ -269,19 +267,17 @@ def execute(self, context: Context):
                 self.column_mapping[column][check], result, tolerance
             )
 
-        for col, checks in self.column_mapping.items():
-            failed_tests.extend(
-                [
-                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
-                    for check, check_values in checks.items()
-                    if not check_values["success"]
-                ]
-            )
+        failed_tests = [
+            f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for col, checks in self.column_mapping.items()
+            for check, check_values in checks.items()
+            if not check_values["success"]
+        ]
         if failed_tests:
-            exception_string = f"""
-                Test failed.\nResults:\n{records!s}\n
-                The following tests have failed:
-                \n{''.join(failed_tests)}"""
+            exception_string = (
+                f"Test failed.\nResults:\n{records!s}\n"
+                f"The following tests have failed:\n{''.join(failed_tests)}"
+            )
             self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
@@ -298,16 +294,14 @@ def _generate_partition_clause(check):
                 return ""
 
         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,
-                    partition_clause=_generate_partition_clause(check),
-                )
-                for check in checks
-            ]
+            self.sql_check_template.format(
+                check_statement=self.column_checks[check].format(column=column),
+                check=check,
+                table=self.table,
+                column=column,
+                partition_clause=_generate_partition_clause(check),
+            )
+            for check in checks
         )
         return checks_sql
 
@@ -481,11 +475,10 @@ def execute(self, context: Context):
             if not check_values["success"]
         ]
         if 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)}
-            """
+            exception_string = (
+                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n"
+                f"The following tests have failed:\n{', '.join(failed_tests)}"
+            )
             self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
@@ -502,15 +495,13 @@ def _generate_partition_clause(check_name):
                 return ""
 
         return "UNION ALL".join(
-            [
-                self.sql_check_template.format(
-                    check_statement=value["check_statement"],
-                    check_name=check_name,
-                    table=self.table,
-                    partition_clause=_generate_partition_clause(check_name),
-                )
-                for check_name, value in self.checks.items()
-            ]
+            self.sql_check_template.format(
+                check_statement=value["check_statement"],
+                check_name=check_name,
+                table=self.table,
+                partition_clause=_generate_partition_clause(check_name),
+            )
+            for check_name, value in self.checks.items()
         )
 
 

From 37c704f9af02a0cc151dfe82422eee9e555b0882 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 12 Oct 2022 16:23:28 -0400
Subject: [PATCH 35/63] Update BigQuery Check operators according to code
 review

---
 .../google/cloud/operators/bigquery.py        | 41 ++++++++-----------
 1 file changed, 18 insertions(+), 23 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 82749bb57edf6..66c5e5f9879dc 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -39,7 +39,6 @@
     SQLTableCheckOperator,
     SQLValueCheckOperator,
     _parse_boolean,
-    _raise_exception,
 )
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
@@ -264,9 +263,7 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
         if not records:
             raise AirflowException("The query returned empty results")
         elif not all(bool(r) for r in records):
-            _raise_exception(
-                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}", self.retry_on_failure
-            )
+            self._raise_exception(f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}")
         self.log.info("Record: %s", event["records"])
         self.log.info("Success.")
 
@@ -623,20 +620,19 @@ def execute(self, context=None):
                 self.column_mapping[column][check], result, tolerance
             )
 
-        for col, checks in self.column_mapping.items():
-            failed_tests.extend(
-                [
-                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
-                    for check, check_values in checks.items()
-                    if not check_values["success"]
-                ]
-            )
+        failed_tests(
+            f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for col, checks in self.column_mapping.items()
+            for check, check_values in checks.items()
+            if not check_values["success"]
+        )
         if 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)
+            exception_string = (
+                f"Test failed.\nResults:\n{records!s}\n"
+                f"The following tests have failed:"
+                f"\n{''.join(failed_tests)}"
+            )
+            self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
 
@@ -734,12 +730,11 @@ def execute(self, context=None):
             if not check_values["success"]
         ]
         if 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)
+            exception_string = (
+                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n"
+                f"The following tests have failed:\n{', '.join(failed_tests)}"
+            )
+            self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
 

From 6368a95d654e14cb302bf34de800584ba24ca36c Mon Sep 17 00:00:00 2001
From: Tzu-ping Chung <uranusjr@gmail.com>
Date: Thu, 13 Oct 2022 13:09:09 +0800
Subject: [PATCH 36/63] Rewrite data-building loop to generator

---
 airflow/providers/common/sql/operators/sql.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index ebac975060396..0a1682be17776 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -241,12 +241,13 @@ def __init__(
         self.column_mapping = column_mapping
         self.partition_clause = partition_clause
 
-        checks_sql_list = []
-        for column, checks in self.column_mapping.items():
-            for check, check_values in checks.items():
-                self._column_mapping_validation(check, check_values)
-            checks_sql_list.append(self._generate_sql_query(column, checks))
-        checks_sql = "UNION ALL".join(checks_sql_list)
+        def _build_checks_sql():
+            for column, checks in self.column_mapping.items():
+                for check, check_values in checks.items():
+                    self._column_mapping_validation(check, check_values)
+                yield self._generate_sql_query(column, checks)
+
+        checks_sql = "UNION ALL".join(_build_checks_sql())
 
         self.sql = f"SELECT col_name, check_type, check_result FROM ({checks_sql}) AS check_columns"
 

From 9367ab90bbed0703d80d3fce57f5f0a9878e6554 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Fri, 14 Oct 2022 11:28:09 -0400
Subject: [PATCH 37/63] Add new accept_none argument to column check operator

The new argument, defaulting to true, will convert Nones returned
from the query to 0s so numeric calculations can be performed
correctly. This allows empty tables to be handled as a row of zeroes.

Additional documentation is also supplied
---
 airflow/providers/common/sql/operators/sql.py         |  6 ++++++
 airflow/providers/google/cloud/operators/bigquery.py  | 11 ++++++++++-
 .../apache-airflow-providers-common-sql/operators.rst |  5 +++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 0a1682be17776..35a89a1e472fd 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -204,6 +204,8 @@ class SQLColumnCheckOperator(BaseSQLOperator):
 
     :param conn_id: the connection ID used to connect to the database
     :param database: name of database which overwrite the defined one in connection
+    :param accept_none: whether or not to accept None values returned by the query. If true, converts None
+        to 0.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -233,6 +235,7 @@ def __init__(
         partition_clause: str | None = None,
         conn_id: str | None = None,
         database: str | None = None,
+        accept_none: bool = True,
         **kwargs,
     ):
         super().__init__(conn_id=conn_id, database=database, **kwargs)
@@ -240,6 +243,7 @@ def __init__(
         self.table = table
         self.column_mapping = column_mapping
         self.partition_clause = partition_clause
+        self.accept_none = accept_none
 
         def _build_checks_sql():
             for column, checks in self.column_mapping.items():
@@ -307,6 +311,8 @@ def _generate_partition_clause(check):
         return checks_sql
 
     def _get_match(self, check_values, record, tolerance=None) -> bool:
+        if record is None and self.accept_none:
+            record = 0
         match_boolean = True
         if "geq_to" in check_values:
             if tolerance is not None:
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 66c5e5f9879dc..74bcfeef6122f 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -559,6 +559,8 @@ def __init__(
         table: str,
         column_mapping: dict,
         partition_clause: str | None = None,
+        database: str | None = None,
+        accept_none: bool = True,
         gcp_conn_id: str = "google_cloud_default",
         use_legacy_sql: bool = True,
         location: str | None = None,
@@ -567,11 +569,18 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(
-            table=table, column_mapping=column_mapping, partition_clause=partition_clause, **kwargs
+            table=table,
+            column_mapping=column_mapping,
+            partition_clause=partition_clause,
+            database=database,
+            accept_none=accept_none,
+            **kwargs,
         )
         self.table = table
         self.column_mapping = column_mapping
         self.partition_clause = partition_clause
+        self.database = database
+        self.accept_none = accept_none
         self.gcp_conn_id = gcp_conn_id
         self.use_legacy_sql = use_legacy_sql
         self.location = location
diff --git a/docs/apache-airflow-providers-common-sql/operators.rst b/docs/apache-airflow-providers-common-sql/operators.rst
index 0283ac0ecaea9..ae3c9072ad02f 100644
--- a/docs/apache-airflow-providers-common-sql/operators.rst
+++ b/docs/apache-airflow-providers-common-sql/operators.rst
@@ -71,6 +71,11 @@ The partition clauses may be given at the operator level as a parameter where it
 at the column level in the column mapping where it partitions all checks for that column, or at the
 check level for a column where it partitions just that check.
 
+A database may also be specified if not using the database from the supplied connection.
+
+The accept_none argument, true by default, will convert None values returned by the query to 0s, allowing
+empty tables to return valid integers.
+
 The below example demonstrates how to instantiate the SQLColumnCheckOperator task.
 
 .. exampleinclude:: /../../tests/system/providers/common/sql/example_sql_column_table_check.py

From 43a33c42f93fb8bff6e8246afe48e2b45f354dd6 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 27 Sep 2022 10:44:58 -0400
Subject: [PATCH 38/63] Fix BigQueryTableCheckOperator test

Signed-off-by: Benji Lampel <benjamin@astronomer.io>
---
 .../providers/google/cloud/bigquery/example_bigquery_queries.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py b/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py
index 789b5fb921229..93ef73fb44687 100644
--- a/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py
+++ b/tests/system/providers/google/cloud/bigquery/example_bigquery_queries.py
@@ -223,7 +223,7 @@
         table_check = BigQueryTableCheckOperator(
             task_id="table_check",
             table=f"{DATASET}.{TABLE_1}",
-            checks={"row_count_check": {"check_statement": {"COUNT(*) = 4"}}},
+            checks={"row_count_check": {"check_statement": "COUNT(*) = 4"}},
         )
         # [END howto_operator_bigquery_table_check]
 

From c761c0569dc052102bafcc3db61af7964b25428c Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 27 Sep 2022 10:50:43 -0400
Subject: [PATCH 39/63] Remove job_id generation in table/col check operators

The job_id is automatically generated by hook.insert_job()
if an empty string is passed, so job_id generation in the
operator is removed in favor of the existing code.

Signed-off-by: Benji Lampel <benjamin@astronomer.io>
---
 .../providers/google/cloud/operators/bigquery.py | 16 ++--------------
 1 file changed, 2 insertions(+), 14 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 8db80d993ac28..7b35ffc12ad90 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -591,13 +591,7 @@ def execute(self, context=None):
             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};"
 
-            job_id = hook.generate_job_id(
-                dag_id=self.dag_id,
-                task_id=self.task_id,
-                logical_date=context["logical_date"],
-                configuration=self.configuration,
-            )
-            job = self._submit_job(hook, job_id=job_id)
+            job = self._submit_job(hook, job_id="")
             context["ti"].xcom_push(key="job_id", value=job.job_id)
             records = list(job.result().to_dataframe().values.flatten())
 
@@ -711,13 +705,7 @@ def execute(self, context=None):
         self.sql = f"SELECT check_name, check_result FROM ({checks_sql}) "
         f"AS check_table {partition_clause_statement};"
 
-        job_id = hook.generate_job_id(
-            dag_id=self.dag_id,
-            task_id=self.task_id,
-            logical_date=context["logical_date"],
-            configuration=self.configuration,
-        )
-        job = self._submit_job(hook, job_id=job_id)
+        job = self._submit_job(hook, job_id="")
         context["ti"].xcom_push(key="job_id", value=job.job_id)
         records = job.result().to_dataframe()
 

From d4758cb202672b856a5c0ee74b83453fdbbf0cd4 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 14:27:27 -0400
Subject: [PATCH 40/63] Rework SQL query building

SQL query building is moved to the init() method of the column
and table check operators to lessen the amount of duplicate code
in the child operator. It also has the added effect of, ideally,
passing a more complete query to OpenLineage.

In doing the above, the column check operator had to be reworked and
now matches the logic of the table check operator in terms of
returning multiple rows and only sending one query to the database.
---
 airflow/providers/common/sql/operators/sql.py | 84 +++++++++++--------
 .../google/cloud/operators/bigquery.py        | 45 ++++------
 2 files changed, 65 insertions(+), 64 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 66984a802f862..1f6158bacc397 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -276,12 +276,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_statment} 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__(
@@ -295,38 +300,50 @@ 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:
@@ -465,8 +482,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__(
@@ -484,16 +501,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()
             ]
         )
@@ -503,6 +515,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:
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 7b35ffc12ad90..b0cbce69853da 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -585,28 +585,27 @@ def execute(self, context=None):
         """Perform checks on the given columns."""
         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};"
 
-            job = self._submit_job(hook, job_id="")
-            context["ti"].xcom_push(key="job_id", value=job.job_id)
-            records = list(job.result().to_dataframe().values.flatten())
+        job = self._submit_job(hook, job_id="")
+        context["ti"].xcom_push(key="job_id", value=job.job_id)
+        records = job.result().to_dataframe()
 
-            if not records:
-                raise AirflowException(f"The following query returned zero rows: {self.sql}")
+        if records.empty:
+            raise AirflowException(f"The following query returned zero rows: {self.sql}")
 
-            self.log.info("Record: %s", records)
+        records.columns = records.columns.str.lower()
+        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.iterrows():
+            column = row[1].get("col_name")
+            check = row[1].get("check_type")
+            result = row[1].get("check_result")
+            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:
@@ -693,18 +692,6 @@ def _submit_job(
     def execute(self, context=None):
         """Execute the given checks on the table."""
         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)
-                for check_name, value in self.checks.items()
-            ]
-        )
-        partition_clause_statement = f"WHERE {self.partition_clause}" if self.partition_clause else ""
-        self.sql = f"SELECT check_name, check_result FROM ({checks_sql}) "
-        f"AS check_table {partition_clause_statement};"
-
         job = self._submit_job(hook, job_id="")
         context["ti"].xcom_push(key="job_id", value=job.job_id)
         records = job.result().to_dataframe()

From cc87c3af275b35f2132144a98814a3fd979a5c41 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 14:52:13 -0400
Subject: [PATCH 41/63] Remove self.sql overwrite parameter in BigQuery check
 operators.

---
 airflow/providers/google/cloud/operators/bigquery.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index b0cbce69853da..486020402bfc5 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -562,8 +562,6 @@ def __init__(
         self.location = location
         self.impersonation_chain = impersonation_chain
         self.labels = labels
-        # OpenLineage needs a valid SQL query with the input/output table(s) to parse
-        self.sql = ""
 
     def _submit_job(
         self,
@@ -670,8 +668,6 @@ def __init__(
         self.location = location
         self.impersonation_chain = impersonation_chain
         self.labels = labels
-        # OpenLineage needs a valid SQL query with the input/output table(s) to parse
-        self.sql = ""
 
     def _submit_job(
         self,

From c5d74228fa168a21856081342b2baf1a8bf88a7f Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 15:25:39 -0400
Subject: [PATCH 42/63] Add option to fail check operators without retries

Adds a new parameter, retry_on_failure, and a new function to
determine if operators should retry or not on test failure.
---
 airflow/providers/common/sql/operators/sql.py | 45 ++++++++++++-------
 .../google/cloud/operators/bigquery.py        | 26 ++++++-----
 2 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 1f6158bacc397..52aa110585409 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -22,7 +22,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, Sequence, SupportsAbs
 
 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, fetch_all_handler
@@ -58,7 +58,13 @@ def _get_failed_checks(checks, col=None):
     ]
 
 
-_PROVIDERS_MATCHER = re.compile(r"airflow\.providers\.(.*)\.hooks.*")
+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 = {
     "amazon": "4.1.0",
@@ -103,12 +109,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):
@@ -347,11 +355,11 @@ def execute(self, context: Context):
 
             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")
 
@@ -530,11 +538,12 @@ def execute(self, context: Context):
 
         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")
 
@@ -594,7 +603,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.")
 
@@ -671,7 +682,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]
@@ -823,7 +834,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")
 
@@ -900,7 +913,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)
 
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 486020402bfc5..24eabe6fea386 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -38,6 +38,7 @@
     SQLTableCheckOperator,
     SQLValueCheckOperator,
     _get_failed_checks,
+    _raise_exception,
     parse_boolean,
 )
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
@@ -248,7 +249,9 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
         if not records:
             raise AirflowException("The query returned empty results")
         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("Record: %s", event["records"])
         self.log.info("Success.")
 
@@ -607,11 +610,11 @@ def execute(self, context=None):
 
             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")
 
@@ -705,11 +708,12 @@ def execute(self, context=None):
 
         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")
 

From 8080fc6765b8461b64a2e66b635e76e5487b61ef Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 28 Sep 2022 16:03:29 -0400
Subject: [PATCH 43/63] Rename and reorder private functions, fix typo

---
 airflow/providers/common/sql/operators/sql.py | 36 +++++++++----------
 .../google/cloud/operators/bigquery.py        |  4 +--
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 52aa110585409..b92e7b99839a8 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -31,7 +31,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.
@@ -286,7 +300,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
 
     sql_check_template = """
         SELECT '{column}' AS col_name, '{check}' AS check_type, {column}_{check} AS check_result
-        FROM (SELECT {check_statment} AS {column}_{check} FROM {table}) AS sq
+        FROM (SELECT {check_statement} AS {column}_{check} FROM {table}) AS sq
     """
 
     column_checks = {
@@ -534,7 +548,7 @@ 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:
@@ -996,7 +1010,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):
@@ -1014,17 +1028,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
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 24eabe6fea386..da2610a72eb2b 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -38,8 +38,8 @@
     SQLTableCheckOperator,
     SQLValueCheckOperator,
     _get_failed_checks,
+    _parse_boolean,
     _raise_exception,
-    parse_boolean,
 )
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
@@ -704,7 +704,7 @@ def execute(self, context=None):
         for row in records.iterrows():
             check = row[1].get("check_name")
             result = row[1].get("check_result")
-            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:

From 625d2b89746970c2f4d0ba57b9821e54156e1a47 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 09:56:33 -0400
Subject: [PATCH 44/63] Update airflow/providers/common/sql/operators/sql.py

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>
---
 airflow/providers/common/sql/operators/sql.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index b92e7b99839a8..c03ad652f3884 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -72,7 +72,7 @@ def _get_failed_checks(checks, col=None):
     ]
 
 
-def _raise_exception(exception_string, retry_on_failure):
+def _raise_exception(exception_string: str, retry_on_failure: bool) -> NoReturn:
     if retry_on_failure:
         raise AirflowException(exception_string)
     raise AirflowFailException(exception_string)

From 4f3f400b4fb8d3e28854707b3d7334de46b6f27f Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 10:35:42 -0400
Subject: [PATCH 45/63] Small updates to helper functions

---
 airflow/providers/common/sql/operators/sql.py | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index c03ad652f3884..028dcf72db663 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -19,7 +19,9 @@
 
 import ast
 import re
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, Sequence, SupportsAbs
+from typing import TYPE_CHECKING, Any, Iterable, Mapping, NoReturn, Sequence, SupportsAbs
+
+from packaging.version import Version
 
 from airflow.compat.functools import cached_property
 from airflow.exceptions import AirflowException, AirflowFailException
@@ -31,18 +33,11 @@
     from airflow.utils.context import Context
 
 
-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
-    """
+def _convert_to_float_if_possible(s: str) -> float | str:
     try:
-        ret = float(s)
+        return float(s)
     except (ValueError, TypeError):
-        ret = s
-    return ret
+        return s
 
 
 def _parse_boolean(val: str) -> str | bool:
@@ -58,7 +53,7 @@ def _parse_boolean(val: str) -> str | bool:
     raise ValueError(f"{val!r} is not a boolean-like string value")
 
 
-def _get_failed_checks(checks, col=None):
+def _get_failed_checks(checks: dict[str, Any], col: str | None = None):
     if col:
         return [
             f"Column: {col}\nCheck: {check},\nCheck Values: {check_values}\n"

From 9400414f233e8183943dc6596a3f7523006ca073 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 10:42:40 -0400
Subject: [PATCH 46/63] Use _raise_exception when query returns 0 rows

---
 airflow/providers/common/sql/operators/sql.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 028dcf72db663..75d76d794f120 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -349,7 +349,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            raise AirflowException(f"The following query returned zero rows: {self.sql}")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         self.log.info("Record: %s", records)
 
@@ -537,7 +537,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            raise AirflowException(f"The following query returned zero rows: {self.sql}")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         self.log.info("Record:\n%s", records)
 
@@ -610,7 +610,7 @@ def execute(self, context: Context):
 
         self.log.info("Record: %s", records)
         if not records:
-            raise AirflowException("The query returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
         elif not all(bool(r) for r in records):
             _raise_exception(
                 f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}", self.retry_on_failure
@@ -662,7 +662,7 @@ def execute(self, context: Context):
         records = self.get_db_hook().get_first(self.sql)
 
         if not records:
-            raise AirflowException("The query returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         pass_value_conv = _convert_to_float_if_possible(self.pass_value)
         is_numeric_value_check = isinstance(pass_value_conv, float)
@@ -763,7 +763,7 @@ def __init__(
         if ratio_formula not in self.ratio_formulas:
             msg_template = "Invalid diff_method: {diff_method}. Supported diff methods are: {diff_methods}"
 
-            raise AirflowException(
+            raise AirflowFailException(
                 msg_template.format(diff_method=ratio_formula, diff_methods=self.ratio_formulas)
             )
         self.ratio_formula = ratio_formula
@@ -788,9 +788,9 @@ def execute(self, context: Context):
         row1 = hook.get_first(self.sql1)
 
         if not row2:
-            raise AirflowException(f"The query {self.sql2} returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql2}", self.retry_on_failure)
         if not row1:
-            raise AirflowException(f"The query {self.sql1} returned None")
+            _raise_exception(f"The following query returned zero rows: {self.sql1}", self.retry_on_failure)
 
         current = dict(zip(self.metrics_sorted, row1))
         reference = dict(zip(self.metrics_sorted, row2))
@@ -888,6 +888,8 @@ def __init__(
     def execute(self, context: Context):
         hook = self.get_db_hook()
         result = hook.get_first(self.sql)[0]
+        if not result:
+            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
 
         if isinstance(self.min_threshold, float):
             lower_bound = self.min_threshold

From 0a794fb8d05648275b899108884adba9d30b1930 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 14:47:46 -0400
Subject: [PATCH 47/63] Update tests and operator code

Updates tests to reflect changes in operator code, and fixed
bugs in operators as well. Mainly moving the code to check for
failed tests into the column and table check operators as it works
slightly differently for each and doesn't make much sense as a
top-level function.
---
 airflow/providers/common/sql/operators/sql.py | 31 ++++-----
 .../common/sql/operators/test_sql.py          | 64 ++++++++++++++-----
 2 files changed, 61 insertions(+), 34 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 75d76d794f120..29adef8a6ba4c 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -53,20 +53,6 @@ def _parse_boolean(val: str) -> str | bool:
     raise ValueError(f"{val!r} is not a boolean-like string value")
 
 
-def _get_failed_checks(checks: dict[str, Any], col: str | None = None):
-    if col:
-        return [
-            f"Column: {col}\nCheck: {check},\nCheck Values: {check_values}\n"
-            for check, check_values in checks.items()
-            if not check_values["success"]
-        ]
-    return [
-        f"\tCheck: {check},\n\tCheck Values: {check_values}\n"
-        for check, check_values in checks.items()
-        if not check_values["success"]
-    ]
-
-
 def _raise_exception(exception_string: str, retry_on_failure: bool) -> NoReturn:
     if retry_on_failure:
         raise AirflowException(exception_string)
@@ -344,8 +330,8 @@ def __init__(
         """
 
     def execute(self, context: Context):
-        hook = self.get_db_hook()
         failed_tests = []
+        hook = self.get_db_hook()
         records = hook.get_records(self.sql)
 
         if not records:
@@ -362,7 +348,14 @@ def execute(self, context: Context):
                 self.column_mapping[column][check], result, tolerance
             )
 
-            failed_tests.extend(_get_failed_checks(self.column_mapping[column], column))
+        for col, checks in self.column_mapping.items():
+            failed_tests.extend(
+                [
+                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+                    for check, check_values in checks.items()
+                    if not check_values["success"]
+                ]
+            )
         if failed_tests:
             exception_string = f"""
                 Test failed.\nResults:\n{records!s}\n
@@ -545,7 +538,11 @@ def execute(self, context: Context):
             check, result = row
             self.checks[check]["success"] = _parse_boolean(str(result))
 
-        failed_tests = _get_failed_checks(self.checks)
+        failed_tests = [
+            f"\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for check, check_values in self.checks.items()
+            if not check_values["success"]
+        ]
         if failed_tests:
             exception_string = f"""
                 Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 46681d468ef17..c41a45daf67b4 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -47,9 +47,6 @@
 
 
 class MockHook:
-    def get_first(self):
-        return
-
     def get_records(self):
         return
 
@@ -110,15 +107,15 @@ class TestColumnCheckOperator:
 
     invalid_column_mapping = {"Y": {"invalid_check_name": {"expectation": 5}}}
 
-    def _construct_operator(self, monkeypatch, column_mapping, return_vals):
-        def get_first_return(*arg):
-            return return_vals
+    def _construct_operator(self, monkeypatch, column_mapping, records):
+        def get_records(*arg):
+            return records
 
         operator = SQLColumnCheckOperator(
             task_id="test_task", table="test_table", column_mapping=column_mapping
         )
         monkeypatch.setattr(operator, "get_db_hook", _get_mock_db_hook)
-        monkeypatch.setattr(MockHook, "get_first", get_first_return)
+        monkeypatch.setattr(MockHook, "get_records", get_records)
         return operator
 
     def test_check_not_in_column_checks(self, monkeypatch):
@@ -126,29 +123,62 @@ def test_check_not_in_column_checks(self, monkeypatch):
             self._construct_operator(monkeypatch, self.invalid_column_mapping, ())
 
     def test_pass_all_checks_exact_check(self, monkeypatch):
-        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 10, 10, 1, 19))
+        records = [
+            ('X', 'null_check', 0),
+            ('X', 'distinct_check', 10),
+            ('X', 'unique_check', 10),
+            ('X', 'min', 1),
+            ('X', 'max', 19),
+        ]
+        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
 
     def test_max_less_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
-            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 10, 10, 1, 21))
+            records = [
+                ('X', 'null_check', 1),
+                ('X', 'distinct_check', 10),
+                ('X', 'unique_check', 10),
+                ('X', 'min', 1),
+                ('X', 'max', 21),
+            ]
+            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
             operator.execute(context=MagicMock())
             assert operator.column_mapping["X"]["max"]["success"] is False
 
     def test_max_greater_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
-            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 10, 10, 1, 9))
+            records = [
+                ('X', 'null_check', 1),
+                ('X', 'distinct_check', 10),
+                ('X', 'unique_check', 10),
+                ('X', 'min', 1),
+                ('X', 'max', 9),
+            ]
+            operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
             operator.execute(context=MagicMock())
             assert operator.column_mapping["X"]["max"]["success"] is False
 
     def test_pass_all_checks_inexact_check(self, monkeypatch):
-        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, (0, 9, 12, 0, 15))
+        records = [
+            ('X', 'null_check', 0),
+            ('X', 'distinct_check', 9),
+            ('X', 'unique_check', 12),
+            ('X', 'min', 0),
+            ('X', 'max', 15),
+        ]
+        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
 
     def test_fail_all_checks_check(self, monkeypatch):
-        operator = operator = self._construct_operator(
-            monkeypatch, self.valid_column_mapping, (1, 12, 11, -1, 20)
-        )
+        records = [
+            ('X', 'null_check', 1),
+            ('X', 'distinct_check', 12),
+            ('X', 'unique_check', 11),
+            ('X', 'min', -1),
+            ('X', 'max', 20),
+        ]
+        operator = operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         with pytest.raises(AirflowException):
             operator.execute(context=MagicMock())
 
@@ -160,9 +190,9 @@ class TestTableCheckOperator:
         "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
     }
 
-    def _construct_operator(self, monkeypatch, checks, return_df):
+    def _construct_operator(self, monkeypatch, checks, records):
         def get_records(*arg):
-            return return_df
+            return records
 
         operator = SQLTableCheckOperator(task_id="test_task", table="test_table", checks=checks)
         monkeypatch.setattr(operator, "get_db_hook", _get_mock_db_hook)
@@ -303,7 +333,7 @@ def setUp(self):
     def test_execute_no_records(self, mock_get_db_hook):
         mock_get_db_hook.return_value.get_first.return_value = []
 
-        with pytest.raises(AirflowException, match=r"The query returned None"):
+        with pytest.raises(AirflowException, match=r"The following query returned zero rows: sql"):
             self._operator.execute({})
 
     @mock.patch.object(SQLCheckOperator, "get_db_hook")

From efa1f798b2b7d6ab1af31294ea6dc0c98702e5f6 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 29 Sep 2022 16:40:25 -0400
Subject: [PATCH 48/63] Fix _failed_checks() to match update in parent
 operators

---
 .../providers/google/cloud/operators/bigquery.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index da2610a72eb2b..b7985c8381469 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -37,7 +37,6 @@
     SQLIntervalCheckOperator,
     SQLTableCheckOperator,
     SQLValueCheckOperator,
-    _get_failed_checks,
     _parse_boolean,
     _raise_exception,
 )
@@ -608,7 +607,14 @@ def execute(self, context=None):
                 self.column_mapping[column][check], result, tolerance
             )
 
-            failed_tests.extend(_get_failed_checks(self.column_mapping[column], column))
+        for col, checks in self.column_mapping.items():
+            failed_tests.extend(
+                [
+                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+                    for check, check_values in checks.items()
+                    if not check_values["success"]
+                ]
+            )
         if failed_tests:
             exception_string = f"""
                 Test failed.\nResults:\n{records!s}\n
@@ -706,7 +712,11 @@ def execute(self, context=None):
             result = row[1].get("check_result")
             self.checks[check]["success"] = _parse_boolean(str(result))
 
-        failed_tests = _get_failed_checks(self.checks)
+        failed_tests = [
+            f"\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for check, check_values in self.checks.items()
+            if not check_values["success"]
+        ]
         if failed_tests:
             exception_string = f"""
                 Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n

From 213d0a236ab36d159ed80728f66ff19734c5f32c Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Fri, 30 Sep 2022 09:33:39 -0400
Subject: [PATCH 49/63] Insert quotes around check_name table op's sql template
 to fix sql error

---
 airflow/providers/common/sql/operators/sql.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 29adef8a6ba4c..ae9994f0e978f 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -492,7 +492,7 @@ class SQLTableCheckOperator(BaseSQLOperator):
     template_fields = ("partition_clause",)
 
     sql_check_template = """
-        SELECT {check_name} AS check_name, MIN({check_name}) AS check_result
+        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
     """
 

From 12c7cefc281e70a1bda539b4f0bbd295a10d3589 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 4 Oct 2022 11:27:07 -0400
Subject: [PATCH 50/63] Remove unnecessary list comprehension

---
 airflow/providers/common/sql/operators/sql.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index ae9994f0e978f..b9e52b8948bd6 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -310,7 +310,6 @@ def __init__(
         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(
@@ -319,7 +318,7 @@ def __init__(
                         table=self.table,
                         column=column,
                     )
-                    for check in checks_list
+                    for check in checks
                 ]
             )
         self.partition_clause = partition_clause

From 1ff14c2f9fde2b3e7ed62026820eeb124dd758e3 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 6 Oct 2022 13:32:56 -0400
Subject: [PATCH 51/63] Added assertions in existing column and table tests

---
 tests/providers/common/sql/operators/test_sql.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index c41a45daf67b4..ffbdbc7d224a7 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -132,6 +132,10 @@ def test_pass_all_checks_exact_check(self, monkeypatch):
         ]
         operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
+        assert [
+            operator.column_mapping["X"][check]["success"] is True
+            for check in [*operator.column_mapping["X"]]
+        ]
 
     def test_max_less_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
@@ -169,6 +173,10 @@ def test_pass_all_checks_inexact_check(self, monkeypatch):
         ]
         operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
+        assert [
+            operator.column_mapping["X"][check]["success"] is True
+            for check in [*operator.column_mapping["X"]]
+        ]
 
     def test_fail_all_checks_check(self, monkeypatch):
         records = [
@@ -240,6 +248,7 @@ def test_pass_all_checks_check(self, monkeypatch):
         records = [("row_count_check", 1), ("column_sum_check", "y")]
         operator = self._construct_operator(monkeypatch, self.checks, records)
         operator.execute(context=MagicMock())
+        assert [operator.checks[check]["success"] is True for check in operator.checks.keys()]
 
     def test_fail_all_checks_check(self, monkeypatch):
         records = [("row_count_check", 0), ("column_sum_check", "n")]

From e36aaf55cf120c5627202e9924a6c8c4eea8a7b3 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Thu, 6 Oct 2022 16:30:00 -0400
Subject: [PATCH 52/63] Add new tests for TableCheckOperator

---
 .../common/sql/operators/test_sql.py          | 46 +++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index ffbdbc7d224a7..04c8997c84f29 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -198,6 +198,30 @@ class TestTableCheckOperator:
         "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
     }
 
+    correct_generate_sql_query_no_partitions = """
+    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
+    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    UNION ALL
+    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
+    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table ) AS sq
+    """
+
+    correct_generate_sql_query_with_partition = """
+    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
+    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10) AS sq
+    UNION ALL
+    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
+    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
+    """
+
+    correct_generate_sql_query_with_partition_and_where = """
+    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
+    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10 AND id = 100) AS sq
+    UNION ALL
+    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
+    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
+    """
+
     def _construct_operator(self, monkeypatch, checks, records):
         def get_records(*arg):
             return records
@@ -256,6 +280,28 @@ def test_fail_all_checks_check(self, monkeypatch):
         with pytest.raises(AirflowException):
             operator.execute(context=MagicMock())
 
+    def test_generate_sql_query_no_partitions(self, monkeypatch):
+        operator = self._construct_operator(monkeypatch, self.checks, ())
+        assert (
+            operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_no_partitions.lstrip()
+        )
+
+    def test_generate_sql_query_with_partitions(self, monkeypatch):
+        operator = self._construct_operator(monkeypatch, self.checks, ())
+        operator.partition_clause = "col_a > 10"
+        assert (
+            operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_partition.lstrip()
+        )
+
+    def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
+        operator = self._construct_operator(monkeypatch, self.checks, ())
+        operator.partition_clause = "col_a > 10"
+        self.checks["row_count_check"]["where"] = "id = 100"
+        assert (
+            operator._generate_sql_query().lstrip()
+            == self.correct_generate_sql_query_with_partition_and_where.lstrip()
+        )
+
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)
 INTERVAL = datetime.timedelta(hours=12)

From ac4fca4217d3fe090923887d6a0289df22fd34dd Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Fri, 7 Oct 2022 10:26:08 -0400
Subject: [PATCH 53/63] Update operator logic and tests

Adds "where" option in checks dictionaries for column and table
operators, which may be renamed. This allows for check-level
partitioning, whereas the partition_clause param will always be
for all checks. New tests are added for this addition.
---
 airflow/providers/common/sql/operators/sql.py | 92 ++++++++++++-------
 .../common/sql/operators/test_sql.py          | 47 +++++++---
 2 files changed, 91 insertions(+), 48 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index b9e52b8948bd6..dd035cbe32f77 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -306,27 +306,16 @@ def __init__(
 
         self.table = table
         self.column_mapping = column_mapping
-        checks_sql = ""
+        self.partition_clause = partition_clause
+
+        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_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
-                ]
-            )
-        self.partition_clause = partition_clause
-        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}
-        """
+            checks_sql.append(self._generate_sql_query(column, checks))
+        checks_sql = " UNION ALL ".join(checks_sql)
+
+        self.sql = f"SELECT col_name, check_type, check_result FROM ({checks_sql}) AS check_columns"
 
     def execute(self, context: Context):
         failed_tests = []
@@ -364,6 +353,31 @@ def execute(self, context: Context):
 
         self.log.info("All tests have passed")
 
+    def _generate_sql_query(self, column, checks):
+        def _generate_partition_clause(check):
+            if self.partition_clause and "where" not in checks[check]:
+                return " WHERE " + self.partition_clause
+            elif not self.partition_clause and "where" in checks[check]:
+                return " WHERE " + checks[check]["where"]
+            elif self.partition_clause and "where" in checks[check]:
+                return " WHERE " + self.partition_clause + " AND " + checks[check]["where"]
+            else:
+                return ""
+
+        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,
+                    partition_clause=_generate_partition_clause(check),
+                )
+                for check in checks
+            ]
+        )
+        return checks_sql
+
     def _get_match(self, check_values, record, tolerance=None) -> bool:
         match_boolean = True
         if "geq_to" in check_values:
@@ -491,8 +505,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} {partition_clause}) AS sq
     """
 
     def __init__(
@@ -510,19 +524,7 @@ def __init__(
         self.table = table
         self.checks = checks
         self.partition_clause = partition_clause
-        checks_sql = " UNION ALL ".join(
-            [
-                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()
-            ]
-        )
-        partition_clause_statement = f"WHERE {self.partition_clause}" if self.partition_clause else ""
-        self.sql = f"""
-            SELECT check_name, check_result FROM ({checks_sql})
-            AS check_table {partition_clause_statement}
-        """
+        self.sql = f"SELECT check_name, check_result FROM ({self._generate_sql_query()}) AS check_table"
 
     def execute(self, context: Context):
         hook = self.get_db_hook()
@@ -552,6 +554,30 @@ def execute(self, context: Context):
 
         self.log.info("All tests have passed")
 
+    def _generate_sql_query(self):
+        def _generate_partition_clause(check_name):
+            if self.partition_clause and "where" not in self.checks[check_name]:
+                return "WHERE " + self.partition_clause
+            elif not self.partition_clause and "where" in self.checks[check_name]:
+                return "WHERE " + self.checks[check_name]["where"]
+            elif self.partition_clause and "where" in self.checks[check_name]:
+                return "WHERE " + self.partition_clause + " AND " + self.checks[check_name]["where"]
+            else:
+                return ""
+
+        checks_sql = "UNION ALL".join(
+            [
+                self.sql_check_template.format(
+                    check_statement=value["check_statement"],
+                    check_name=check_name,
+                    table=self.table,
+                    partition_clause=_generate_partition_clause(check_name),
+                )
+                for check_name, value in self.checks.items()
+            ]
+        )
+        return checks_sql
+
 
 class SQLCheckOperator(BaseSQLOperator):
     """
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 04c8997c84f29..a0990cc50371c 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -193,34 +193,44 @@ def test_fail_all_checks_check(self, monkeypatch):
 
 class TestTableCheckOperator:
 
+    count_check = "COUNT(*) == 1000"
+    sum_check = "col_a + col_b < col_c"
     checks = {
-        "row_count_check": {"check_statement": "COUNT(*) == 1000"},
-        "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
+        "row_count_check": {"check_statement": f"{count_check}"},
+        "column_sum_check": {"check_statement": f"{sum_check}"},
     }
 
-    correct_generate_sql_query_no_partitions = """
+    correct_generate_sql_query_no_partitions = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table ) AS sq
     """
 
-    correct_generate_sql_query_with_partition = """
+    correct_generate_sql_query_with_partition = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
-    """
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
+    """  # noqa 501
 
-    correct_generate_sql_query_with_partition_and_where = """
+    correct_generate_sql_query_with_partition_and_where = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN COUNT(*) == 1000 THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10 AND id = 100) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10 AND id = 100) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN col_a + col_b < col_c THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
-    """
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
+    """  # noqa 501
+
+    correct_generate_sql_query_with_where = f"""
+    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    UNION ALL
+    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE id = 100) AS sq
+    """  # noqa 501
 
     def _construct_operator(self, monkeypatch, checks, records):
         def get_records(*arg):
@@ -294,14 +304,21 @@ def test_generate_sql_query_with_partitions(self, monkeypatch):
         )
 
     def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
-        operator = self._construct_operator(monkeypatch, self.checks, ())
+        checks = self.checks
+        checks["row_count_check"]["where"] = "id = 100"
+        operator = self._construct_operator(monkeypatch, checks, ())
         operator.partition_clause = "col_a > 10"
-        self.checks["row_count_check"]["where"] = "id = 100"
         assert (
             operator._generate_sql_query().lstrip()
             == self.correct_generate_sql_query_with_partition_and_where.lstrip()
         )
 
+    def test_generate_sql_query_with_check_partition(self, monkeypatch):
+        checks = self.checks
+        checks["column_sum_check"]["where"] = "id = 100"
+        operator = self._construct_operator(monkeypatch, checks, ())
+        assert operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_where.lstrip()
+
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)
 INTERVAL = datetime.timedelta(hours=12)

From 8beb8ab22a0e0f3e0a7e579a7d0e4c65465e686f Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Fri, 7 Oct 2022 11:35:16 -0400
Subject: [PATCH 54/63] Add testing for column check operator and line edits

Cleans up operator and adds testing for new generator function.
---
 airflow/providers/common/sql/operators/sql.py |  19 +--
 .../common/sql/operators/test_sql.py          | 119 +++++++++++++++---
 2 files changed, 112 insertions(+), 26 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index dd035cbe32f77..4cd13214fbcab 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -281,7 +281,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
 
     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
+        FROM (SELECT {check_statement} AS {column}_{check} FROM {table} {partition_clause}) AS sq
     """
 
     column_checks = {
@@ -308,12 +308,12 @@ def __init__(
         self.column_mapping = column_mapping
         self.partition_clause = partition_clause
 
-        checks_sql = []
+        checks_sql_list = []
         for column, checks in self.column_mapping.items():
             for check, check_values in checks.items():
                 self._column_mapping_validation(check, check_values)
-            checks_sql.append(self._generate_sql_query(column, checks))
-        checks_sql = " UNION ALL ".join(checks_sql)
+            checks_sql_list.append(self._generate_sql_query(column, checks))
+        checks_sql = "UNION ALL".join(checks_sql_list)
 
         self.sql = f"SELECT col_name, check_type, check_result FROM ({checks_sql}) AS check_columns"
 
@@ -356,15 +356,15 @@ def execute(self, context: Context):
     def _generate_sql_query(self, column, checks):
         def _generate_partition_clause(check):
             if self.partition_clause and "where" not in checks[check]:
-                return " WHERE " + self.partition_clause
+                return "WHERE " + self.partition_clause
             elif not self.partition_clause and "where" in checks[check]:
-                return " WHERE " + checks[check]["where"]
+                return "WHERE " + checks[check]["where"]
             elif self.partition_clause and "where" in checks[check]:
-                return " WHERE " + self.partition_clause + " AND " + checks[check]["where"]
+                return "WHERE " + self.partition_clause + " AND " + checks[check]["where"]
             else:
                 return ""
 
-        checks_sql = " UNION ALL ".join(
+        checks_sql = "UNION ALL".join(
             [
                 self.sql_check_template.format(
                     check_statement=self.column_checks[check].format(column=column),
@@ -506,7 +506,8 @@ class SQLTableCheckOperator(BaseSQLOperator):
 
     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} {partition_clause}) AS sq
+    FROM (SELECT CASE WHEN {check_statement} THEN 1 ELSE 0 END AS {check_name}
+          FROM {table} {partition_clause}) AS sq
     """
 
     def __init__(
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index a0990cc50371c..7413014b84925 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -105,8 +105,47 @@ class TestColumnCheckOperator:
         }
     }
 
+    short_valid_column_mapping = {
+        "X": {
+            "null_check": {"equal_to": 0},
+            "distinct_check": {"equal_to": 10, "tolerance": 0.1},
+        }
+    }
+
     invalid_column_mapping = {"Y": {"invalid_check_name": {"expectation": 5}}}
 
+    correct_generate_sql_query_no_partitions = """
+        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
+        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table ) AS sq
+    UNION ALL
+        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
+        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table ) AS sq
+    """
+
+    correct_generate_sql_query_with_partition = """
+        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
+        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table WHERE Y > 1) AS sq
+    UNION ALL
+        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
+        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table WHERE Y > 1) AS sq
+    """  # noqa 501
+
+    correct_generate_sql_query_with_partition_and_where = """
+        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
+        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table WHERE Y > 1 AND Z < 100) AS sq
+    UNION ALL
+        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
+        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table WHERE Y > 1) AS sq
+    """  # noqa 501
+
+    correct_generate_sql_query_with_where = """
+        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
+        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table ) AS sq
+    UNION ALL
+        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
+        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table WHERE Z < 100) AS sq
+    """  # 501
+
     def _construct_operator(self, monkeypatch, column_mapping, records):
         def get_records(*arg):
             return records
@@ -190,6 +229,44 @@ def test_fail_all_checks_check(self, monkeypatch):
         with pytest.raises(AirflowException):
             operator.execute(context=MagicMock())
 
+    def test_generate_sql_query_no_partitions(self, monkeypatch):
+        checks = self.short_valid_column_mapping["X"]
+        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
+        assert (
+            operator._generate_sql_query("X", checks).lstrip()
+            == self.correct_generate_sql_query_no_partitions.lstrip()
+        )
+
+    def test_generate_sql_query_with_partitions(self, monkeypatch):
+        checks = self.short_valid_column_mapping["X"]
+        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
+        operator.partition_clause = "Y > 1"
+        assert (
+            operator._generate_sql_query("X", checks).lstrip()
+            == self.correct_generate_sql_query_with_partition.lstrip()
+        )
+
+    def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
+        self.short_valid_column_mapping["X"]["null_check"]["where"] = "Z < 100"
+        checks = self.short_valid_column_mapping["X"]
+        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
+        operator.partition_clause = "Y > 1"
+        assert (
+            operator._generate_sql_query("X", checks).lstrip()
+            == self.correct_generate_sql_query_with_partition_and_where.lstrip()
+        )
+        del self.short_valid_column_mapping["X"]["null_check"]["where"]
+
+    def test_generate_sql_query_with_check_partition(self, monkeypatch):
+        self.short_valid_column_mapping["X"]["distinct_check"]["where"] = "Z < 100"
+        checks = self.short_valid_column_mapping["X"]
+        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
+        assert (
+            operator._generate_sql_query("X", checks).lstrip()
+            == self.correct_generate_sql_query_with_where.lstrip()
+        )
+        del self.short_valid_column_mapping["X"]["distinct_check"]["where"]
+
 
 class TestTableCheckOperator:
 
@@ -202,35 +279,43 @@ class TestTableCheckOperator:
 
     correct_generate_sql_query_no_partitions = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
+          FROM test_table ) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
+          FROM test_table ) AS sq
     """
 
     correct_generate_sql_query_with_partition = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
+          FROM test_table WHERE col_a > 10) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
-    """  # noqa 501
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
+          FROM test_table WHERE col_a > 10) AS sq
+    """
 
     correct_generate_sql_query_with_partition_and_where = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table WHERE col_a > 10 AND id = 100) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
+          FROM test_table WHERE col_a > 10 AND id = 100) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE col_a > 10) AS sq
-    """  # noqa 501
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
+          FROM test_table WHERE col_a > 10) AS sq
+    """
 
     correct_generate_sql_query_with_where = f"""
     SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
-    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check FROM test_table ) AS sq
+    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
+          FROM test_table ) AS sq
     UNION ALL
     SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
-    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check FROM test_table WHERE id = 100) AS sq
-    """  # noqa 501
+    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
+          FROM test_table WHERE id = 100) AS sq
+    """
 
     def _construct_operator(self, monkeypatch, checks, records):
         def get_records(*arg):
@@ -304,20 +389,20 @@ def test_generate_sql_query_with_partitions(self, monkeypatch):
         )
 
     def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
-        checks = self.checks
-        checks["row_count_check"]["where"] = "id = 100"
-        operator = self._construct_operator(monkeypatch, checks, ())
+        self.checks["row_count_check"]["where"] = "id = 100"
+        operator = self._construct_operator(monkeypatch, self.checks, ())
         operator.partition_clause = "col_a > 10"
         assert (
             operator._generate_sql_query().lstrip()
             == self.correct_generate_sql_query_with_partition_and_where.lstrip()
         )
+        del self.checks["row_count_check"]["where"]
 
     def test_generate_sql_query_with_check_partition(self, monkeypatch):
-        checks = self.checks
-        checks["column_sum_check"]["where"] = "id = 100"
-        operator = self._construct_operator(monkeypatch, checks, ())
+        self.checks["column_sum_check"]["where"] = "id = 100"
+        operator = self._construct_operator(monkeypatch, self.checks, ())
         assert operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_where.lstrip()
+        del self.checks["column_sum_check"]["where"]
 
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)

From 385317f67237907066c2d815c061f5ca674f1b52 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 11 Oct 2022 09:51:00 -0400
Subject: [PATCH 55/63] Change name 'where' to 'partition_clause' in check
 dictionaries

---
 airflow/providers/common/sql/operators/sql.py | 31 ++++++++++---------
 .../common/sql/operators/test_sql.py          | 16 +++++-----
 2 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 4cd13214fbcab..4f202bac4f66a 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -259,6 +259,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
                     "tolerance": 0.2,
                 },
                 "max": {"less_than": 1000, "geq_to": 10, "tolerance": 0.01},
+                "partition_clause": "foreign_key IS NOT NULL",
             }
         }
 
@@ -355,12 +356,12 @@ def execute(self, context: Context):
 
     def _generate_sql_query(self, column, checks):
         def _generate_partition_clause(check):
-            if self.partition_clause and "where" not in checks[check]:
+            if self.partition_clause and "partition_clause" not in checks[check]:
                 return "WHERE " + self.partition_clause
-            elif not self.partition_clause and "where" in checks[check]:
-                return "WHERE " + checks[check]["where"]
-            elif self.partition_clause and "where" in checks[check]:
-                return "WHERE " + self.partition_clause + " AND " + checks[check]["where"]
+            elif not self.partition_clause and "partition_clause" in checks[check]:
+                return "WHERE " + checks[check]["partition_clause"]
+            elif self.partition_clause and "partition_clause" in checks[check]:
+                return "WHERE " + self.partition_clause + " AND " + checks[check]["partition_clause"]
             else:
                 return ""
 
@@ -477,13 +478,14 @@ class SQLTableCheckOperator(BaseSQLOperator):
     Checks should be written to return a boolean result.
 
     :param table: the table to run checks on
-    :param checks: the dictionary of checks, e.g.:
+    :param checks: the dictionary of checks, where check names are followed by a dictionary containing at
+        least a check statement, and optionally a partition clause, e.g.:
 
     .. code-block:: python
 
         {
             "row_count_check": {"check_statement": "COUNT(*) = 1000"},
-            "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
+            "column_sum_check": {"check_statement": "col_a + col_b < col_c", "partition_clause": "col_a IS NOT NULL"},
         }
 
 
@@ -557,16 +559,18 @@ def execute(self, context: Context):
 
     def _generate_sql_query(self):
         def _generate_partition_clause(check_name):
-            if self.partition_clause and "where" not in self.checks[check_name]:
+            if self.partition_clause and "partition_clause" not in self.checks[check_name]:
                 return "WHERE " + self.partition_clause
-            elif not self.partition_clause and "where" in self.checks[check_name]:
-                return "WHERE " + self.checks[check_name]["where"]
-            elif self.partition_clause and "where" in self.checks[check_name]:
-                return "WHERE " + self.partition_clause + " AND " + self.checks[check_name]["where"]
+            elif not self.partition_clause and "partition_clause" in self.checks[check_name]:
+                return "WHERE " + self.checks[check_name]["partition_clause"]
+            elif self.partition_clause and "partition_clause" in self.checks[check_name]:
+                return (
+                    "WHERE " + self.partition_clause + " AND " + self.checks[check_name]["partition_clause"]
+                )
             else:
                 return ""
 
-        checks_sql = "UNION ALL".join(
+        return "UNION ALL".join(
             [
                 self.sql_check_template.format(
                     check_statement=value["check_statement"],
@@ -577,7 +581,6 @@ def _generate_partition_clause(check_name):
                 for check_name, value in self.checks.items()
             ]
         )
-        return checks_sql
 
 
 class SQLCheckOperator(BaseSQLOperator):
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 7413014b84925..7611713423544 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -247,7 +247,7 @@ def test_generate_sql_query_with_partitions(self, monkeypatch):
         )
 
     def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
-        self.short_valid_column_mapping["X"]["null_check"]["where"] = "Z < 100"
+        self.short_valid_column_mapping["X"]["null_check"]["partition_clause"] = "Z < 100"
         checks = self.short_valid_column_mapping["X"]
         operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
         operator.partition_clause = "Y > 1"
@@ -255,17 +255,17 @@ def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatc
             operator._generate_sql_query("X", checks).lstrip()
             == self.correct_generate_sql_query_with_partition_and_where.lstrip()
         )
-        del self.short_valid_column_mapping["X"]["null_check"]["where"]
+        del self.short_valid_column_mapping["X"]["null_check"]["partition_clause"]
 
     def test_generate_sql_query_with_check_partition(self, monkeypatch):
-        self.short_valid_column_mapping["X"]["distinct_check"]["where"] = "Z < 100"
+        self.short_valid_column_mapping["X"]["distinct_check"]["partition_clause"] = "Z < 100"
         checks = self.short_valid_column_mapping["X"]
         operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
         assert (
             operator._generate_sql_query("X", checks).lstrip()
             == self.correct_generate_sql_query_with_where.lstrip()
         )
-        del self.short_valid_column_mapping["X"]["distinct_check"]["where"]
+        del self.short_valid_column_mapping["X"]["distinct_check"]["partition_clause"]
 
 
 class TestTableCheckOperator:
@@ -389,20 +389,20 @@ def test_generate_sql_query_with_partitions(self, monkeypatch):
         )
 
     def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
-        self.checks["row_count_check"]["where"] = "id = 100"
+        self.checks["row_count_check"]["partition_clause"] = "id = 100"
         operator = self._construct_operator(monkeypatch, self.checks, ())
         operator.partition_clause = "col_a > 10"
         assert (
             operator._generate_sql_query().lstrip()
             == self.correct_generate_sql_query_with_partition_and_where.lstrip()
         )
-        del self.checks["row_count_check"]["where"]
+        del self.checks["row_count_check"]["partition_clause"]
 
     def test_generate_sql_query_with_check_partition(self, monkeypatch):
-        self.checks["column_sum_check"]["where"] = "id = 100"
+        self.checks["column_sum_check"]["partition_clause"] = "id = 100"
         operator = self._construct_operator(monkeypatch, self.checks, ())
         assert operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_where.lstrip()
-        del self.checks["column_sum_check"]["where"]
+        del self.checks["column_sum_check"]["partition_clause"]
 
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)

From 979def57894063010cac91b16f4238f0d0fde486 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 11 Oct 2022 12:35:59 -0400
Subject: [PATCH 56/63] Update docs and use f strings

---
 airflow/providers/common/sql/operators/sql.py | 17 ++++++-----
 .../operators.rst                             | 28 +++++++++++--------
 2 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 4f202bac4f66a..d43136db7ecde 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -242,6 +242,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
     - geq_to: value that results should be greater than or equal to
     - leq_to: value that results should be less than or equal to
     - tolerance: the percentage that the result may be off from the expected value
+    - partition_clause: an extra clause passed into a WHERE statement to partition data
 
     :param table: the table to run checks on
     :param column_mapping: the dictionary of columns and their associated checks, e.g.
@@ -252,6 +253,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
             "col_name": {
                 "null_check": {
                     "equal_to": 0,
+                    "partition_clause": "foreign_key IS NOT NULL",
                 },
                 "min": {
                     "greater_than": 5,
@@ -259,7 +261,6 @@ class SQLColumnCheckOperator(BaseSQLOperator):
                     "tolerance": 0.2,
                 },
                 "max": {"less_than": 1000, "geq_to": 10, "tolerance": 0.01},
-                "partition_clause": "foreign_key IS NOT NULL",
             }
         }
 
@@ -357,11 +358,11 @@ def execute(self, context: Context):
     def _generate_sql_query(self, column, checks):
         def _generate_partition_clause(check):
             if self.partition_clause and "partition_clause" not in checks[check]:
-                return "WHERE " + self.partition_clause
+                return f"WHERE {self.partition_clause}"
             elif not self.partition_clause and "partition_clause" in checks[check]:
-                return "WHERE " + checks[check]["partition_clause"]
+                return f"WHERE {checks[check]['partition_clause']}"
             elif self.partition_clause and "partition_clause" in checks[check]:
-                return "WHERE " + self.partition_clause + " AND " + checks[check]["partition_clause"]
+                return f"WHERE {self.partition_clause} AND {checks[check]['partition_clause']}"
             else:
                 return ""
 
@@ -560,13 +561,11 @@ def execute(self, context: Context):
     def _generate_sql_query(self):
         def _generate_partition_clause(check_name):
             if self.partition_clause and "partition_clause" not in self.checks[check_name]:
-                return "WHERE " + self.partition_clause
+                return f"WHERE {self.partition_clause}"
             elif not self.partition_clause and "partition_clause" in self.checks[check_name]:
-                return "WHERE " + self.checks[check_name]["partition_clause"]
+                return f"WHERE {self.checks[check_name]['partition_clause']}"
             elif self.partition_clause and "partition_clause" in self.checks[check_name]:
-                return (
-                    "WHERE " + self.partition_clause + " AND " + self.checks[check_name]["partition_clause"]
-                )
+                return f"WHERE {self.partition_clause} AND {self.checks[check_name]['partition_clause']}"
             else:
                 return ""
 
diff --git a/docs/apache-airflow-providers-common-sql/operators.rst b/docs/apache-airflow-providers-common-sql/operators.rst
index e10759117e0b1..97bf9f4b6f6c1 100644
--- a/docs/apache-airflow-providers-common-sql/operators.rst
+++ b/docs/apache-airflow-providers-common-sql/operators.rst
@@ -51,16 +51,14 @@ Check SQL Table Columns
 
 Use the :class:`~airflow.providers.common.sql.operators.sql.SQLColumnCheckOperator` to run data quality
 checks against columns of a given table. As well as a connection ID and table, a column_mapping
-describing the relationship between columns and tests to run must be supplied. An example column
-mapping is a set of three nested dictionaries and looks like:
+describing the relationship between columns and tests to run must be supplied. An example column mapping
+is a set of three nested dictionaries and looks like:
 
 .. code-block:: python
 
         column_mapping = {
             "col_name": {
-                "null_check": {
-                    "equal_to": 0,
-                },
+                "null_check": {"equal_to": 0, "partition_clause": "other_col LIKE 'this'"},
                 "min": {
                     "greater_than": 5,
                     "leq_to": 10,
@@ -79,8 +77,8 @@ The valid checks are:
 - min: checks the minimum value in the column
 - max: checks the maximum value in the column
 
-Each entry in the check's dictionary is either a condition for success of the check or the tolerance. The
-conditions for success are:
+Each entry in the check's dictionary is either a condition for success of the check, the tolerance,
+or a partition clause. The conditions for success are:
 
 - greater_than
 - geq_to
@@ -92,7 +90,9 @@ When specifying conditions, equal_to is not compatible with other conditions. Bo
 bound condition may be specified in the same check. The tolerance is a percentage that the result may
 be out of bounds but still considered successful.
 
-
+The partition clauses may be given at the operator level as a parameter where it partitions all checks,
+at the column level in the column mapping where it partitions all checks for that column, or at the
+check level for a column where it partitions just that check.
 
 The below example demonstrates how to instantiate the SQLColumnCheckOperator task.
 
@@ -119,14 +119,20 @@ checks argument is a set of two nested dictionaries and looks like:
                 "row_count_check": {
                     "check_statement": "COUNT(*) = 1000",
                 },
-                "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
+                "column_sum_check": {
+                    "check_statement": "col_a + col_b < col_c",
+                    "partition_clause": "col_a IS NOT NULL",
+                },
             },
         )
 
 The first set of keys are the check names, which are referenced in the templated query the operator builds.
-The dictionary key under the check name must be check_statement, with the value a SQL statement that
+A dictionary key under the check name must include check_statement and the value a SQL statement that
 resolves to a boolean (this can be any string or int that resolves to a boolean in
-airflow.operators.sql.parse_boolean).
+airflow.operators.sql.parse_boolean). The other possible key to supply is partition_clause, which is a
+check level statement that will partition the data in the table using a WHERE clause for that check.
+This statement is compatible with the parameter partition_clause, where the latter filters across all
+checks.
 
 The below example demonstrates how to instantiate the SQLTableCheckOperator task.
 

From 586f6b45a37362a17739f894e48bda9202fc7a22 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Tue, 11 Oct 2022 12:37:56 -0400
Subject: [PATCH 57/63] Edit operator docstring

---
 airflow/providers/common/sql/operators/sql.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index d43136db7ecde..42035965d2ed5 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -486,7 +486,8 @@ class SQLTableCheckOperator(BaseSQLOperator):
 
         {
             "row_count_check": {"check_statement": "COUNT(*) = 1000"},
-            "column_sum_check": {"check_statement": "col_a + col_b < col_c", "partition_clause": "col_a IS NOT NULL"},
+            "column_sum_check": {"check_statement": "col_a + col_b < col_c"},
+            "third_check": {"check_statement": "MIN(col) = 1", "partition_clause": "col IS NOT NULL"},
         }
 
 

From c060da0626c2332bd299c66d19f28a1640fc689d Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 12 Oct 2022 09:11:38 -0400
Subject: [PATCH 58/63] Move _raise_exception to base class to simplify method

---
 airflow/providers/common/sql/operators/sql.py | 41 ++++++++-----------
 1 file changed, 18 insertions(+), 23 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index 42035965d2ed5..d231f267dd01c 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -53,12 +53,6 @@ def _parse_boolean(val: str) -> str | bool:
     raise ValueError(f"{val!r} is not a boolean-like string value")
 
 
-def _raise_exception(exception_string: str, retry_on_failure: bool) -> NoReturn:
-    if retry_on_failure:
-        raise AirflowException(exception_string)
-    raise AirflowFailException(exception_string)
-
-
 _PROVIDERS_MATCHER = re.compile(r'airflow\.providers\.(.*)\.hooks.*')
 
 _MIN_SUPPORTED_PROVIDERS_VERSION = {
@@ -158,6 +152,11 @@ def get_db_hook(self) -> DbApiHook:
         """
         return self._hook
 
+    def _raise_exception(self, exception_string: str) -> NoReturn:
+        if self.retry_on_failure:
+            raise AirflowException(exception_string)
+        raise AirflowFailException(exception_string)
+
 
 class SQLExecuteQueryOperator(BaseSQLOperator):
     """
@@ -325,7 +324,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
 
         self.log.info("Record: %s", records)
 
@@ -351,7 +350,7 @@ def execute(self, context: Context):
                 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._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
 
@@ -536,7 +535,7 @@ def execute(self, context: Context):
         records = hook.get_records(self.sql)
 
         if not records:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
 
         self.log.info("Record:\n%s", records)
 
@@ -555,7 +554,7 @@ def execute(self, context: Context):
                 The following tests have failed:
                 \n{', '.join(failed_tests)}
             """
-            _raise_exception(exception_string, self.retry_on_failure)
+            self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
 
@@ -636,11 +635,9 @@ def execute(self, context: Context):
 
         self.log.info("Record: %s", records)
         if not records:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
         elif not all(bool(r) for r in records):
-            _raise_exception(
-                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}", self.retry_on_failure
-            )
+            self._raise_exception(f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}")
 
         self.log.info("Success.")
 
@@ -688,7 +685,7 @@ def execute(self, context: Context):
         records = self.get_db_hook().get_first(self.sql)
 
         if not records:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
 
         pass_value_conv = _convert_to_float_if_possible(self.pass_value)
         is_numeric_value_check = isinstance(pass_value_conv, float)
@@ -717,7 +714,7 @@ def execute(self, context: Context):
             tests = []
 
         if not all(tests):
-            _raise_exception(error_msg, self.retry_on_failure)
+            self._raise_exception(error_msg)
 
     def _to_float(self, records):
         return [float(record) for record in records]
@@ -814,9 +811,9 @@ def execute(self, context: Context):
         row1 = hook.get_first(self.sql1)
 
         if not row2:
-            _raise_exception(f"The following query returned zero rows: {self.sql2}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql2}")
         if not row1:
-            _raise_exception(f"The following query returned zero rows: {self.sql1}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql1}")
 
         current = dict(zip(self.metrics_sorted, row1))
         reference = dict(zip(self.metrics_sorted, row2))
@@ -869,9 +866,7 @@ def execute(self, context: Context):
                     ratios[k],
                     self.metrics_thresholds[k],
                 )
-            _raise_exception(
-                f"The following tests have failed:\n {', '.join(sorted(failed_tests))}", self.retry_on_failure
-            )
+            self._raise_exception(f"The following tests have failed:\n {', '.join(sorted(failed_tests))}")
 
         self.log.info("All tests have passed")
 
@@ -915,7 +910,7 @@ def execute(self, context: Context):
         hook = self.get_db_hook()
         result = hook.get_first(self.sql)[0]
         if not result:
-            _raise_exception(f"The following query returned zero rows: {self.sql}", self.retry_on_failure)
+            self._raise_exception(f"The following query returned zero rows: {self.sql}")
 
         if isinstance(self.min_threshold, float):
             lower_bound = self.min_threshold
@@ -950,7 +945,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_exception(error_msg, self.retry_on_failure)
+            self._raise_exception(error_msg)
 
         self.log.info("Test %s Successful.", self.task_id)
 

From 8498cbfc87bab3768691edbec3613a96e6069fa4 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 12 Oct 2022 10:03:19 -0400
Subject: [PATCH 59/63] Updates from code review

---
 airflow/providers/common/sql/operators/sql.py | 69 ++++++++-----------
 1 file changed, 30 insertions(+), 39 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index d231f267dd01c..acb31efa3b002 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -319,7 +319,6 @@ def __init__(
         self.sql = f"SELECT col_name, check_type, check_result FROM ({checks_sql}) AS check_columns"
 
     def execute(self, context: Context):
-        failed_tests = []
         hook = self.get_db_hook()
         records = hook.get_records(self.sql)
 
@@ -328,8 +327,7 @@ def execute(self, context: Context):
 
         self.log.info("Record: %s", records)
 
-        for row in records:
-            column, check, result = row
+        for column, check, result in records:
             tolerance = self.column_mapping[column][check].get("tolerance")
 
             self.column_mapping[column][check]["result"] = result
@@ -337,19 +335,17 @@ def execute(self, context: Context):
                 self.column_mapping[column][check], result, tolerance
             )
 
-        for col, checks in self.column_mapping.items():
-            failed_tests.extend(
-                [
-                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
-                    for check, check_values in checks.items()
-                    if not check_values["success"]
-                ]
-            )
+        failed_tests = [
+            f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for col, checks in self.column_mapping.items()
+            for check, check_values in checks.items()
+            if not check_values["success"]
+        ]
         if failed_tests:
-            exception_string = f"""
-                Test failed.\nResults:\n{records!s}\n
-                The following tests have failed:
-                \n{''.join(failed_tests)}"""
+            exception_string = (
+                f"Test failed.\nResults:\n{records!s}\n"
+                f"The following tests have failed:\n{''.join(failed_tests)}"
+            )
             self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
@@ -366,16 +362,14 @@ def _generate_partition_clause(check):
                 return ""
 
         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,
-                    partition_clause=_generate_partition_clause(check),
-                )
-                for check in checks
-            ]
+            self.sql_check_template.format(
+                check_statement=self.column_checks[check].format(column=column),
+                check=check,
+                table=self.table,
+                column=column,
+                partition_clause=_generate_partition_clause(check),
+            )
+            for check in checks
         )
         return checks_sql
 
@@ -549,11 +543,10 @@ def execute(self, context: Context):
             if not check_values["success"]
         ]
         if 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)}
-            """
+            exception_string = (
+                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n"
+                f"The following tests have failed:\n{', '.join(failed_tests)}"
+            )
             self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
@@ -570,15 +563,13 @@ def _generate_partition_clause(check_name):
                 return ""
 
         return "UNION ALL".join(
-            [
-                self.sql_check_template.format(
-                    check_statement=value["check_statement"],
-                    check_name=check_name,
-                    table=self.table,
-                    partition_clause=_generate_partition_clause(check_name),
-                )
-                for check_name, value in self.checks.items()
-            ]
+            self.sql_check_template.format(
+                check_statement=value["check_statement"],
+                check_name=check_name,
+                table=self.table,
+                partition_clause=_generate_partition_clause(check_name),
+            )
+            for check_name, value in self.checks.items()
         )
 
 

From a15e1a23b95058018f4017d289382847267ce24b Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 12 Oct 2022 16:23:28 -0400
Subject: [PATCH 60/63] Update BigQuery Check operators according to code
 review

---
 .../google/cloud/operators/bigquery.py        | 41 ++++++++-----------
 1 file changed, 18 insertions(+), 23 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index b7985c8381469..9e711d2678e90 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -38,7 +38,6 @@
     SQLTableCheckOperator,
     SQLValueCheckOperator,
     _parse_boolean,
-    _raise_exception,
 )
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
@@ -248,9 +247,7 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
         if not records:
             raise AirflowException("The query returned empty results")
         elif not all(bool(r) for r in records):
-            _raise_exception(
-                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}", self.retry_on_failure
-            )
+            self._raise_exception(f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}")
         self.log.info("Record: %s", event["records"])
         self.log.info("Success.")
 
@@ -607,20 +604,19 @@ def execute(self, context=None):
                 self.column_mapping[column][check], result, tolerance
             )
 
-        for col, checks in self.column_mapping.items():
-            failed_tests.extend(
-                [
-                    f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
-                    for check, check_values in checks.items()
-                    if not check_values["success"]
-                ]
-            )
+        failed_tests(
+            f"Column: {col}\n\tCheck: {check},\n\tCheck Values: {check_values}\n"
+            for col, checks in self.column_mapping.items()
+            for check, check_values in checks.items()
+            if not check_values["success"]
+        )
         if 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)
+            exception_string = (
+                f"Test failed.\nResults:\n{records!s}\n"
+                f"The following tests have failed:"
+                f"\n{''.join(failed_tests)}"
+            )
+            self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
 
@@ -718,12 +714,11 @@ def execute(self, context=None):
             if not check_values["success"]
         ]
         if 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)
+            exception_string = (
+                f"Test failed.\nQuery:\n{self.sql}\nResults:\n{records!s}\n"
+                f"The following tests have failed:\n{', '.join(failed_tests)}"
+            )
+            self._raise_exception(exception_string)
 
         self.log.info("All tests have passed")
 

From 0f9bce3089750f2e82077be5e516d44d7a3eae02 Mon Sep 17 00:00:00 2001
From: Tzu-ping Chung <uranusjr@gmail.com>
Date: Thu, 13 Oct 2022 13:09:09 +0800
Subject: [PATCH 61/63] Rewrite data-building loop to generator

---
 airflow/providers/common/sql/operators/sql.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index acb31efa3b002..f408949932972 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -309,12 +309,13 @@ def __init__(
         self.column_mapping = column_mapping
         self.partition_clause = partition_clause
 
-        checks_sql_list = []
-        for column, checks in self.column_mapping.items():
-            for check, check_values in checks.items():
-                self._column_mapping_validation(check, check_values)
-            checks_sql_list.append(self._generate_sql_query(column, checks))
-        checks_sql = "UNION ALL".join(checks_sql_list)
+        def _build_checks_sql():
+            for column, checks in self.column_mapping.items():
+                for check, check_values in checks.items():
+                    self._column_mapping_validation(check, check_values)
+                yield self._generate_sql_query(column, checks)
+
+        checks_sql = "UNION ALL".join(_build_checks_sql())
 
         self.sql = f"SELECT col_name, check_type, check_result FROM ({checks_sql}) AS check_columns"
 

From 3140a02472de53b3a9bfa0b79be596e5c0bfa683 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Fri, 14 Oct 2022 11:28:09 -0400
Subject: [PATCH 62/63] Add new accept_none argument to column check operator

The new argument, defaulting to true, will convert Nones returned
from the query to 0s so numeric calculations can be performed
correctly. This allows empty tables to be handled as a row of zeroes.

Additional documentation is also supplied
---
 airflow/providers/common/sql/operators/sql.py         |  6 ++++++
 airflow/providers/google/cloud/operators/bigquery.py  | 11 ++++++++++-
 .../apache-airflow-providers-common-sql/operators.rst |  5 +++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index f408949932972..a1fd93fe06b42 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -272,6 +272,8 @@ class SQLColumnCheckOperator(BaseSQLOperator):
 
     :param conn_id: the connection ID used to connect to the database
     :param database: name of database which overwrite the defined one in connection
+    :param accept_none: whether or not to accept None values returned by the query. If true, converts None
+        to 0.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -301,6 +303,7 @@ def __init__(
         partition_clause: str | None = None,
         conn_id: str | None = None,
         database: str | None = None,
+        accept_none: bool = True,
         **kwargs,
     ):
         super().__init__(conn_id=conn_id, database=database, **kwargs)
@@ -308,6 +311,7 @@ def __init__(
         self.table = table
         self.column_mapping = column_mapping
         self.partition_clause = partition_clause
+        self.accept_none = accept_none
 
         def _build_checks_sql():
             for column, checks in self.column_mapping.items():
@@ -375,6 +379,8 @@ def _generate_partition_clause(check):
         return checks_sql
 
     def _get_match(self, check_values, record, tolerance=None) -> bool:
+        if record is None and self.accept_none:
+            record = 0
         match_boolean = True
         if "geq_to" in check_values:
             if tolerance is not None:
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 9e711d2678e90..cec55844b6244 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -543,6 +543,8 @@ def __init__(
         table: str,
         column_mapping: dict,
         partition_clause: str | None = None,
+        database: str | None = None,
+        accept_none: bool = True,
         gcp_conn_id: str = "google_cloud_default",
         use_legacy_sql: bool = True,
         location: str | None = None,
@@ -551,11 +553,18 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(
-            table=table, column_mapping=column_mapping, partition_clause=partition_clause, **kwargs
+            table=table,
+            column_mapping=column_mapping,
+            partition_clause=partition_clause,
+            database=database,
+            accept_none=accept_none,
+            **kwargs,
         )
         self.table = table
         self.column_mapping = column_mapping
         self.partition_clause = partition_clause
+        self.database = database
+        self.accept_none = accept_none
         self.gcp_conn_id = gcp_conn_id
         self.use_legacy_sql = use_legacy_sql
         self.location = location
diff --git a/docs/apache-airflow-providers-common-sql/operators.rst b/docs/apache-airflow-providers-common-sql/operators.rst
index 97bf9f4b6f6c1..bc725be418c03 100644
--- a/docs/apache-airflow-providers-common-sql/operators.rst
+++ b/docs/apache-airflow-providers-common-sql/operators.rst
@@ -94,6 +94,11 @@ The partition clauses may be given at the operator level as a parameter where it
 at the column level in the column mapping where it partitions all checks for that column, or at the
 check level for a column where it partitions just that check.
 
+A database may also be specified if not using the database from the supplied connection.
+
+The accept_none argument, true by default, will convert None values returned by the query to 0s, allowing
+empty tables to return valid integers.
+
 The below example demonstrates how to instantiate the SQLColumnCheckOperator task.
 
 .. exampleinclude:: /../../tests/system/providers/common/sql/example_sql_column_table_check.py

From 09c7b51bf39246114e1bb0ba21d181ff1b5919c0 Mon Sep 17 00:00:00 2001
From: Benji Lampel <benjamin@astronomer.io>
Date: Wed, 26 Oct 2022 10:04:23 -0400
Subject: [PATCH 63/63] Fix formatting issues after rebase

---
 airflow/providers/common/sql/operators/sql.py |  6 +--
 .../common/sql/operators/test_sql.py          | 50 +++++++++----------
 2 files changed, 27 insertions(+), 29 deletions(-)

diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py
index a1fd93fe06b42..9b3aa868dd8e8 100644
--- a/airflow/providers/common/sql/operators/sql.py
+++ b/airflow/providers/common/sql/operators/sql.py
@@ -19,9 +19,7 @@
 
 import ast
 import re
-from typing import TYPE_CHECKING, Any, Iterable, Mapping, NoReturn, Sequence, SupportsAbs
-
-from packaging.version import Version
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, NoReturn, Sequence, SupportsAbs
 
 from airflow.compat.functools import cached_property
 from airflow.exceptions import AirflowException, AirflowFailException
@@ -53,7 +51,7 @@ def _parse_boolean(val: str) -> str | bool:
     raise ValueError(f"{val!r} is not a boolean-like string value")
 
 
-_PROVIDERS_MATCHER = re.compile(r'airflow\.providers\.(.*)\.hooks.*')
+_PROVIDERS_MATCHER = re.compile(r"airflow\.providers\.(.*)\.hooks.*")
 
 _MIN_SUPPORTED_PROVIDERS_VERSION = {
     "amazon": "4.1.0",
diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py
index 7611713423544..2980326602935 100644
--- a/tests/providers/common/sql/operators/test_sql.py
+++ b/tests/providers/common/sql/operators/test_sql.py
@@ -163,11 +163,11 @@ def test_check_not_in_column_checks(self, monkeypatch):
 
     def test_pass_all_checks_exact_check(self, monkeypatch):
         records = [
-            ('X', 'null_check', 0),
-            ('X', 'distinct_check', 10),
-            ('X', 'unique_check', 10),
-            ('X', 'min', 1),
-            ('X', 'max', 19),
+            ("X", "null_check", 0),
+            ("X", "distinct_check", 10),
+            ("X", "unique_check", 10),
+            ("X", "min", 1),
+            ("X", "max", 19),
         ]
         operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
@@ -179,11 +179,11 @@ def test_pass_all_checks_exact_check(self, monkeypatch):
     def test_max_less_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
             records = [
-                ('X', 'null_check', 1),
-                ('X', 'distinct_check', 10),
-                ('X', 'unique_check', 10),
-                ('X', 'min', 1),
-                ('X', 'max', 21),
+                ("X", "null_check", 1),
+                ("X", "distinct_check", 10),
+                ("X", "unique_check", 10),
+                ("X", "min", 1),
+                ("X", "max", 21),
             ]
             operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
             operator.execute(context=MagicMock())
@@ -192,11 +192,11 @@ def test_max_less_than_fails_check(self, monkeypatch):
     def test_max_greater_than_fails_check(self, monkeypatch):
         with pytest.raises(AirflowException):
             records = [
-                ('X', 'null_check', 1),
-                ('X', 'distinct_check', 10),
-                ('X', 'unique_check', 10),
-                ('X', 'min', 1),
-                ('X', 'max', 9),
+                ("X", "null_check", 1),
+                ("X", "distinct_check", 10),
+                ("X", "unique_check", 10),
+                ("X", "min", 1),
+                ("X", "max", 9),
             ]
             operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
             operator.execute(context=MagicMock())
@@ -204,11 +204,11 @@ def test_max_greater_than_fails_check(self, monkeypatch):
 
     def test_pass_all_checks_inexact_check(self, monkeypatch):
         records = [
-            ('X', 'null_check', 0),
-            ('X', 'distinct_check', 9),
-            ('X', 'unique_check', 12),
-            ('X', 'min', 0),
-            ('X', 'max', 15),
+            ("X", "null_check", 0),
+            ("X", "distinct_check", 9),
+            ("X", "unique_check", 12),
+            ("X", "min", 0),
+            ("X", "max", 15),
         ]
         operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         operator.execute(context=MagicMock())
@@ -219,11 +219,11 @@ def test_pass_all_checks_inexact_check(self, monkeypatch):
 
     def test_fail_all_checks_check(self, monkeypatch):
         records = [
-            ('X', 'null_check', 1),
-            ('X', 'distinct_check', 12),
-            ('X', 'unique_check', 11),
-            ('X', 'min', -1),
-            ('X', 'max', 20),
+            ("X", "null_check", 1),
+            ("X", "distinct_check", 12),
+            ("X", "unique_check", 11),
+            ("X", "min", -1),
+            ("X", "max", 20),
         ]
         operator = operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
         with pytest.raises(AirflowException):
