Merge branch 'misc-fixes' into 'enterprise'

feat: skip view sampling on unsupported flavors + capture object_type

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

Author: ci bot

testgen/commands/queries/execute_tests_query.py

@@ -7,7 +7,7 @@
 import pandas as pd
 
 from testgen.common import read_template_sql_file
-from testgen.common.clean_sql import concat_columns
+from testgen.common.clean_sql import concat_columns, null_if_empty
 from testgen.common.database.database_service import (
     fetch_dict_from_db,
     get_flavor_service,
@@ -348,16 +348,27 @@ def _get_params(self, test_def: TestExecutionDef | None = None) -> dict:
                 "CONCAT_COLUMNS": concat_columns(test_def.column_name, self.null_value) if test_def.column_name else "",
                 "SKIP_ERRORS": test_def.skip_errors or 0,
                 "CUSTOM_QUERY": test_def.custom_query,
-                "BASELINE_CT": test_def.baseline_ct,
-                "BASELINE_UNIQUE_CT": test_def.baseline_unique_ct,
+                # Numeric baseline params feed division/arithmetic in CAT measures.
+                # Render empty values as NULL so the measure evaluates to NULL instead of producing
+                # invalid SQL like CAST( AS FLOAT). Cannot default to 0 — several are denominators.
+                "BASELINE_CT": null_if_empty(test_def.baseline_ct),
+                "BASELINE_UNIQUE_CT": null_if_empty(test_def.baseline_unique_ct),
+                # BASELINE_VALUE is a quoted literal, an unquoted number, or an IN-list depending on
+                # test type — no single empty-substitution is valid SQL, so it is templated as-is.
                 "BASELINE_VALUE": test_def.baseline_value,
-                "BASELINE_VALUE_CT": test_def.baseline_value_ct,
+                "BASELINE_VALUE_CT": null_if_empty(test_def.baseline_value_ct),
                 "THRESHOLD_VALUE": test_def.threshold_value or 0,
-                "BASELINE_SUM": test_def.baseline_sum,
-                "BASELINE_AVG": test_def.baseline_avg,
-                "BASELINE_SD": test_def.baseline_sd,
-                "LOWER_TOLERANCE": "NULL" if test_def.lower_tolerance in (None, "") else test_def.lower_tolerance,
-                "UPPER_TOLERANCE": "NULL" if test_def.upper_tolerance in (None, "") else test_def.upper_tolerance,
+                # Freshness_Trend uses BASELINE_SUM as a quoted timestamp (the template's
+                # NULLIF('{...}', '') handles empty); Other tests use it as a numeric sum.
+                "BASELINE_SUM": (
+                    test_def.baseline_sum
+                    if test_def.test_type == "Freshness_Trend"
+                    else null_if_empty(test_def.baseline_sum)
+                ),
+                "BASELINE_AVG": null_if_empty(test_def.baseline_avg),
+                "BASELINE_SD": null_if_empty(test_def.baseline_sd),
+                "LOWER_TOLERANCE": null_if_empty(test_def.lower_tolerance),
+                "UPPER_TOLERANCE": null_if_empty(test_def.upper_tolerance),
                 # SUBSET_CONDITION should be replaced after CUSTOM_QUERY
                 # since the latter may contain the former
                 "SUBSET_CONDITION": test_def.subset_condition or "1=1",

testgen/commands/queries/refresh_data_chars_query.py

@@ -30,6 +30,7 @@ class RefreshDataCharsSQL:
         "general_type",
         "column_type",
         "db_data_type",
+        "object_type",
         "approx_record_ct",
         "record_ct",
     )
@@ -157,6 +158,7 @@ def get_staging_data_chars(self, data_chars: list[ColumnChars], run_date: dateti
                 column.general_type,
                 column.column_type,
                 column.db_data_type,
+                column.object_type,
                 column.approx_record_ct,
                 column.record_ct,
             ]

testgen/commands/run_profiling.py

@@ -19,7 +19,7 @@
     write_to_app_db,
 )
 from testgen.common.database.column_chars import ColumnChars
-from testgen.common.database.database_service import ThreadedProgress
+from testgen.common.database.database_service import ThreadedProgress, get_flavor_service
 from testgen.common.job_context import job_context
 from testgen.common.mixpanel_service import MixpanelService
 from testgen.common.models import get_current_session, with_database_session
