Merge pull request #47 from DataKitchen/release/4.20.4

Release/4.20.4

Author: datakitchen-devops

pyproject.toml

@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dataops-testgen"
-version = "4.16.3"
+version = "4.20.4"
 description = "DataKitchen's Data Quality DataOps TestGen"
 authors = [
     { "name" = "DataKitchen, Inc.", "email" = "[email protected]" },

testgen/__main__.py

@@ -41,8 +41,10 @@
     get_tg_schema,
     version_service,
 )
+from testgen.common.models import with_database_session
+from testgen.common.models.profiling_run import ProfilingRun
+from testgen.common.models.test_run import TestRun
 from testgen.scheduler import register_scheduler_job, run_scheduler
-from testgen.ui.queries import profiling_run_queries, test_run_queries
 from testgen.utils import plugins
 
 LOG = logging.getLogger("testgen")
@@ -72,9 +74,9 @@ def invoke(self, ctx: Context):
     cls=CliGroup,
     help=f"""
     {VERSION_DATA.edition} {VERSION_DATA.current or ""}
-    
+
     {f"New version available! {VERSION_DATA.latest}" if VERSION_DATA.latest != VERSION_DATA.current else ""}
-    
+
     Schema revision: {get_schema_revision()}
     """
 )
@@ -625,11 +627,16 @@ def run_ui():
     use_ssl = os.path.isfile(settings.SSL_CERT_FILE) and os.path.isfile(settings.SSL_KEY_FILE)
 
     patch_streamlit.patch(force=True)
-    try:
-        profiling_run_queries.cancel_all_running()
-        test_run_queries.cancel_all_running()
-    except Exception:
-        LOG.warning("Failed to cancel 'Running' profiling/test runs")
+
+    @with_database_session
+    def cancel_all_running():
+        try:
+            ProfilingRun.cancel_all_running()
+            TestRun.cancel_all_running()
+        except Exception:
+            LOG.warning("Failed to cancel 'Running' profiling/test runs")
+
+    cancel_all_running()
 
     try:
         app_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ui/app.py")

testgen/commands/queries/execute_cat_tests_query.py

@@ -1,11 +1,19 @@
-import typing
+from typing import ClassVar, TypedDict
 
 from testgen.commands.queries.rollup_scores_query import CRollupScoresSQL
 from testgen.common import date_service, read_template_sql_file
-from testgen.common.database import database_service
+from testgen.common.database.database_service import get_flavor_service, replace_params
 from testgen.common.read_file import replace_templated_functions
 
 
+class CATTestParams(TypedDict):
+    schema_name: str
+    table_name: str
+    cat_sequence: int
+    test_measures: str
+    test_conditions: str
+
+
 class CCATExecutionSQL:
     project_code = ""
     flavor = ""
@@ -16,11 +24,9 @@ class CCATExecutionSQL:
     table_groups_id = ""
     max_query_chars = ""
     exception_message = ""
-
-    # Test Set Parameters
     target_schema = ""
     target_table = ""
-    dctTestParms: typing.ClassVar = {}
+    cat_test_params: ClassVar[CATTestParams] = {}
 
     _rollup_scores_sql: CRollupScoresSQL = None
 
