vdk-impala: Introduce checks for insert template (#2198)

sbuldeev · pre-commit-ci[bot] · web-flow · commit b869ddbe5b77 · 2023-06-15T14:27:49.000+03:00
Why: Add functionality to the insert processing templates in order to allow quality checks to be made before the data is inserted into the target table. Currently, the checks done on the insert template processing step are not covering if the semantics of the data is correct. Therefore, bad data could go into the target table which could be unwanted behavior. More details explained in #1361 What: -Adding functionality to handle the insert template behavior if the user provides checks Tests: provided positive and negative regression tests Signed-off-by: Stefan Buldeev sbuldeev@vmware.com --------- Signed-off-by: Stefan Buldeev sbuldeev@vmware.com Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/data_quality_exception.py b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/data_quality_exception.py
@@ -0,0 +1,23 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+
+class DataQualityException(Exception):
+    """
+    Exception raised for errors with the quality of the data.
+
+    Attributes:
+        checked_object -- Object that the quality checks are ran against
+        target_table -- DWH table where target data is loaded
+        source_view -- View from which the raw data is loaded from
+    """
+
+    def __init__(self, checked_object, target_table, source_view):
+        self.checked_object = checked_object
+        self.target_table = target_table
+        self.source_view = source_view
+        self.message = f"""What happened: Error occurred while performing quality checks.\n
+        Why it happened: Object: {checked_object} is not passing the quality checks.\n
+        Consequences: The source view data will not be processed to the target table - {target_table}.\n
+        Countermeasures: Check the source view: {source_view} what data is trying to be processed."""
+        super().__init__(self.message)
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/dimension/scd1/02-handle-quality-checks.py b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/dimension/scd1/02-handle-quality-checks.py
@@ -3,6 +3,7 @@
 import os
 
 from vdk.api.job_input import IJobInput
+from vdk.plugin.impala.templates.data_quality_exception import DataQualityException
 from vdk.plugin.impala.templates.utility import align_stg_table_with_target
 from vdk.plugin.impala.templates.utility import get_file_content
 from vdk.plugin.impala.templates.utility import get_staging_table_name
@@ -34,10 +35,9 @@ def run(job_input: IJobInput):
         staging_table_name = get_staging_table_name(target_schema, target_table)
 
         staging_table = f"{staging_schema}.{staging_table_name}"
+        target_table_full_name = f"{target_schema}.{target_table}"
 
-        align_stg_table_with_target(
-            f"{target_schema}.{target_table}", staging_table, job_input
-        )
+        align_stg_table_with_target(target_table_full_name, staging_table, job_input)
 
         insert_into_staging = insert_query.format(
             target_schema=staging_schema,
@@ -58,7 +58,11 @@ def run(job_input: IJobInput):
             )
             job_input.execute_query(insert_into_target)
         else:
-            raise Exception("The data is not passing the quality checks!")
+            raise DataQualityException(
+                checked_object=staging_table,
+                source_view=f"{source_schema}.{source_view}",
+                target_table=target_table_full_name,
+            )
 
     else:
         job_input.execute_query(insert_query)
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/dimension/scd1/README.md b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/dimension/scd1/README.md
@@ -16,7 +16,7 @@ In summary, it overwrites the target table with the source data.
 - target_table         - SC Data Warehouse table of DW type 'Slowly Changing Dimension Type 1', where target data is loaded
 - source_schema        - SC Data Lake schema, where source raw data is loaded from
 - source_view          - SC Data Lake view, where source raw data is loaded from
-- check                - (Optional) Callback function responsible for checking the quality of the data
+- check                - (Optional) Callback function responsible for checking the quality of the data. Takes in a table name as a parameter which will be used for data validation
 - staging_schema       - (Optional) Schema where the checks will be executed. If not provided target_schema will be used as default
 
 ### Prerequisites:
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/00-fact-snapshot-definition.py b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/00-fact-snapshot-definition.py
@@ -1,5 +1,8 @@
 # Copyright 2021-2023 VMware, Inc.
 # SPDX-License-Identifier: Apache-2.0
+from typing import Callable
+from typing import Optional
+
 from pydantic import BaseModel
 from vdk.api.job_input import IJobInput
 from vdk.plugin.impala.templates.template_arguments_validator import (
@@ -12,6 +15,8 @@ class FactDailySnapshotParams(BaseModel):
     target_table: str
     source_schema: str
     source_view: str
+    check: Optional[Callable[[str], bool]]
+    staging_schema: Optional[str]
 
 
 class FactDailySnapshot(TemplateArgumentsValidator):
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/02-handle-quality-checks.py b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/02-handle-quality-checks.py
@@ -0,0 +1,84 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+from vdk.api.job_input import IJobInput
+from vdk.plugin.impala.templates.data_quality_exception import DataQualityException
+from vdk.plugin.impala.templates.utility import align_stg_table_with_target
+from vdk.plugin.impala.templates.utility import get_file_content
+from vdk.plugin.impala.templates.utility import get_staging_table_name
+
+SQL_FILES_FOLDER = (
+    os.path.dirname(os.path.abspath(__file__)) + "/02-requisite-sql-scripts"
+)
+
+"""
+This step is intened to handle quality checks if such provided
+and stop the data from being populated into the target table if the check has negative outcome.
+Otherwise the data will be directly processed according to the used template type
+"""
+
+
+def run(job_input: IJobInput):
+    job_arguments = job_input.get_arguments()
+
+    check = job_arguments.get("check")
+    partition_clause = job_arguments["_vdk_template_insert_partition_clause"]
+    source_schema = job_arguments.get("source_schema")
+    source_view = job_arguments.get("source_view")
+    target_schema = job_arguments.get("target_schema")
+    target_table = job_arguments.get("target_table")
+    insert_query = get_file_content(SQL_FILES_FOLDER, "02-insert-into-target.sql")
+
+    if check:
+        staging_schema = job_arguments.get("staging_schema", target_schema)
+        staging_table_name = get_staging_table_name(target_schema, target_table)
+
+        staging_table = f"{staging_schema}.{staging_table_name}"
+        target_table_full_name = f"{target_schema}.{target_table}"
+
+        align_stg_table_with_target(target_table_full_name, staging_table, job_input)
+
+        insert_into_staging = insert_query.format(
+            target_schema=staging_schema,
+            target_table=staging_table_name,
+            _vdk_template_insert_partition_clause=partition_clause,
+            source_schema=source_schema,
+            source_view=source_view,
+        )
+        job_input.execute_query(insert_into_staging)
+
+        view_schema = staging_schema
+        view_name = f"vw_{staging_table_name}"
+        create_view_query = get_file_content(
+            SQL_FILES_FOLDER, "02-create-consolidated-view.sql"
+        )
+        create_view = create_view_query.format(
+            view_schema=view_schema,
+            view_name=view_name,
+            target_schema=target_schema,
+            target_table=target_table,
+            staging_schema=staging_schema,
+            staging_table_name=staging_table_name,
+        )
+        job_input.execute_query(create_view)
+
+        view_full_name = f"{view_schema}.{view_name}"
+        if check(view_full_name):
+            insert_into_target = insert_query.format(
+                source_schema=staging_schema,
+                source_view=staging_table_name,
+                _vdk_template_insert_partition_clause=partition_clause,
+                target_schema=target_schema,
+                target_table=target_table,
+            )
+            job_input.execute_query(insert_into_target)
+        else:
+            raise DataQualityException(
+                checked_object=view_full_name,
+                source_view=f"{source_schema}.{source_view}",
+                target_table=target_table_full_name,
+            )
+
+    else:
+        job_input.execute_query(insert_query)
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/02-requisite-sql-scripts/02-create-consolidated-view.sql b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/02-requisite-sql-scripts/02-create-consolidated-view.sql
@@ -0,0 +1,8 @@
+CREATE VIEW IF NOT EXISTS {view_schema}.{view_name}
+AS (
+    SELECT *
+    FROM {target_schema}.{target_table}
+    UNION ALL
+    SELECT *
+    FROM {staging_schema}.{staging_table_name}
+)
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/02-requisite-sql-scripts/02-insert-into-target.sql b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/02-requisite-sql-scripts/02-insert-into-target.sql
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/README.md b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/insert/README.md
@@ -12,6 +12,8 @@ This template can be used to load raw data from Data Lake to target Table in Dat
 - target_table    - Data Warehouse table of DW type table, where target data is loaded
 - source_schema   - Data Lake schema, where source raw data is loaded from
 - source_view     - Data Lake view, where source raw data is loaded from
+- check           - (Optional) Callback function responsible for checking the quality of the data. Takes in a table name as a parameter which will be used for data validation
+- staging_schema  - (Optional) Schema where the checks will be executed. If not provided target_schema will be used as default
 
 ### Prerequisites:
 
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/snapshot/02-handle-quality-checks.py b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/snapshot/02-handle-quality-checks.py
@@ -3,6 +3,7 @@
 import os
 
 from vdk.api.job_input import IJobInput
+from vdk.plugin.impala.templates.data_quality_exception import DataQualityException
 from vdk.plugin.impala.templates.utility import align_stg_table_with_target
 from vdk.plugin.impala.templates.utility import get_file_content
 from vdk.plugin.impala.templates.utility import get_staging_table_name
@@ -38,10 +39,9 @@ def run(job_input: IJobInput):
         staging_table_name = get_staging_table_name(target_schema, target_table)
 
         staging_table = f"{staging_schema}.{staging_table_name}"
+        target_table_full_name = f"{target_schema}.{target_table}"
 
-        align_stg_table_with_target(
-            f"{target_schema}.{target_table}", staging_table, job_input
-        )
+        align_stg_table_with_target(target_table_full_name, staging_table, job_input)
 
         insert_into_staging = insert_query.format(
             current_target_schema=staging_schema,
@@ -65,7 +65,11 @@ def run(job_input: IJobInput):
             )
             job_input.execute_query(insert_into_target)
         else:
-            raise Exception("The data is not passing the quality checks!")
+            raise DataQualityException(
+                checked_object=staging_table,
+                source_view=f"{source_schema}.{source_view}",
+                target_table=target_table_full_name,
+            )
 
     else:
         insert_query = insert_query.replace(
diff --git a/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/snapshot/README.md b/projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/templates/load/fact/snapshot/README.md
@@ -19,7 +19,7 @@ truncating all present target table records observed after t1.
 - source_schema   - SC Data Lake schema, where source raw data is loaded from
 - source_view     - SC Data Lake view, where source raw data is loaded from
 - last_arrival_ts - Timestamp column, on which increments to target_table are done
-- check           - (Optional) Callback function responsible for checking the quality of the data
+- check           - (Optional) Callback function responsible for checking the quality of the data. Takes in a table name as a parameter which will be used for data validation
 - staging_schema  - (Optional) Schema where the checks will be executed. If not provided target_schema will be used as default
 
 ### Prerequisites:
diff --git a/projects/vdk-plugins/vdk-impala/tests/functional/jobs/insert_template_job/02_run_insert_template.py b/projects/vdk-plugins/vdk-impala/tests/functional/jobs/insert_template_job/02_run_insert_template.py
@@ -1,5 +1,6 @@
 # Copyright 2021-2023 VMware, Inc.
 # SPDX-License-Identifier: Apache-2.0
+from functional import test_utility
 from vdk.api.job_input import IJobInput
 
 __author__ = "VMware, Inc."
@@ -9,7 +10,13 @@
 
 
 def run(job_input: IJobInput) -> None:
+    template_args = job_input.get_arguments()
+    check = template_args.get("check")
+    if check:
+        check = test_utility.setup_testing_check(check)
+        template_args["check"] = check
+
     job_input.execute_template(
         template_name="insert",
-        template_args=job_input.get_arguments(),
+        template_args=template_args,
     )
diff --git a/projects/vdk-plugins/vdk-impala/tests/functional/template_regression_test.py b/projects/vdk-plugins/vdk-impala/tests/functional/template_regression_test.py
@@ -766,3 +766,63 @@ def test_insert_partition(self) -> None:
         self.assertSetEqual(
             actual, expected, f"Elements in {expect_table} and {target_table} differ."
         )
+
+    def test_insert_checks_positive(self) -> None:
+        test_schema = "vdkprototypes"
+        staging_schema = "staging_vdkprototypes"
+        source_view = "vw_fact_vmc_utilization_cpu_mem_every5min_daily_check_positive"
+        target_table = "dw_fact_vmc_utilization_cpu_mem_every5min_daily_check_positive"
+        expect_table = "ex_fact_vmc_utilization_cpu_mem_every5min_daily_check_positive"
+
+        res = self._run_job(
+            "insert_template_job",
+            {
+                "source_schema": test_schema,
+                "source_view": source_view,
+                "target_schema": test_schema,
+                "target_table": target_table,
+                "expect_schema": test_schema,
+                "expect_table": expect_table,
+                "check": "use_positive_check",
+                "staging_schema": staging_schema,
+            },
+        )
+
+        assert not res.exception
+        actual_rs = self._run_query(f"SELECT * FROM {test_schema}.{target_table}")
+        expected_rs = self._run_query(f"SELECT * FROM {test_schema}.{expect_table}")
+        assert actual_rs.output and expected_rs.output
+
+        actual = {x for x in actual_rs.output.split("\n")}
+        expected = {x for x in expected_rs.output.split("\n")}
+
+        self.assertSetEqual(
+            actual, expected, f"Elements in {expect_table} and {target_table} differ."
+        )
+
+    def test_insert_checks_negative(self) -> None:
+        test_schema = "vdkprototypes"
+        staging_schema = "staging_vdkprototypes"
+        source_view = "vw_fact_vmc_utilization_cpu_mem_every5min_daily_check_negative"
+        target_table = "dw_fact_vmc_utilization_cpu_mem_every5min_daily_check_negative"
+        expect_table = "ex_fact_vmc_utilization_cpu_mem_every5min_daily_check_negative"
+
+        res = self._run_job(
+            "insert_template_job",
+            {
+                "source_schema": test_schema,
+                "source_view": source_view,
+                "target_schema": test_schema,
+                "target_table": target_table,
+                "expect_schema": test_schema,
+                "expect_table": expect_table,
+                "check": "use_negative_check",
+                "staging_schema": staging_schema,
+            },
+        )
+
+        assert res.exception
+        actual_rs = self._run_query(f"SELECT * FROM {test_schema}.{target_table}")
+        expected_rs = self._run_query(f"SELECT * FROM {test_schema}.{expect_table}")
+        assert actual_rs.output and expected_rs.output
+        assert actual_rs.output != expected_rs.output