@@ -86,8 +86,9 @@ def run_profiling(
         if data_chars:
             sql_generator = ProfilingSQL(connection, table_group, profiling_run)
 
-            _run_column_profiling(sql_generator, data_chars)
-            _run_frequency_analysis(sql_generator)
+            sampling_params = _compute_sampling_params(sql_generator, data_chars)
+            _run_column_profiling(sql_generator, data_chars, sampling_params)
+            _run_frequency_analysis(sql_generator, sampling_params)
             _run_hygiene_issue_detection(sql_generator)
 
             # if table_group.profile_do_pair_rules == "Y":
@@ -143,26 +144,40 @@ def _exclude_xde_columns(data_chars: list[ColumnChars], table_group_id: UUID) ->
     return filtered
 
 
-def _run_column_profiling(sql_generator: ProfilingSQL, data_chars: list[ColumnChars]) -> None:
+def _compute_sampling_params(
+    sql_generator: ProfilingSQL, data_chars: list[ColumnChars]
+) -> dict[str, TableSampling]:
+    table_group = sql_generator.table_group
+    sampling_params: dict[str, TableSampling] = {}
+    if not table_group.profile_use_sampling:
+        return sampling_params
+
+    sampleable_types = get_flavor_service(sql_generator.flavor).sampleable_object_types
+    for column in data_chars:
+        if sampling_params.get(column.table_name):
+            continue
+        if sampleable_types is not None and column.object_type not in sampleable_types:
+            continue
+        result = calculate_sampling_params(
+            table_name=column.table_name,
+            record_count=column.record_ct,
+            sample_percent_raw=table_group.profile_sample_percent,
+            min_sample=table_group.profile_sample_min_count,
+        )
+        if result:
+            sampling_params[column.table_name] = result
+    return sampling_params
+
+
+def _run_column_profiling(
+    sql_generator: ProfilingSQL, data_chars: list[ColumnChars], sampling_params: dict[str, TableSampling]
+) -> None:
     profiling_run = sql_generator.profiling_run
     profiling_run.set_progress("col_profiling", "Running")
     profiling_run.save()
     get_current_session().commit()
 
     LOG.info(f"Running column profiling queries: {len(data_chars)}")
-    table_group = sql_generator.table_group
-    sampling_params: dict[str, TableSampling] = {}
-    if table_group.profile_use_sampling:
-        for column in data_chars:
-            if not sampling_params.get(column.table_name):
-                result = calculate_sampling_params(
-                    table_name=column.table_name,
-                    record_count=column.record_ct,
-                    sample_percent_raw=table_group.profile_sample_percent,
-                    min_sample=table_group.profile_sample_min_count,
-                )
-                if result:
-                    sampling_params[column.table_name] = result
 
     def update_column_progress(progress: ThreadedProgress) -> None:
         profiling_run.set_progress(
@@ -218,7 +233,7 @@ def update_column_progress(progress: ThreadedProgress) -> None:
     )
 
 
-def _run_frequency_analysis(sql_generator: ProfilingSQL) -> None:
+def _run_frequency_analysis(sql_generator: ProfilingSQL, sampling_params: dict[str, TableSampling]) -> None:
     profiling_run = sql_generator.profiling_run
     profiling_run.set_progress("freq_analysis", "Running")
     profiling_run.save()
@@ -227,7 +242,7 @@ def _run_frequency_analysis(sql_generator: ProfilingSQL) -> None:
     error_data = None
     try:
         LOG.info("Selecting columns for frequency analysis")
-        frequency_columns = fetch_dict_from_db(*sql_generator.get_frequency_analysis_columns())
+        frequency_columns = [ColumnChars(**column) for column in fetch_dict_from_db(*sql_generator.get_frequency_analysis_columns())]
 
         if frequency_columns:
             LOG.info(f"Running frequency analysis queries: {len(frequency_columns)}")
@@ -240,7 +255,10 @@ def update_frequency_progress(progress: ThreadedProgress) -> None:
                 get_current_session().commit()
 
             frequency_results, result_columns, error_data = fetch_from_db_threaded(
-                [sql_generator.run_frequency_analysis(ColumnChars(**column)) for column in frequency_columns],
+                [
+                    sql_generator.run_frequency_analysis(column, sampling_params.get(column.table_name))
+                    for column in frequency_columns
+                ],
                 use_target_db=True,
                 max_threads=sql_generator.connection.max_threads,
                 progress_callback=update_frequency_progress,

testgen/common/clean_sql.py

@@ -50,6 +50,19 @@ def quote_identifiers(identifiers: str, flavor: str) -> str:
     return ", ".join(quoted_values)
 
 
+def null_if_empty(value: object) -> object:
+    """Return the literal ``"NULL"`` when ``value`` is empty (``None`` or ``""``), else ``value``.
+
+    For numeric test parameters substituted into SQL templates (baseline counts,
+    averages, standard deviations, tolerances). An empty substitution produces invalid SQL
+    such as ``CAST( AS FLOAT)``; ``"NULL"`` makes the surrounding expression evaluate to NULL
+    instead. A real ``0`` is preserved (``0 in (None, "")`` is ``False``). Not suitable for
+    params used as quoted literals or IN-lists (e.g. ``BASELINE_VALUE``, ``BASELINE_SUM`` for
+    Freshness_Trend), where ``"NULL"`` would not be valid SQL.
+    """
+    return "NULL" if value in (None, "") else value
+
+
 def concat_columns(columns: str, null_value: str):
     # Prepares SQL expression to concatenate comma-separated column list
     expression = ""

testgen/common/database/column_chars.py

@@ -1,4 +1,13 @@
 import dataclasses
+from enum import StrEnum
+
+
+class ObjectType(StrEnum):
+    TABLE = "TABLE"
+    VIEW = "VIEW"
+    MATERIALIZED_VIEW = "MATERIALIZED_VIEW"
+    EXTERNAL = "EXTERNAL"
+    OTHER = "OTHER"
 
 
 @dataclasses.dataclass
@@ -11,6 +20,7 @@ class ColumnChars:
     column_type: str | None = None
     db_data_type: str | None = None
     is_decimal: bool = False
+    object_type: ObjectType | None = None
     approx_record_ct: int | None = None
     # This should not default to 0 since we don't always retrieve actual row counts
     # UI relies on the null value to know that the approx_record_ct should be displayed instead

testgen/common/database/flavor/flavor_service.py

@@ -6,7 +6,7 @@
 from sqlalchemy import create_engine as sqlalchemy_create_engine
 from sqlalchemy.engine.base import Engine
 
-from testgen.common.database.column_chars import ColumnChars
+from testgen.common.database.column_chars import ColumnChars, ObjectType
 from testgen.common.encrypt import DecryptText
 
 SQLFlavor = Literal["redshift", "redshift_spectrum", "snowflake", "mssql", "postgresql", "databricks", "bigquery", "oracle", "sap_hana", "salesforce_data360"]
@@ -116,6 +116,7 @@ def row_limit_clauses(self, n: int) -> tuple[str, str]:
 
     qualifies_table_refs_with_schema = True
     metadata_via_api = False
+    sampleable_object_types: frozenset[ObjectType] | None = None
 
     def get_schema_columns(self, _params: ResolvedConnectionParams, _schema: str) -> list[ColumnChars] | None:
         """Return column metadata without querying information_schema.

testgen/common/database/flavor/mssql_flavor_service.py

@@ -1,6 +1,7 @@
 from sqlalchemy.engine import URL
 
 from testgen import settings
+from testgen.common.database.column_chars import ObjectType
 from testgen.common.database.flavor.flavor_service import FlavorService, ResolvedConnectionParams
 
 
@@ -10,6 +11,8 @@ class MssqlFlavorService(FlavorService):
     escaped_underscore = "[_]"
     row_limiting_clause = "top"
     url_scheme = "mssql+pyodbc"
+    # TABLESAMPLE is rejected on views; SQL Server has no materialized views, so only base tables.
+    sampleable_object_types = frozenset({ObjectType.TABLE})
 
     def get_connection_string_from_fields(self, params: ResolvedConnectionParams) -> str:
         connection_url = URL.create(

testgen/common/database/flavor/postgresql_flavor_service.py

@@ -1,5 +1,6 @@
 from urllib.parse import quote_plus
 
+from testgen.common.database.column_chars import ObjectType
 from testgen.common.database.flavor.flavor_service import ResolvedConnectionParams
 from testgen.common.database.flavor.redshift_flavor_service import RedshiftFlavorService
 
@@ -8,6 +9,8 @@ class PostgresqlFlavorService(RedshiftFlavorService):
 
     escaped_underscore = "\\_"
     url_scheme = "postgresql"
+    # TABLESAMPLE applies only to tables and materialized views
+    sampleable_object_types = frozenset({ObjectType.TABLE, ObjectType.MATERIALIZED_VIEW})
 
     def get_connection_string_from_fields(self, params: ResolvedConnectionParams) -> str:
         if params.host.startswith("/"):

testgen/common/models/data_table.py

@@ -17,6 +17,7 @@
 )
 from sqlalchemy.dialects import postgresql
 
+from testgen.common.database.column_chars import ObjectType
 from testgen.common.models import get_current_session
 from testgen.common.models.data_column import DataColumnChars
 from testgen.common.models.entity import Entity
@@ -42,6 +43,7 @@ class TableProfilingOverview:
     table_groups_id: UUID
     schema_name: str | None
     table_name: str
+    object_type: ObjectType | None
     record_ct: int | None
     column_ct: int | None
     dq_score_profiling: float | None
@@ -61,6 +63,7 @@ class DataTable(Entity):
     table_groups_id: UUID = Column(postgresql.UUID(as_uuid=True), ForeignKey("table_groups.id"))
     schema_name: str | None = Column(String)
     table_name: str = Column(String)
+    object_type: ObjectType | None = Column(String)
     column_ct: int | None = Column(BigInteger)
     record_ct: int | None = Column(BigInteger)
     approx_record_ct: int | None = Column(BigInteger)
@@ -110,6 +113,7 @@ def get_profiling_overview(
                 cls.table_groups_id,
                 cls.schema_name,
                 cls.table_name,
+                cls.object_type,
                 cls.record_ct,
                 cls.column_ct,
                 cls.dq_score_profiling,

testgen/common/source_data_service.py

@@ -11,7 +11,7 @@
 import pandas as pd
 from sqlalchemy import text
 
-from testgen.common.clean_sql import concat_columns
+from testgen.common.clean_sql import concat_columns, null_if_empty
 from testgen.common.database.database_service import get_flavor_service, replace_params
 from testgen.common.date_service import parse_fuzzy_date
 from testgen.common.models import get_current_session
@@ -207,12 +207,14 @@ def _build_query_standard(issue_data: dict, limit: int) -> str | None:
         if (parsed_test_date := parse_fuzzy_date(issue_data["test_date"]))
         else None,
         "CUSTOM_QUERY": test_definition.custom_query,
+        # BASELINE_VALUE varies (quoted literal / number / IN-list) by test type — templated as-is.
+        # The numeric baselines render empty as NULL to avoid invalid SQL like CAST( AS FLOAT).
         "BASELINE_VALUE": test_definition.baseline_value,
-        "BASELINE_CT": test_definition.baseline_ct,
-        "BASELINE_AVG": test_definition.baseline_avg,
-        "BASELINE_SD": test_definition.baseline_sd,
-        "LOWER_TOLERANCE": "NULL" if test_definition.lower_tolerance in (None, "") else test_definition.lower_tolerance,
-        "UPPER_TOLERANCE": "NULL" if test_definition.upper_tolerance in (None, "") else test_definition.upper_tolerance,
+        "BASELINE_CT": null_if_empty(test_definition.baseline_ct),
+        "BASELINE_AVG": null_if_empty(test_definition.baseline_avg),
+        "BASELINE_SD": null_if_empty(test_definition.baseline_sd),
+        "LOWER_TOLERANCE": null_if_empty(test_definition.lower_tolerance),
+        "UPPER_TOLERANCE": null_if_empty(test_definition.upper_tolerance),
         "THRESHOLD_VALUE": test_definition.threshold_value or 0,
         # SUBSET_CONDITION should be replaced after CUSTOM_QUERY
         # since the latter may contain the former