|
| 1 | +"""Shared data catalog convenience service. |
| 2 | +
|
| 3 | +Generates CREATE TABLE scripts from profiled column metadata and fetches |
| 4 | +sample rows from source tables. Used by both the Streamlit UI and MCP tools. |
| 5 | +""" |
| 6 | +import logging |
| 7 | +from dataclasses import dataclass |
| 8 | +from typing import Literal |
| 9 | +from uuid import UUID |
| 10 | + |
| 11 | +import pandas as pd |
| 12 | + |
| 13 | +from testgen.common.database.database_service import get_flavor_service |
| 14 | +from testgen.common.database.flavor.flavor_service import FlavorService |
| 15 | +from testgen.common.models.connection import Connection |
| 16 | +from testgen.common.models.data_column import CreateScriptColumn, DataColumnChars |
| 17 | +from testgen.common.pii_masking import get_pii_columns, mask_source_data_pii |
| 18 | +from testgen.ui.services.database_service import fetch_from_target_db |
| 19 | +from testgen.utils import to_dataframe |
| 20 | + |
| 21 | +LOG = logging.getLogger("testgen") |
| 22 | + |
| 23 | + |
| 24 | +@dataclass |
| 25 | +class TableSampleResult: |
| 26 | + status: Literal["OK", "ND", "ERR"] |
| 27 | + message: str | None = None |
| 28 | + df: pd.DataFrame | None = None |
| 29 | + pii_redacted: bool = False |
| 30 | + |
| 31 | + |
| 32 | +def render_create_table_script( |
| 33 | + schema_name: str, |
| 34 | + table_name: str, |
| 35 | + columns: list[CreateScriptColumn], |
| 36 | + flavor_service: FlavorService, |
| 37 | + *, |
| 38 | + annotate_changes: bool = False, |
| 39 | +) -> str: |
| 40 | + """Render CREATE TABLE DDL from profiled columns, quoting identifiers for the flavor. |
| 41 | +
|
| 42 | + Column types use the profiling-derived suggestion, falling back to the original |
| 43 | + database type. ``annotate_changes`` appends ``-- WAS <type>`` comments where the |
| 44 | + suggestion differs from the original type. |
| 45 | + """ |
| 46 | + quote = flavor_service.quote_character |
| 47 | + table_ref = flavor_service.get_table_ref(schema_name, table_name) |
| 48 | + quoted_names = [f"{quote}{col.column_name}{quote}" for col in columns] |
| 49 | + |
| 50 | + name_width = max(len(name) for name in quoted_names) |
| 51 | + type_width = max(len(col.datatype_suggestion or col.db_data_type or "") for col in columns) |
| 52 | + |
| 53 | + col_defs = [] |
| 54 | + for index, (col, name) in enumerate(zip(columns, quoted_names, strict=True)): |
| 55 | + col_type = col.datatype_suggestion or col.db_data_type or "" |
| 56 | + separator = "" if index == len(columns) - 1 else "," |
| 57 | + line = f"{name:<{name_width}} {col_type:<{type_width}}{separator}" |
| 58 | + if ( |
| 59 | + annotate_changes |
| 60 | + and col.db_data_type |
| 61 | + and col.datatype_suggestion |
| 62 | + and col.db_data_type.lower() != col.datatype_suggestion.lower() |
| 63 | + ): |
| 64 | + line = f"{line} -- WAS {col.db_data_type}" |
| 65 | + col_defs.append(line.rstrip()) |
| 66 | + |
| 67 | + body = "\n ".join(col_defs) |
| 68 | + return f"CREATE TABLE {table_ref} (\n {body}\n);" |
| 69 | + |
| 70 | + |
| 71 | +def build_create_table_script( |
| 72 | + table_group_id: UUID, table_name: str, *, annotate_changes: bool = False, |
| 73 | +) -> str | None: |
| 74 | + """Build a CREATE TABLE script for a profiled table, or ``None`` if it is not in the catalog.""" |
| 75 | + schema_name, columns = DataColumnChars.list_for_create_script(table_group_id, table_name) |
| 76 | + if not columns or schema_name is None: |
| 77 | + return None |
| 78 | + connection = Connection.get_by_table_group(table_group_id) |
| 79 | + if connection is None: |
| 80 | + return None |
| 81 | + flavor_service = get_flavor_service(connection.sql_flavor) |
| 82 | + return render_create_table_script( |
| 83 | + schema_name, table_name, columns, flavor_service, annotate_changes=annotate_changes, |
| 84 | + ) |
| 85 | + |
| 86 | + |
| 87 | +def fetch_table_sample( |
| 88 | + connection: Connection, |
| 89 | + table_group_id: UUID, |
| 90 | + schema_name: str, |
| 91 | + table_name: str, |
| 92 | + *, |
| 93 | + limit: int, |
| 94 | + mask_pii: bool, |
| 95 | + column_name: str | None = None, |
| 96 | +) -> TableSampleResult: |
| 97 | + """Fetch distinct sample rows from a source table, masking PII columns when requested.""" |
| 98 | + flavor_service = get_flavor_service(connection.sql_flavor) |
| 99 | + prefix, suffix = flavor_service.row_limit_clauses(limit) |
| 100 | + quote = flavor_service.quote_character |
| 101 | + table_ref = flavor_service.get_table_ref(schema_name, table_name) |
| 102 | + columns_expr = f"{quote}{column_name}{quote}" if column_name else "*" |
| 103 | + query = f"SELECT DISTINCT {prefix} {columns_expr} FROM {table_ref} {suffix}".strip() |
| 104 | + |
| 105 | + try: |
| 106 | + results = fetch_from_target_db(connection, query) |
| 107 | + except Exception: |
| 108 | + LOG.exception("Table sample fetch encountered an error.") |
| 109 | + return TableSampleResult("ERR", message="The sample data could not be loaded.") |
| 110 | + |
| 111 | + if not results: |
| 112 | + return TableSampleResult("ND") |
| 113 | + |
| 114 | + df = to_dataframe(results) |
| 115 | + pii_redacted = False |
| 116 | + if mask_pii: |
| 117 | + pii_columns = get_pii_columns(str(table_group_id), schema_name, table_name) |
| 118 | + pii_redacted = mask_source_data_pii(df, pii_columns) |
| 119 | + return TableSampleResult("OK", df=df, pii_redacted=pii_redacted) |
0 commit comments