Merge branch 'chip/weird_table_name_handler' into 'enterprise'

Added ID_SEPARTOR to test queries to support non-standard schema and table names

See merge request dkinternal/testgen/dataops-testgen!329

Author: ci bot

testgen/commands/queries/execute_cat_tests_query.py

@@ -17,7 +17,6 @@ class CATTestParams(TypedDict):
 class CCATExecutionSQL:
     project_code = ""
     flavor = ""
-    concat_operator = ""
     test_suite = ""
     run_date = ""
     test_run_id = ""
@@ -35,8 +34,7 @@ def __init__(self, strProjectCode, strTestSuiteId, strTestSuite, strSQLFlavor, m
         self.test_suite_id = strTestSuiteId
         self.test_suite = strTestSuite
         self.project_code = strProjectCode
-        flavor_service = get_flavor_service(strSQLFlavor)
-        self.concat_operator = flavor_service.get_concat_operator()
+        self.flavor_service = get_flavor_service(strSQLFlavor)
         self.flavor = strSQLFlavor
         self.max_query_chars = max_query_chars
         self.today = date_service.get_now_as_string_with_offset(minutes_offset)
@@ -58,9 +56,9 @@ def _get_query(self, template_file_name: str, sub_directory: str | None = "exec_
             "TEST_SUITE_ID": self.test_suite_id,
             "TABLE_GROUPS_ID": self.table_groups_id,
             "SQL_FLAVOR": self.flavor,
-            "ID_SEPARATOR": "`" if self.flavor in ("databricks", "bigquery") else '"',
-            "VARCHAR_TYPE": "STRING" if self.flavor == "bigquery" else "VARCHAR(1000)",
-            "CONCAT_OPERATOR": self.concat_operator,
+            "QUOTE": self.flavor_service.quote_character,
+            "VARCHAR_TYPE": self.flavor_service.varchar_type,
+            "CONCAT_OPERATOR": self.flavor_service.concat_operator,
             "SCHEMA_NAME": self.target_schema,
             "TABLE_NAME": self.target_table,
             "NOW_DATE": "GETDATE()",
@@ -74,7 +72,7 @@ def _get_query(self, template_file_name: str, sub_directory: str | None = "exec_
         query = replace_params(query, params)
         query = replace_templated_functions(query, self.flavor)
 
-        if no_bind and self.flavor != "databricks":
+        if no_bind:
             # Adding escape character where ':' is referenced
             query = query.replace(":", "\\:")
 

testgen/commands/queries/execute_tests_query.py

@@ -2,7 +2,7 @@
 
 from testgen.common import date_service, read_template_sql_file
 from testgen.common.clean_sql import CleanSQL, ConcatColumnList, quote_identifiers
-from testgen.common.database.database_service import replace_params
+from testgen.common.database.database_service import get_flavor_service, replace_params
 
 
 class TestParams(TypedDict):
@@ -55,6 +55,7 @@ class CTestExecutionSQL:
     def __init__(self, strProjectCode, strFlavor, strTestSuiteId, strTestSuite, minutes_offset=0):
         self.project_code = strProjectCode
         self.flavor = strFlavor
+        self.flavor_service = get_flavor_service(strFlavor)
         self.test_suite_id = strTestSuiteId
         self.test_suite = strTestSuite
         self.today = date_service.get_now_as_string_with_offset(minutes_offset)
@@ -101,13 +102,14 @@ def _get_query(
             "TEST_SUITE_ID": self.test_suite_id,
             "TEST_SUITE": self.test_suite,
             "SQL_FLAVOR": self.flavor,
+            "QUOTE": self.flavor_service.quote_character,
             "TEST_RUN_ID": self.test_run_id,
             "INPUT_PARAMETERS": self._get_input_parameters(),
             "RUN_DATE": self.run_date,
             "EXCEPTION_MESSAGE": self.exception_message,
             "START_TIME": self.today,
             "PROCESS_ID": self.process_id,
-            "VARCHAR_TYPE": "STRING" if self.flavor in ("databricks", "bigquery") else "VARCHAR",
+            "VARCHAR_TYPE": self.flavor_service.varchar_type,
             "NOW_TIMESTAMP": date_service.get_now_as_string_with_offset(self.minutes_offset),
             **{key.upper(): value or "" for key, value in self.test_params.items()},
         }
@@ -128,12 +130,12 @@ def _get_query(
 
             subset_condition = self.test_params["subset_condition"]
             params["SUBSET_DISPLAY"] = subset_condition.replace(
-                "'", r"\'" if self.flavor == "bigquery" else "''"
+                "'", self.flavor_service.escaped_single_quote
             ) if subset_condition else ""
 
         query = replace_params(query, params)
 
-        if no_bind and self.flavor != "databricks":
+        if no_bind:
             # Adding escape character where ':' is referenced
             query = query.replace(":", "\\:")
 

testgen/commands/queries/generate_tests_query.py

@@ -2,7 +2,7 @@
 from typing import ClassVar, TypedDict
 
 from testgen.common import CleanSQL, date_service, read_template_sql_file
-from testgen.common.database.database_service import replace_params
+from testgen.common.database.database_service import get_flavor_service, replace_params
 from testgen.common.read_file import get_template_files
 
 LOG = logging.getLogger("testgen")
@@ -29,7 +29,10 @@ class CDeriveTestsSQL:
 
     _use_clean = False
 
-    def __init__(self):
+    def __init__(self, flavor):
+        self.sql_flavor = flavor
+        self.flavor_service = get_flavor_service(flavor)
+
         today = date_service.get_now_as_string()
         self.run_date = today
         self.as_of_date = today
@@ -47,7 +50,7 @@ def _get_params(self) -> dict:
             "GENERATION_SET": self.generation_set,
             "AS_OF_DATE": self.as_of_date,
             "DATA_SCHEMA": self.data_schema,
-            "ID_SEPARATOR": "`" if self.sql_flavor in ("databricks", "bigquery") else '"',
+            "QUOTE": self.flavor_service.quote_character,
         }
 
     def _get_query(self, template_file_name: str, sub_directory: str | None = "generation") -> tuple[str, dict]:

testgen/commands/queries/refresh_data_chars_query.py

@@ -1,5 +1,5 @@
 from testgen.common import read_template_sql_file
-from testgen.common.database.database_service import replace_params
+from testgen.common.database.database_service import get_flavor_service, replace_params
 from testgen.common.database.flavor.flavor_service import SQLFlavor
 from testgen.utils import chunk_queries
 
@@ -46,36 +46,30 @@ def _get_query(self, template_file_name: str, sub_directory: str | None = "data_
 
     def _get_table_criteria(self) -> str:
         table_criteria = ""
-
-        table_ref = "c.table_name"
-        escape_clause = ""
-        escaped_underscore = "\\_"
-        if self.sql_flavor.startswith("mssql"):
-            escaped_underscore = "[_]"
-        elif self.sql_flavor == "snowflake":
-            escaped_underscore = "\\\\_"
-            escape_clause = "ESCAPE '\\\\'"
-        elif self.sql_flavor == "redshift":
-            escaped_underscore = "\\\\_"
-        elif self.sql_flavor == "redshift_spectrum":
-            table_ref = "c.tablename" 
-
+        flavor_service = get_flavor_service(self.sql_flavor)
+        
         if self.profiling_table_set:
-            table_criteria += f" AND {table_ref} IN ({self.profiling_table_set})"
+            table_criteria += f" AND c.{flavor_service.ddf_table_ref} IN ({self.profiling_table_set})"
 
         if self.profiling_include_mask:
-            include_table_names = [ item.strip().replace("_", escaped_underscore) for item in self.profiling_include_mask.split(",") ]
+            include_table_names = [
+                item.strip().replace("_", flavor_service.escaped_underscore)
+                for item in self.profiling_include_mask.split(",")
+            ]
             table_criteria += f"""
             AND (
-                {" OR ".join([ f"({table_ref} LIKE '{item}' {escape_clause})" for item in include_table_names ])}
+                {" OR ".join([ f"(c.{flavor_service.ddf_table_ref} LIKE '{item}' {flavor_service.escape_clause})" for item in include_table_names ])}
             )
             """
 
         if self.profiling_exclude_mask:
-            exclude_table_names = [ item.strip().replace("_", escaped_underscore) for item in self.profiling_exclude_mask.split(",") ]
+            exclude_table_names = [
+                item.strip().replace("_", flavor_service.escaped_underscore)
+                for item in self.profiling_exclude_mask.split(",")
+            ]
             table_criteria += f"""
             AND NOT (
-                {" OR ".join([ f"({table_ref} LIKE '{item}' {escape_clause})" for item in exclude_table_names ])}
+                {" OR ".join([ f"(c.{flavor_service.ddf_table_ref} LIKE '{item}' {flavor_service.escape_clause})" for item in exclude_table_names ])}
             )
             """
 

testgen/commands/queries/test_parameter_validation_query.py

@@ -1,7 +1,7 @@
 import typing
 
 from testgen.common import CleanSQL, date_service, read_template_sql_file
-from testgen.common.database.database_service import replace_params
+from testgen.common.database.database_service import get_flavor_service, replace_params
 
 
 class CTestParamValidationSQL:
@@ -19,6 +19,7 @@ class CTestParamValidationSQL:
 
     def __init__(self, strFlavor, strTestSuiteId):
         self.flavor = strFlavor
+        self.flavor_service = get_flavor_service(strFlavor)
         self.test_suite_id = strTestSuiteId
         self.today = date_service.get_now_as_string()
 
@@ -35,8 +36,8 @@ def _get_query(self, template_file_name: str, sub_directory: str | None = "valid
             "CAT_TEST_IDS": tuple(self.test_ids or []),
             "START_TIME": self.today,
             "NOW_TIMESTAMP": date_service.get_now_as_string(),
-            "COLUMNS_TABLE": f"{self.tg_schema}.INFORMATION_SCHEMA.COLUMNS" if self.flavor == "bigquery" else "information_schema.columns",
-            "ID_SEPARATOR": "`" if self.flavor in ("databricks", "bigquery") else '"',
+            "DATA_SCHEMA": self.tg_schema,
+            "QUOTE": self.flavor_service.quote_character,
         }
         query = replace_params(query, params)
         return query, params

testgen/commands/run_generate_tests.py

@@ -19,7 +19,7 @@ def run_test_gen_queries(table_group_id: str, test_suite: str, generation_set: s
     connection = Connection.get_by_table_group(table_group_id)
     set_target_db_params(connection.__dict__)
 
-    clsTests = CDeriveTestsSQL()
+    clsTests = CDeriveTestsSQL(connection.sql_flavor)
 
     LOG.info(f"CurrentStep: Retrieving General Parameters for Test Suite {test_suite}")
     params = get_test_generation_params(table_group_id, test_suite)
@@ -32,7 +32,6 @@ def run_test_gen_queries(table_group_id: str, test_suite: str, generation_set: s
     clsTests.test_suite_id = params["test_suite_id"] if params["test_suite_id"] else ""
     clsTests.connection_id = str(connection.connection_id)
     clsTests.table_groups_id = table_group_id
-    clsTests.sql_flavor = connection.sql_flavor
     clsTests.data_schema = params["table_group_schema"]
     if params["profiling_as_of_date"] is not None:
         clsTests.as_of_date = params["profiling_as_of_date"].strftime("%Y-%m-%d %H:%M:%S")

testgen/commands/run_refresh_data_chars.py

@@ -7,6 +7,7 @@
     execute_db_queries,
     fetch_dict_from_db,
     fetch_from_db_threaded,
+    get_flavor_service,
     write_to_app_db,
 )
 from testgen.common.get_pipeline_parms import TestExecutionParams
@@ -18,12 +19,14 @@
 def run_refresh_data_chars_queries(params: TestExecutionParams, run_date: str, spinner: Spinner=None):
     LOG.info("CurrentStep: Initializing Data Characteristics Refresh")
     sql_generator = CRefreshDataCharsSQL(params, run_date, STAGING_TABLE)
+    flavor_service = get_flavor_service(params["sql_flavor"])
+    quote = flavor_service.quote_character
 
     LOG.info("CurrentStep: Getting DDF for table group")
     ddf_results = fetch_dict_from_db(*sql_generator.GetDDFQuery(), use_target_db=True)
 
     distinct_tables = {
-        f"{item['table_schema']}.{item['table_name']}"
+        f"{quote}{item['table_schema']}{quote}.{quote}{item['table_name']}{quote}"
         for item in ddf_results
     }
     if distinct_tables:
@@ -65,7 +68,7 @@ def run_refresh_data_chars_queries(params: TestExecutionParams, run_date: str, s
             item["general_type"],
             item["column_type"],
             item["db_data_type"],
-            count_map.get(f"{item['table_schema']}.{item['table_name']}", 0),
+            count_map.get(f"{quote}{item['table_schema']}{quote}.{quote}{item['table_name']}{quote}", 0),
         ]
         for item in ddf_results
     ]

testgen/common/clean_sql.py

@@ -1,5 +1,7 @@
 import re
 
+from testgen.common.database.database_service import get_flavor_service
+
 
 def CleanSQL(strInput: str) -> str:
     # Use regular expression to remove comment text fenced by /*...*/
@@ -24,20 +26,22 @@ def quote_identifiers(identifiers: str, flavor: str) -> str:
         "by",
         "having",
     ]
+    flavor_service = get_flavor_service(flavor)
+    quote = flavor_service.quote_character
 
     quoted_values = []
     for value in identifiers.split(","):
         value = value.strip()
-        if value.startswith('"') and value.endswith('"'):
+        if value.startswith(quote) and value.endswith(quote):
             quoted_values.append(value)
         elif any(
-            (flavor == "snowflake" and c.lower())
-            or (flavor != "snowflake" and c.isupper())
+            (flavor_service.default_uppercase and c.lower())
+            or (not flavor_service.default_uppercase and c.isupper())
             or c.isspace()
             or value.lower() in keywords
             for c in value
         ):
-            quoted_values.append(f'"{value}"')
+            quoted_values.append(f"{quote}{value}{quote}")
         else:
             quoted_values.append(value)
     return ", ".join(quoted_values)

testgen/common/database/flavor/bigquery_flavor_service.py

@@ -4,6 +4,11 @@
 
 
 class BigqueryFlavorService(FlavorService):
+
+    quote_character = "`"
+    escaped_single_quote = "\\'"
+    varchar_type = "STRING"
+
     def get_connection_string_head(self):
         return "bigquery://"
 

testgen/common/database/flavor/databricks_flavor_service.py

@@ -5,6 +5,10 @@
 
 class DatabricksFlavorService(FlavorService):
 
+    quote_character = "`"
+    escaped_single_quote = "\\'"
+    varchar_type = "STRING"
+
     def get_connection_string_head(self):
         return f"{self.flavor}://{self.username}:{quote_plus(self.password)}@"