@@ -29,7 +35,7 @@ def __init__(self, strProjectCode, strTestSuiteId, strTestSuite, strSQLFlavor, m
         self.test_suite_id = strTestSuiteId
         self.test_suite = strTestSuite
         self.project_code = strProjectCode
-        flavor_service = database_service.get_flavor_service(strSQLFlavor)
+        flavor_service = get_flavor_service(strSQLFlavor)
         self.concat_operator = flavor_service.get_concat_operator()
         self.flavor = strSQLFlavor
         self.max_query_chars = max_query_chars
@@ -41,83 +47,78 @@ def _get_rollup_scores_sql(self) -> CRollupScoresSQL:
             self._rollup_scores_sql = CRollupScoresSQL(self.test_run_id, self.table_groups_id)
 
         return self._rollup_scores_sql
-
-    def _ReplaceParms(self, strInputString):
-        strInputString = strInputString.replace("{MAX_QUERY_CHARS}", str(self.max_query_chars))
-        strInputString = strInputString.replace("{TEST_RUN_ID}", self.test_run_id)
-        strInputString = strInputString.replace("{PROJECT_CODE}", self.project_code)
-        strInputString = strInputString.replace("{TEST_SUITE}", self.test_suite)
-        strInputString = strInputString.replace("{TEST_SUITE_ID}", self.test_suite_id)
-        strInputString = strInputString.replace("{TABLE_GROUPS_ID}", self.table_groups_id)
-
-        strInputString = strInputString.replace("{SQL_FLAVOR}", self.flavor)
-        strInputString = strInputString.replace("{ID_SEPARATOR}", "`" if self.flavor == "databricks" else '"')
-        strInputString = strInputString.replace("{CONCAT_OPERATOR}", self.concat_operator)
-
-        strInputString = strInputString.replace("{SCHEMA_NAME}", self.target_schema)
-        strInputString = strInputString.replace("{TABLE_NAME}", self.target_table)
-
-        strInputString = strInputString.replace("{RUN_DATE}", self.run_date)
-        strInputString = strInputString.replace("{NOW_DATE}", "GETDATE()")
-        strInputString = strInputString.replace("{START_TIME}", self.today)
-        strInputString = strInputString.replace(
-            "{NOW}", date_service.get_now_as_string_with_offset(self.minutes_offset)
-        )
-        strInputString = strInputString.replace("{EXCEPTION_MESSAGE}", self.exception_message.strip())
-
-        for parm, value in self.dctTestParms.items():
-            strInputString = strInputString.replace("{" + parm.upper() + "}", str(value))
-
-        strInputString = strInputString.replace("{RUN_DATE}", self.run_date)
-
-        strInputString = replace_templated_functions(strInputString, self.flavor)
-
-        if self.flavor != "databricks":
+    
+    def _get_query(self, template_file_name: str, sub_directory: str | None = "exec_cat_tests", no_bind: bool = False) -> tuple[str, dict | None]:
+        query = read_template_sql_file(template_file_name, sub_directory)
+        params = {
+            "MAX_QUERY_CHARS": self.max_query_chars,
+            "TEST_RUN_ID": self.test_run_id,
+            "PROJECT_CODE": self.project_code,
+            "TEST_SUITE": self.test_suite,
+            "TEST_SUITE_ID": self.test_suite_id,
+            "TABLE_GROUPS_ID": self.table_groups_id,
+            "SQL_FLAVOR": self.flavor,
+            "ID_SEPARATOR": "`" if self.flavor == "databricks" else '"',
+            "CONCAT_OPERATOR": self.concat_operator,
+            "SCHEMA_NAME": self.target_schema,
+            "TABLE_NAME": self.target_table,
+            "NOW_DATE": "GETDATE()",
+            "START_TIME": self.today,
+            "NOW_TIMESTAMP": date_service.get_now_as_string_with_offset(self.minutes_offset),
+            "EXCEPTION_MESSAGE": self.exception_message.strip(),
+            **{key.upper(): value for key, value in self.cat_test_params.items()},
+            # This has to be replaced at the end
+            "RUN_DATE": self.run_date,
+        }
+        query = replace_params(query, params)
+        query = replace_templated_functions(query, self.flavor)
+
+        if no_bind and self.flavor != "databricks":
             # Adding escape character where ':' is referenced
-            strInputString = strInputString.replace(":", "\\:")
+            query = query.replace(":", "\\:")
 
-        return strInputString
+        return query, None if no_bind else params
 
-    def GetDistinctTablesSQL(self):
-        # Runs on DK DB
-        strQ = self._ReplaceParms(read_template_sql_file("ex_cat_get_distinct_tables.sql", "exec_cat_tests"))
-        return strQ
+    def GetDistinctTablesSQL(self) -> tuple[str, dict]:
+        # Runs on App database
+        return self._get_query("ex_cat_get_distinct_tables.sql")
 
-    def GetAggregateTableTestSQL(self):
-        # Runs on DK DB
-        strQ = self._ReplaceParms(read_template_sql_file("ex_cat_build_agg_table_tests.sql", "exec_cat_tests"))
-        return strQ
+    def GetAggregateTableTestSQL(self) -> tuple[str, None]:
+        # Runs on App database
+        return self._get_query("ex_cat_build_agg_table_tests.sql", no_bind=True)
 
-    def GetAggregateTestParmsSQL(self):
-        # Runs on DK DB
-        strQ = self._ReplaceParms(read_template_sql_file("ex_cat_retrieve_agg_test_parms.sql", "exec_cat_tests"))
-        return strQ
+    def GetAggregateTestParmsSQL(self) -> tuple[str, dict]:
+        # Runs on App database
+        return self._get_query("ex_cat_retrieve_agg_test_parms.sql")
 
-    def PrepCATQuerySQL(self):
-        strQ = self._ReplaceParms(read_template_sql_file("ex_cat_test_query.sql", "exec_cat_tests"))
-        return strQ
+    def PrepCATQuerySQL(self) -> tuple[str, None]:
+        # Runs on Target database
+        return self._get_query("ex_cat_test_query.sql", no_bind=True)
 
-    def GetCATResultsParseSQL(self):
-        strQ = self._ReplaceParms(read_template_sql_file("ex_cat_results_parse.sql", "exec_cat_tests"))
-        return strQ
+    def GetCATResultsParseSQL(self) -> tuple[str, dict]:
+        # Runs on App database
+        return self._get_query("ex_cat_results_parse.sql")
 
-    def FinalizeTestResultsSQL(self):
-        strQ = self._ReplaceParms(read_template_sql_file("ex_finalize_test_run_results.sql", "execution"))
-        return strQ
+    def FinalizeTestResultsSQL(self) -> tuple[str, dict]:
+        # Runs on App database
+        return self._get_query("ex_finalize_test_run_results.sql", "execution")
 
-    def PushTestRunStatusUpdateSQL(self):
-        strQ = self._ReplaceParms(read_template_sql_file("ex_update_test_record_in_testrun_table.sql", "execution"))
-        return strQ
+    def PushTestRunStatusUpdateSQL(self) -> tuple[str, dict]:
+        # Runs on App database
+        return self._get_query("ex_update_test_record_in_testrun_table.sql", "execution")
 
-    def FinalizeTestSuiteUpdateSQL(self):
-        strQ = self._ReplaceParms(read_template_sql_file("ex_update_test_suite.sql", "execution"))
-        return strQ
+    def FinalizeTestSuiteUpdateSQL(self) -> tuple[str, dict]:
+        # Runs on App database
+        return self._get_query("ex_update_test_suite.sql", "execution")
 
-    def CalcPrevalenceTestResultsSQL(self):
-        return self._ReplaceParms(read_template_sql_file("ex_calc_prevalence_test_results.sql", "execution"))
+    def CalcPrevalenceTestResultsSQL(self) -> tuple[str, None]:
+        # Runs on App database
+        return self._get_query("ex_calc_prevalence_test_results.sql", "execution", no_bind=True)
 
-    def TestScoringRollupRunSQL(self):
+    def TestScoringRollupRunSQL(self) -> tuple[str, dict]:
+        # Runs on App database
         return self._get_rollup_scores_sql().GetRollupScoresTestRunQuery()
 
-    def TestScoringRollupTableGroupSQL(self):
+    def TestScoringRollupTableGroupSQL(self) -> tuple[str, dict]:
+        # Runs on App database
         return self._get_rollup_scores_sql().GetRollupScoresTestTableGroupQuery()