fix(profiling): skip sampling on views where the sample clause rejects it

SQL Server and PostgreSQL reject TABLESAMPLE on views ("can only be used with
local tables" / "tables and materialized views"), so profiling a table group
containing a view with sampling enabled errored those columns. Add a
samples_views flavor capability (False for mssql/postgresql) and an object_type
signal from those flavors' DDF, and skip sampling for views on those flavors —
they are profiled in full.

Also compute the per-run sampling params once and share them between column
profiling and frequency analysis. Frequency analysis previously never sampled
(the secondary query's TABLESAMPLE branch was unreachable); it now samples the
same tables, with the view-skip applied. Sample-scale frequency counts are left
unscaled.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: aarthy-dk

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/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: str | None = None # ObjectType values
     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 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:
         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

@@ -42,6 +42,7 @@ class TableProfilingOverview:
     table_groups_id: UUID
     schema_name: str | None
     table_name: str
+    object_type: str | None
     record_ct: int | None
     column_ct: int | None
     dq_score_profiling: float | None
@@ -61,6 +62,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: str | None = Column(String)
     column_ct: int | None = Column(BigInteger)
     record_ct: int | None = Column(BigInteger)
     approx_record_ct: int | None = Column(BigInteger)
@@ -110,6 +112,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/template/data_chars/data_chars_update.sql

@@ -8,6 +8,7 @@ WITH new_chars AS (
       schema_name,
       table_name,
       run_date,
+      MAX(object_type) AS object_type,
       MAX(approx_record_ct) AS approx_record_ct,
       MAX(record_ct) AS record_ct,
       COUNT(*) AS column_ct
@@ -20,7 +21,8 @@ WITH new_chars AS (
 ),
 updated_records AS (
    UPDATE data_table_chars
-   SET approx_record_ct = n.approx_record_ct,
+   SET object_type = n.object_type,
+      approx_record_ct = n.approx_record_ct,
       record_ct = n.record_ct,
       column_ct = n.column_ct,
       last_refresh_date = n.run_date,
@@ -55,6 +57,7 @@ WITH new_chars AS (
       schema_name,
       table_name,
       run_date,
+      MAX(object_type) AS object_type,
       MAX(approx_record_ct) AS approx_record_ct,
       MAX(record_ct) AS record_ct,
       COUNT(*) AS column_ct
@@ -72,6 +75,7 @@ inserted_records AS (
          table_name,
          add_date,
          last_refresh_date,
+         object_type,
          approx_record_ct,
          record_ct,
          column_ct
@@ -81,6 +85,7 @@ inserted_records AS (
       n.table_name,
       n.run_date,
       n.run_date,
+      n.object_type,
       n.approx_record_ct,
       n.record_ct,
       n.column_ct

testgen/template/dbsetup/030_initialize_new_schema_structure.sql

@@ -35,6 +35,7 @@ CREATE TABLE stg_data_chars_updates (
    general_type          VARCHAR(1),
    column_type           VARCHAR(50),
    db_data_type          VARCHAR(50),
+   object_type           VARCHAR(20),
    approx_record_ct      BIGINT,
    record_ct             BIGINT
 );
@@ -419,6 +420,7 @@ CREATE TABLE data_table_chars (
    table_groups_id       UUID,
    schema_name           VARCHAR(50),
    table_name            VARCHAR(120),
+   object_type           VARCHAR(20),
    functional_table_type VARCHAR(50),
    description           VARCHAR(1000),
    critical_data_element BOOLEAN,

testgen/template/dbupgrade/0193_incremental_upgrade.sql

@@ -0,0 +1,7 @@
+SET SEARCH_PATH TO {SCHEMA_NAME};
+
+ALTER TABLE data_table_chars
+    ADD COLUMN IF NOT EXISTS object_type VARCHAR(20);
+
+ALTER TABLE stg_data_chars_updates
+    ADD COLUMN IF NOT EXISTS object_type VARCHAR(20);