From cc49b65901c467fbad0c7deac52d31992601598e Mon Sep 17 00:00:00 2001 From: Peter <101368063+ClassicMMT@users.noreply.github.com> Date: Thu, 16 Jul 2026 12:54:15 +1200 Subject: [PATCH 1/3] feat: safe data preview --- HISTORY.rst | 10 + datamasque/client/__init__.py | 66 ++++ datamasque/client/models/discovery.py | 25 +- datamasque/client/models/safe_data_preview.py | 338 ++++++++++++++++++ tests/test_discovery.py | 37 ++ tests/test_safe_data_preview.py | 262 ++++++++++++++ 6 files changed, 737 insertions(+), 1 deletion(-) create mode 100644 datamasque/client/models/safe_data_preview.py create mode 100644 tests/test_safe_data_preview.py diff --git a/HISTORY.rst b/HISTORY.rst index cf9720f..534eb63 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -2,6 +2,16 @@ History ======= +1.1.8 (unreleased) +------------------ + +* Added typed Safe Data Preview support: + + * ``safe_data_preview`` on ``InDataDiscoveryConfig`` (via ``SafeDataPreviewOptions``) + to configure or disable the preview. + * ``safe_data_preview`` on schema-discovery result columns and file-discovery locators, + typed by ``kind``. + 1.1.7 (2026-07-14) ------------------ diff --git a/datamasque/client/__init__.py b/datamasque/client/__init__.py index e2275e3..0229cd3 100644 --- a/datamasque/client/__init__.py +++ b/datamasque/client/__init__.py @@ -116,6 +116,40 @@ RunInfo, UnfinishedRun, ) +from datamasque.client.models.safe_data_preview import ( + BooleanPreview, + BooleanStatistics, + ColumnKind, + ColumnPreview, + CommonStatistics, + FirstCharsStatistics, + LengthEntry, + LengthsStatistics, + MaskedFormEntry, + NumericBin, + NumericHistograms, + NumericPreview, + NumericStatistics, + NumericSummaries, + NumericTemporalDisclosureLevel, + PatternComposition, + PatternEntry, + PatternsStatistics, + SafeDataPreview, + SafeDataPreviewOptions, + StringDisclosureLevel, + StringPreview, + StringStatistics, + TemporalBin, + TemporalHistograms, + TemporalPreview, + TemporalStatistics, + TemporalSummaries, + UnknownPreview, + UnsupportedPreview, + UnsupportedPreviewReason, + UnsupportedStatistics, +) from datamasque.client.models.status import ( AsyncRulesetGenerationTaskStatus, MaskingRunStatus, @@ -131,6 +165,11 @@ "AsyncRulesetGenerationInProgressError", "AsyncRulesetGenerationTaskStatus", "AzureConnectionConfig", + "BooleanPreview", + "BooleanStatistics", + "ColumnKind", + "ColumnPreview", + "CommonStatistics", "ConnectionConfig", "ConnectionId", "ConstraintColumns", @@ -171,6 +210,7 @@ "FileId", "FileOrContent", "FileRulesetGenerationRequest", + "FirstCharsStatistics", "ForeignKeyRef", "GitSnapshot", "HashColumnsTableConfig", @@ -186,16 +226,28 @@ "InvalidLibraryError", "InvalidRulesetError", "JsonPath", + "LengthEntry", + "LengthsStatistics", "LicenseInfo", "Locator", "MaskType", + "MaskedFormEntry", "MaskingRunOptions", "MaskingRunRequest", "MaskingRunStatus", "MongoConnectionConfig", "MountedShareConnectionConfig", "MssqlLinkedServerConnectionConfig", + "NumericBin", + "NumericHistograms", + "NumericPreview", + "NumericStatistics", + "NumericSummaries", + "NumericTemporalDisclosureLevel", "OracleWalletFile", + "PatternComposition", + "PatternEntry", + "PatternsStatistics", "ReferencingForeignKey", "Ruleset", "RulesetGenerationRequest", @@ -213,6 +265,8 @@ "RunInfo", "RunNotCancellableError", "S3ConnectionConfig", + "SafeDataPreview", + "SafeDataPreviewOptions", "SchemaDiscoveryColumn", "SchemaDiscoveryFromConfigRequest", "SchemaDiscoveryPage", @@ -228,9 +282,21 @@ "SseConfig", "SseSelection", "SslZipFile", + "StringDisclosureLevel", + "StringPreview", + "StringStatistics", "SwitchableLicenseMetadata", "TableConstraints", + "TemporalBin", + "TemporalHistograms", + "TemporalPreview", + "TemporalStatistics", + "TemporalSummaries", "UnfinishedRun", + "UnknownPreview", + "UnsupportedPreview", + "UnsupportedPreviewReason", + "UnsupportedStatistics", "User", "UserId", "UserRole", diff --git a/datamasque/client/models/discovery.py b/datamasque/client/models/discovery.py index 23a5a88..3b69fe7 100644 --- a/datamasque/client/models/discovery.py +++ b/datamasque/client/models/discovery.py @@ -3,13 +3,19 @@ from enum import Enum from typing import Any, Optional, Union -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator from datamasque.client.models.connection import ConnectionConfig, ConnectionId, unwrap_connection_id from datamasque.client.models.data_selection import HashColumnsTableConfig, Locator, UserSelection from datamasque.client.models.discovery_config import DiscoveryConfig, DiscoveryConfigId, unwrap_discovery_config_id from datamasque.client.models.pagination import Page from datamasque.client.models.runs import RunConnectionRef +from datamasque.client.models.safe_data_preview import ( + ColumnPreview, + SafeDataPreview, + SafeDataPreviewOptions, + parse_safe_data_preview, +) class InDataDiscoveryRule(BaseModel): @@ -32,6 +38,7 @@ class InDataDiscoveryConfig(BaseModel): non_sensitive_rules: Optional[list[InDataDiscoveryRule]] = None ignore_rules: Optional[list[InDataDiscoveryRule]] = None force: Optional[bool] = None + safe_data_preview: Optional[SafeDataPreviewOptions] = None class SchemaDiscoveryRequest(BaseModel): @@ -294,6 +301,14 @@ class SchemaDiscoveryColumn(BaseModel): unique_index_names: list[str] referencing_foreign_keys: list[ReferencingForeignKey] constraint: str # Primary or Unique, or empty string if column does not participate in a PK/UK + safe_data_preview: Optional[SafeDataPreview] = None + + @field_validator("safe_data_preview", mode="before") + @classmethod + def _parse_safe_data_preview( + cls, value: Union[dict[str, JsonValue], ColumnPreview, None] + ) -> Optional[ColumnPreview]: + return parse_safe_data_preview(value) class SchemaDiscoveryResult(BaseModel): @@ -356,6 +371,14 @@ class FileDiscoveryLocatorResult(BaseModel): locator: Locator matches: list[FileDiscoveryMatch] data_types: list[str] + safe_data_preview: Optional[SafeDataPreview] = None + + @field_validator("safe_data_preview", mode="before") + @classmethod + def _parse_safe_data_preview( + cls, value: Union[dict[str, JsonValue], ColumnPreview, None] + ) -> Optional[ColumnPreview]: + return parse_safe_data_preview(value) class FileDiscoveryFile(BaseModel): diff --git a/datamasque/client/models/safe_data_preview.py b/datamasque/client/models/safe_data_preview.py new file mode 100644 index 0000000..a772eaa --- /dev/null +++ b/datamasque/client/models/safe_data_preview.py @@ -0,0 +1,338 @@ +"""Typed request and response shapes for Safe Data Preview, part of In-Data Discovery.""" + +from enum import Enum +from typing import Literal, Optional, Union + +from pydantic import ( + BaseModel, + ConfigDict, + JsonValue, + SerializeAsAny, + ValidationError, +) + + +class StringDisclosureLevel(Enum): + """How much detail a string Safe Data Preview reveals.""" + + lengths = "lengths" + patterns = "patterns" + first_chars = "first_chars" + + +class NumericTemporalDisclosureLevel(Enum): + """How much detail a numeric or temporal Safe Data Preview reveals.""" + + summaries = "summaries" + histograms = "histograms" + + +class SafeDataPreviewOptions(BaseModel): + """Configures whether Safe Data Preview runs and how much detail each column reveals.""" + + model_config = ConfigDict(extra="forbid") + + enabled: Optional[bool] = None + min_group_size: Optional[int] = None + string_level: Optional[StringDisclosureLevel] = None + numeric_level: Optional[NumericTemporalDisclosureLevel] = None + temporal_level: Optional[NumericTemporalDisclosureLevel] = None + + +class ColumnKind(str, Enum): + """Which shape a column's Safe Data Preview takes.""" + + boolean = "boolean" + numeric = "numeric" + string = "string" + temporal = "temporal" + unsupported = "unsupported" + + +class UnsupportedPreviewReason(str, Enum): + """Known reasons an `UnsupportedPreview` gives for declining a column.""" + + no_rows = "No rows sampled" + all_null = "All sampled values are null" + all_identical = "All sampled values are identical" + below_min_group_size = "A value count is below the minimum group size" + uninterpretable = "Values couldn't be interpreted as a single consistent type" + is_binary = "Binary data isn't previewed" + datatype_not_previewed = "This data type has no Safe Data Preview" + too_many_columns = "Table has too many columns for Safe Data Preview" + unavailable = "This column could not be previewed" + + +class CommonStatistics(BaseModel): + """Row counts shared by every preview kind.""" + + model_config = ConfigDict(extra="allow") + + count_row: int + count_null: int + count_distinct: Optional[int] = None + + +class LengthEntry(BaseModel): + """A value length and how many sampled rows had it.""" + + model_config = ConfigDict(extra="allow") + + length: int + count: int + + +class LengthsStatistics(BaseModel): + """Distribution of value lengths.""" + + model_config = ConfigDict(extra="allow") + + min: int + max: int + mean: float + median: float + most_common: list[LengthEntry] + + +class PatternEntry(BaseModel): + """A masked character pattern and how many sampled rows matched it.""" + + model_config = ConfigDict(extra="allow") + + pattern: str + count: int + + +class PatternComposition(BaseModel): + """The proportion of letter, digit, and other characters across sampled values.""" + + model_config = ConfigDict(extra="allow") + + letter: float + digit: float + other: float + + +class PatternsStatistics(BaseModel): + """The most common character patterns and the overall character composition.""" + + model_config = ConfigDict(extra="allow") + + top: list[PatternEntry] + composition: PatternComposition + + +class MaskedFormEntry(BaseModel): + """A masked leading form of a value and how many sampled rows had it.""" + + model_config = ConfigDict(extra="allow") + + masked: str + count: int + + +class FirstCharsStatistics(BaseModel): + """The most common masked leading forms.""" + + model_config = ConfigDict(extra="allow") + + top: list[MaskedFormEntry] + + +class StringStatistics(BaseModel): + """String preview statistics; `patterns`/`first_chars` appear only at the matching disclosure level.""" + + model_config = ConfigDict(extra="allow") + + lengths: LengthsStatistics + patterns: Optional[PatternsStatistics] = None + first_chars: Optional[FirstCharsStatistics] = None + + +class NumericSummaries(BaseModel): + """Numeric mean and percentile summary.""" + + model_config = ConfigDict(extra="allow") + + mean: float + q1: float + q2: float + q3: float + p5: float + p95: float + + +class NumericBin(BaseModel): + """One half-open interval `[lower_bound, upper_bound)` of a numeric histogram and its row count.""" + + model_config = ConfigDict(extra="allow") + + lower_bound: float + upper_bound: float + count: int + + +class NumericHistograms(BaseModel): + """The bins of a numeric histogram.""" + + model_config = ConfigDict(extra="allow") + + bins: list[NumericBin] + + +class NumericStatistics(BaseModel): + """Numeric preview statistics; `histograms` is present only at the histogram disclosure level.""" + + model_config = ConfigDict(extra="allow") + + summaries: NumericSummaries + histograms: Optional[NumericHistograms] = None + + +class TemporalSummaries(BaseModel): + """Temporal mean and percentile summary, each an ISO-formatted string.""" + + model_config = ConfigDict(extra="allow") + + mean: str + q1: str + q2: str + q3: str + p5: str + p95: str + + +class TemporalBin(BaseModel): + """One interval of a temporal histogram, with ISO-formatted bounds and a human-readable label.""" + + model_config = ConfigDict(extra="allow") + + lower_bound: str + upper_bound: str + label: str + count: int + + +class TemporalHistograms(BaseModel): + """The bins of a temporal histogram.""" + + model_config = ConfigDict(extra="allow") + + bins: list[TemporalBin] + + +class TemporalStatistics(BaseModel): + """Temporal preview statistics; `histograms` is present only at the histogram disclosure level.""" + + model_config = ConfigDict(extra="allow") + + summaries: TemporalSummaries + histograms: Optional[TemporalHistograms] = None + + +class BooleanStatistics(BaseModel): + """Counts of true and false values.""" + + model_config = ConfigDict(extra="allow") + + count_true: int + count_false: int + + +class UnsupportedStatistics(BaseModel): + """Why a column could not be previewed. See `UnsupportedPreviewReason` for known values.""" + + model_config = ConfigDict(extra="allow") + + reason: str + + +class ColumnPreview(BaseModel): + """Fields present on every Safe Data Preview.""" + + model_config = ConfigDict(extra="allow") + + # For grouped file discovery, the file whose sample produced this preview; None for database columns. + sampled_from: Optional[str] = None + + +class StringPreview(ColumnPreview): + """A preview of a string column.""" + + kind: Literal[ColumnKind.string] = ColumnKind.string + statistics_common: CommonStatistics + statistics_kind: StringStatistics + + +class NumericPreview(ColumnPreview): + """A preview of a numeric column.""" + + kind: Literal[ColumnKind.numeric] = ColumnKind.numeric + statistics_common: CommonStatistics + statistics_kind: NumericStatistics + + +class TemporalPreview(ColumnPreview): + """A preview of a date/time column.""" + + kind: Literal[ColumnKind.temporal] = ColumnKind.temporal + statistics_common: CommonStatistics + statistics_kind: TemporalStatistics + + +class BooleanPreview(ColumnPreview): + """A preview of a boolean column.""" + + kind: Literal[ColumnKind.boolean] = ColumnKind.boolean + statistics_common: CommonStatistics + statistics_kind: BooleanStatistics + + +class UnsupportedPreview(ColumnPreview): + """A column the preview ran against but declined; `statistics_kind.reason` says why.""" + + kind: Literal[ColumnKind.unsupported] = ColumnKind.unsupported + statistics_common: CommonStatistics + statistics_kind: UnsupportedStatistics + + +class UnknownPreview(ColumnPreview): + """A Safe Data Preview whose `kind` or shape this client can't parse; raw fields stay on `model_extra`.""" + + kind: Optional[str] = None + statistics_common: Optional[CommonStatistics] = None + + +PREVIEW_BY_KIND: dict[str, type[ColumnPreview]] = { + ColumnKind.string.value: StringPreview, + ColumnKind.numeric.value: NumericPreview, + ColumnKind.temporal.value: TemporalPreview, + ColumnKind.boolean.value: BooleanPreview, + ColumnKind.unsupported.value: UnsupportedPreview, +} + + +def parse_safe_data_preview(value: Union[dict[str, JsonValue], ColumnPreview, None]) -> Optional[ColumnPreview]: + """Parse a raw preview into its typed variant, or `UnknownPreview` if the kind or shape is unrecognised.""" + + if value is None: + return None + if isinstance(value, ColumnPreview): + return value + if not isinstance(value, dict): + return None + kind = value.get("kind") + preview_type = PREVIEW_BY_KIND.get(kind) if isinstance(kind, str) else None + if preview_type is not None: + try: + return preview_type.model_validate(value) + except ValidationError: + pass + try: + return UnknownPreview.model_validate(value) + except ValidationError: + # Fallback to raw + return UnknownPreview.model_construct(**value) # type: ignore[arg-type] + + +SafeDataPreview = SerializeAsAny[ColumnPreview] diff --git a/tests/test_discovery.py b/tests/test_discovery.py index 6121430..a9706d7 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -22,10 +22,13 @@ InDataDiscoveryConfig, RulesetGenerationRequest, RunId, + SafeDataPreviewOptions, SchemaDiscoveryFromConfigRequest, SchemaDiscoveryPage, SchemaDiscoveryRequest, SchemaDiscoveryResult, + StringDisclosureLevel, + StringPreview, ) from datamasque.client.exceptions import ( AsyncRulesetGenerationInProgressError, @@ -648,6 +651,18 @@ def test_schema_discovery_request_model_dump_includes_set_fields(): } +def test_schema_discovery_request_model_dump_includes_safe_data_preview(): + req = SchemaDiscoveryRequest( + connection="conn-1", + in_data_discovery=InDataDiscoveryConfig( + enabled=True, + safe_data_preview=SafeDataPreviewOptions(enabled=False, string_level=StringDisclosureLevel.lengths), + ), + ) + dumped = req.model_dump(exclude_none=True, mode="json") + assert dumped["in_data_discovery"]["safe_data_preview"] == {"enabled": False, "string_level": "lengths"} + + def test_discovery_requests_accept_connection_config_objects(): """All three discovery request models accept a full `ConnectionConfig` and extract its `id`.""" connection = DatabaseConnectionConfig( @@ -770,6 +785,28 @@ def _schema_discovery_row(row_id: int, column_name: str, table_name: str = "user } +def test_schema_discovery_result_parses_safe_data_preview(): + row = _schema_discovery_row(1, "email") + row["data"]["safe_data_preview"] = { + "kind": "string", + "sampled_from": "data/people.csv", + "statistics_common": {"count_row": 1000, "count_null": 0, "count_distinct": 988}, + "statistics_kind": { + "lengths": {"min": 5, "max": 30, "mean": 18.0, "median": 18.0, "most_common": []}, + }, + } + result = SchemaDiscoveryResult.model_validate(row) + preview = result.data.safe_data_preview + assert isinstance(preview, StringPreview) + assert preview.sampled_from == "data/people.csv" + assert preview.statistics_common.count_distinct == 988 + + +def test_schema_discovery_result_without_safe_data_preview(): + result = SchemaDiscoveryResult.model_validate(_schema_discovery_row(1, "email")) + assert result.data.safe_data_preview is None + + def test_list_schema_discovery_results_follows_pagination(client): run_id = RunId(42) page1 = { diff --git a/tests/test_safe_data_preview.py b/tests/test_safe_data_preview.py new file mode 100644 index 0000000..6d08109 --- /dev/null +++ b/tests/test_safe_data_preview.py @@ -0,0 +1,262 @@ +"""Tests for the typed Safe Data Preview models and their tolerant parsing.""" + +from pydantic import BaseModel + +from datamasque.client import FileDiscoveryResult, SchemaDiscoveryColumn +from datamasque.client.models.safe_data_preview import ( + BooleanPreview, + NumericPreview, + StringPreview, + TemporalPreview, + UnknownPreview, + UnsupportedPreview, + UnsupportedPreviewReason, + parse_safe_data_preview, +) + +_STRING_PREVIEW = { + "kind": "string", + "sampled_from": None, + "statistics_common": {"count_row": 1000, "count_null": 12, "count_distinct": 988}, + "statistics_kind": { + "lengths": { + "min": 9, + "max": 34, + "mean": 21.4, + "median": 21.0, + "most_common": [{"length": 20, "count": 140}, {"length": 22, "count": 118}], + }, + "patterns": { + "top": [ + {"pattern": "aaaa@aaaa.aaa", "count": 512}, + {"pattern": "aaaa.aaaa@aaaa.aaa", "count": 300}, + ], + "composition": {"letter": 0.78, "digit": 0.05, "other": 0.17}, + }, + "first_chars": {"top": [{"masked": "j*******", "count": 40}, {"masked": "s*******", "count": 33}]}, + }, +} + +_NUMERIC_PREVIEW = { + "kind": "numeric", + "sampled_from": None, + "statistics_common": {"count_row": 1000, "count_null": 0, "count_distinct": 71}, + "statistics_kind": { + "summaries": {"mean": 41.3, "q1": 29.0, "q2": 41.0, "q3": 54.0, "p5": 19.0, "p95": 68.0}, + "histograms": { + "bins": [ + {"lower_bound": 18.0, "upper_bound": 28.0, "count": 210}, + {"lower_bound": 28.0, "upper_bound": 38.0, "count": 240}, + ] + }, + }, +} + +_TEMPORAL_PREVIEW = { + "kind": "temporal", + "sampled_from": None, + "statistics_common": {"count_row": 500, "count_null": 3, "count_distinct": 480}, + "statistics_kind": { + "summaries": { + "mean": "1987-04-12", + "q1": "1975-01-01", + "q2": "1988-06-30", + "q3": "2001-03-15", + "p5": "1960-01-01", + "p95": "2010-12-31", + }, + "histograms": { + "bins": [ + {"lower_bound": "1970-01-01", "upper_bound": "1980-01-01", "label": "1970s", "count": 88}, + {"lower_bound": "1980-01-01", "upper_bound": "1990-01-01", "label": "1980s", "count": 142}, + ] + }, + }, +} + +_BOOLEAN_PREVIEW = { + "kind": "boolean", + "sampled_from": None, + "statistics_common": {"count_row": 1000, "count_null": 0, "count_distinct": 2}, + "statistics_kind": {"count_true": 640, "count_false": 360}, +} + +_UNSUPPORTED_PREVIEW = { + "kind": "unsupported", + "sampled_from": None, + "statistics_common": {"count_row": 0, "count_null": 0, "count_distinct": None}, + "statistics_kind": {"reason": "Binary data isn't previewed"}, +} + + +def _assert_no_unexpected_extras(value: object) -> None: + """Recursively assert no field landed in a model's `model_extra` -- i.e. everything is typed.""" + if isinstance(value, BaseModel): + assert not value.model_extra, f"unexpected extra fields on {type(value).__name__}: {value.model_extra}" + for field_value in value.__dict__.values(): + _assert_no_unexpected_extras(field_value) + elif isinstance(value, list): + for item in value: + _assert_no_unexpected_extras(item) + + +def test_each_preview_parses_to_its_typed_variant(): + assert type(parse_safe_data_preview(_STRING_PREVIEW)) is StringPreview + assert type(parse_safe_data_preview(_NUMERIC_PREVIEW)) is NumericPreview + assert type(parse_safe_data_preview(_TEMPORAL_PREVIEW)) is TemporalPreview + assert type(parse_safe_data_preview(_BOOLEAN_PREVIEW)) is BooleanPreview + assert type(parse_safe_data_preview(_UNSUPPORTED_PREVIEW)) is UnsupportedPreview + + +def test_each_preview_round_trips(): + assert parse_safe_data_preview(_STRING_PREVIEW).model_dump(mode="json") == _STRING_PREVIEW + assert parse_safe_data_preview(_NUMERIC_PREVIEW).model_dump(mode="json") == _NUMERIC_PREVIEW + assert parse_safe_data_preview(_TEMPORAL_PREVIEW).model_dump(mode="json") == _TEMPORAL_PREVIEW + assert parse_safe_data_preview(_BOOLEAN_PREVIEW).model_dump(mode="json") == _BOOLEAN_PREVIEW + assert parse_safe_data_preview(_UNSUPPORTED_PREVIEW).model_dump(mode="json") == _UNSUPPORTED_PREVIEW + + +def test_each_preview_has_no_untyped_fields(): + _assert_no_unexpected_extras(parse_safe_data_preview(_STRING_PREVIEW)) + _assert_no_unexpected_extras(parse_safe_data_preview(_NUMERIC_PREVIEW)) + _assert_no_unexpected_extras(parse_safe_data_preview(_TEMPORAL_PREVIEW)) + _assert_no_unexpected_extras(parse_safe_data_preview(_BOOLEAN_PREVIEW)) + _assert_no_unexpected_extras(parse_safe_data_preview(_UNSUPPORTED_PREVIEW)) + + +def test_numeric_preview_typed_access(): + preview = parse_safe_data_preview(_NUMERIC_PREVIEW) + assert isinstance(preview, NumericPreview) + assert preview.statistics_common.count_row == 1000 + assert preview.statistics_kind.summaries.mean == 41.3 + assert preview.statistics_kind.histograms.bins[0].count == 210 + + +def test_string_preview_typed_access(): + preview = parse_safe_data_preview(_STRING_PREVIEW) + assert isinstance(preview, StringPreview) + assert preview.statistics_kind.lengths.most_common[0].length == 20 + assert preview.statistics_kind.patterns.composition.letter == 0.78 + assert preview.statistics_kind.first_chars.top[0].masked == "j*******" + + +def test_unsupported_preview_reason_reads_as_str_and_matches_known_constant(): + preview = parse_safe_data_preview(_UNSUPPORTED_PREVIEW) + assert isinstance(preview, UnsupportedPreview) + assert UnsupportedPreviewReason(preview.statistics_kind.reason) is UnsupportedPreviewReason.is_binary + assert preview.statistics_kind.reason == "Binary data isn't previewed" + + +def test_none_passes_through(): + assert parse_safe_data_preview(None) is None + + +def test_already_parsed_model_passes_through(): + preview = parse_safe_data_preview(_BOOLEAN_PREVIEW) + assert parse_safe_data_preview(preview) is preview + + +def test_unknown_kind_degrades(): + payload = { + "kind": "geospatial", + "sampled_from": None, + "statistics_common": {"count_row": 5, "count_null": 0, "count_distinct": 5}, + "statistics_kind": {"crs": "EPSG:4326"}, + } + preview = parse_safe_data_preview(payload) + assert isinstance(preview, UnknownPreview) + assert preview.kind == "geospatial" + assert preview.model_extra["statistics_kind"] == {"crs": "EPSG:4326"} + + +def test_known_kind_with_broken_shape_degrades(): + # `numeric`, but missing the required `statistics_kind`. + payload = { + "kind": "numeric", + "sampled_from": None, + "statistics_common": {"count_row": 5, "count_null": 0, "count_distinct": 5}, + } + preview = parse_safe_data_preview(payload) + assert isinstance(preview, UnknownPreview) + assert preview.kind == "numeric" + + +def test_unparseable_payload_falls_back_without_raising(): + payload = {"kind": 123, "sampled_from": "f.csv", "statistics_common": {"count_row": "oops"}} + preview = parse_safe_data_preview(payload) + assert isinstance(preview, UnknownPreview) + assert preview.kind == 123 + + +def test_non_dict_preview_is_none(): + assert parse_safe_data_preview("not a preview") is None + assert parse_safe_data_preview(42) is None + + +def _column_data(**overrides: object) -> dict[str, object]: + data: dict[str, object] = { + "data_type": "text", + "foreign_keys": [], + "discovery_matches": [], + "constraint_columns": [], + "unique_index_names": [], + "referencing_foreign_keys": [], + "constraint": "", + } + data.update(overrides) + return data + + +def test_column_null_preview_is_none(): + column = SchemaDiscoveryColumn.model_validate(_column_data(safe_data_preview=None)) + assert column.safe_data_preview is None + + +def test_column_absent_preview_is_none(): + column = SchemaDiscoveryColumn.model_validate(_column_data()) + assert column.safe_data_preview is None + + +def test_column_unsupported_preview_is_distinct_from_none(): + column = SchemaDiscoveryColumn.model_validate(_column_data(safe_data_preview=_UNSUPPORTED_PREVIEW)) + assert column.safe_data_preview is not None + assert isinstance(column.safe_data_preview, UnsupportedPreview) + + +def test_column_unknown_preview_does_not_break_the_result(): + column = SchemaDiscoveryColumn.model_validate( + _column_data(safe_data_preview={"kind": "geospatial", "statistics_common": {"count_row": 1, "count_null": 0}}) + ) + assert isinstance(column.safe_data_preview, UnknownPreview) + assert column.data_type == "text" # the rest of the column still parsed + + +def test_column_model_dump_preserves_preview_fields(): + # SerializeAsAny guard: dumping through the ColumnPreview-typed field must keep the concrete fields. + column = SchemaDiscoveryColumn.model_validate(_column_data(safe_data_preview=_NUMERIC_PREVIEW)) + dumped = column.model_dump(mode="json")["safe_data_preview"] + assert dumped["kind"] == "numeric" + assert dumped["statistics_kind"]["summaries"]["mean"] == 41.3 + + +def _file_discovery_payload(locator_extra: dict[str, object]) -> dict[str, object]: + return { + "id": 7, + "connection": {"id": "conn-1", "name": "my files"}, + "file_type": "csv", + "files": [{"path": "data/people.csv", "file_type": "csv"}], + "results": [{"locator": "age", "data_types": ["integer"], "matches": [], **locator_extra}], + } + + +def test_file_locator_preview_parses_with_sampled_from(): + payload = _file_discovery_payload({"safe_data_preview": {**_NUMERIC_PREVIEW, "sampled_from": "data/people.csv"}}) + result = FileDiscoveryResult.model_validate(payload) + preview = result.results[0].safe_data_preview + assert isinstance(preview, NumericPreview) + assert preview.sampled_from == "data/people.csv" + + +def test_file_locator_without_preview_is_none(): + result = FileDiscoveryResult.model_validate(_file_discovery_payload({})) + assert result.results[0].safe_data_preview is None From 5bcf818cf625349006058e747d1e541fdc6b9c20 Mon Sep 17 00:00:00 2001 From: Peter <101368063+ClassicMMT@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:08:03 +1200 Subject: [PATCH 2/3] feat: add row_count to schema-discovery table_metadata --- HISTORY.rst | 2 ++ datamasque/client/models/discovery.py | 1 + tests/test_discovery.py | 2 ++ tests/test_safe_data_preview.py | 3 +++ 4 files changed, 8 insertions(+) diff --git a/HISTORY.rst b/HISTORY.rst index 534eb63..c310632 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -12,6 +12,8 @@ History * ``safe_data_preview`` on schema-discovery result columns and file-discovery locators, typed by ``kind``. +* Added ``row_count`` to ``TableConstraints`` in schema-discovery ``table_metadata``. + 1.1.7 (2026-07-14) ------------------ diff --git a/datamasque/client/models/discovery.py b/datamasque/client/models/discovery.py index 3b69fe7..686c5e9 100644 --- a/datamasque/client/models/discovery.py +++ b/datamasque/client/models/discovery.py @@ -339,6 +339,7 @@ class TableConstraints(BaseModel): primary_keys: Optional[list[ConstraintColumns]] = None unique_keys: Optional[list[ConstraintColumns]] = None foreign_keys: Optional[list[ConstraintColumns]] = None + row_count: Optional[int] = None class SchemaDiscoveryPage(Page[SchemaDiscoveryResult]): diff --git a/tests/test_discovery.py b/tests/test_discovery.py index a9706d7..f23fc4d 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -870,6 +870,7 @@ def test_get_schema_discovery_page_returns_page_with_table_metadata(client): "primary_keys": [{"columns": ["id"]}], "unique_keys": [{"columns": ["email"]}], "foreign_keys": [], + "row_count": 12345, }, }, }, @@ -885,6 +886,7 @@ def test_get_schema_discovery_page_returns_page_with_table_metadata(client): assert isinstance(page, SchemaDiscoveryPage) assert [r.column for r in page.results] == ["email"] assert page.table_metadata["public"]["users"].primary_keys[0].columns == ["id"] + assert page.table_metadata["public"]["users"].row_count == 12345 assert m.last_request.qs == {"limit": ["10"], "offset": ["20"]} diff --git a/tests/test_safe_data_preview.py b/tests/test_safe_data_preview.py index 6d08109..100f81e 100644 --- a/tests/test_safe_data_preview.py +++ b/tests/test_safe_data_preview.py @@ -5,6 +5,7 @@ from datamasque.client import FileDiscoveryResult, SchemaDiscoveryColumn from datamasque.client.models.safe_data_preview import ( BooleanPreview, + CommonStatistics, NumericPreview, StringPreview, TemporalPreview, @@ -179,6 +180,8 @@ def test_known_kind_with_broken_shape_degrades(): preview = parse_safe_data_preview(payload) assert isinstance(preview, UnknownPreview) assert preview.kind == "numeric" + assert isinstance(preview.statistics_common, CommonStatistics) + assert preview.statistics_common.count_null == 0 def test_unparseable_payload_falls_back_without_raising(): From f01e0b20d666b2bf5c86a0a9fb0723b915c2b573 Mon Sep 17 00:00:00 2001 From: Peter <101368063+ClassicMMT@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:25:09 +1200 Subject: [PATCH 3/3] Fix: Remove UnknownPreview and Clean Up --- HISTORY.rst | 6 +- datamasque/client/__init__.py | 4 - datamasque/client/models/discovery.py | 18 +- datamasque/client/models/safe_data_preview.py | 74 +------- tests/test_discovery.py | 19 -- tests/test_safe_data_preview.py | 163 ++++-------------- 6 files changed, 50 insertions(+), 234 deletions(-) diff --git a/HISTORY.rst b/HISTORY.rst index c310632..4cceae4 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -7,10 +7,8 @@ History * Added typed Safe Data Preview support: - * ``safe_data_preview`` on ``InDataDiscoveryConfig`` (via ``SafeDataPreviewOptions``) - to configure or disable the preview. - * ``safe_data_preview`` on schema-discovery result columns and file-discovery locators, - typed by ``kind``. + * ``safe_data_preview`` on ``InDataDiscoveryConfig`` (via ``SafeDataPreviewOptions``) to configure or disable it. + * ``safe_data_preview`` on schema-discovery result columns and file-discovery locators, typed by ``kind``. * Added ``row_count`` to ``TableConstraints`` in schema-discovery ``table_metadata``. diff --git a/datamasque/client/__init__.py b/datamasque/client/__init__.py index 0229cd3..6ec133f 100644 --- a/datamasque/client/__init__.py +++ b/datamasque/client/__init__.py @@ -145,9 +145,7 @@ TemporalPreview, TemporalStatistics, TemporalSummaries, - UnknownPreview, UnsupportedPreview, - UnsupportedPreviewReason, UnsupportedStatistics, ) from datamasque.client.models.status import ( @@ -293,9 +291,7 @@ "TemporalStatistics", "TemporalSummaries", "UnfinishedRun", - "UnknownPreview", "UnsupportedPreview", - "UnsupportedPreviewReason", "UnsupportedStatistics", "User", "UserId", diff --git a/datamasque/client/models/discovery.py b/datamasque/client/models/discovery.py index 686c5e9..7fe8e6a 100644 --- a/datamasque/client/models/discovery.py +++ b/datamasque/client/models/discovery.py @@ -3,7 +3,7 @@ from enum import Enum from typing import Any, Optional, Union -from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from datamasque.client.models.connection import ConnectionConfig, ConnectionId, unwrap_connection_id from datamasque.client.models.data_selection import HashColumnsTableConfig, Locator, UserSelection @@ -11,10 +11,8 @@ from datamasque.client.models.pagination import Page from datamasque.client.models.runs import RunConnectionRef from datamasque.client.models.safe_data_preview import ( - ColumnPreview, SafeDataPreview, SafeDataPreviewOptions, - parse_safe_data_preview, ) @@ -303,13 +301,6 @@ class SchemaDiscoveryColumn(BaseModel): constraint: str # Primary or Unique, or empty string if column does not participate in a PK/UK safe_data_preview: Optional[SafeDataPreview] = None - @field_validator("safe_data_preview", mode="before") - @classmethod - def _parse_safe_data_preview( - cls, value: Union[dict[str, JsonValue], ColumnPreview, None] - ) -> Optional[ColumnPreview]: - return parse_safe_data_preview(value) - class SchemaDiscoveryResult(BaseModel): """A single row in the v2 schema discovery results.""" @@ -374,13 +365,6 @@ class FileDiscoveryLocatorResult(BaseModel): data_types: list[str] safe_data_preview: Optional[SafeDataPreview] = None - @field_validator("safe_data_preview", mode="before") - @classmethod - def _parse_safe_data_preview( - cls, value: Union[dict[str, JsonValue], ColumnPreview, None] - ) -> Optional[ColumnPreview]: - return parse_safe_data_preview(value) - class FileDiscoveryFile(BaseModel): """A file entry in a file discovery result.""" diff --git a/datamasque/client/models/safe_data_preview.py b/datamasque/client/models/safe_data_preview.py index a772eaa..258f0a2 100644 --- a/datamasque/client/models/safe_data_preview.py +++ b/datamasque/client/models/safe_data_preview.py @@ -1,14 +1,12 @@ """Typed request and response shapes for Safe Data Preview, part of In-Data Discovery.""" from enum import Enum -from typing import Literal, Optional, Union +from typing import Annotated, Literal, Optional, Union from pydantic import ( BaseModel, ConfigDict, - JsonValue, - SerializeAsAny, - ValidationError, + Field, ) @@ -49,20 +47,6 @@ class ColumnKind(str, Enum): unsupported = "unsupported" -class UnsupportedPreviewReason(str, Enum): - """Known reasons an `UnsupportedPreview` gives for declining a column.""" - - no_rows = "No rows sampled" - all_null = "All sampled values are null" - all_identical = "All sampled values are identical" - below_min_group_size = "A value count is below the minimum group size" - uninterpretable = "Values couldn't be interpreted as a single consistent type" - is_binary = "Binary data isn't previewed" - datatype_not_previewed = "This data type has no Safe Data Preview" - too_many_columns = "Table has too many columns for Safe Data Preview" - unavailable = "This column could not be previewed" - - class CommonStatistics(BaseModel): """Row counts shared by every preview kind.""" @@ -140,7 +124,7 @@ class FirstCharsStatistics(BaseModel): class StringStatistics(BaseModel): - """String preview statistics; `patterns`/`first_chars` appear only at the matching disclosure level.""" + """String preview statistics; `patterns`/`first_chars` appear at or above certain disclosure levels.""" model_config = ConfigDict(extra="allow") @@ -240,7 +224,7 @@ class BooleanStatistics(BaseModel): class UnsupportedStatistics(BaseModel): - """Why a column could not be previewed. See `UnsupportedPreviewReason` for known values.""" + """Why a column could not be previewed.""" model_config = ConfigDict(extra="allow") @@ -254,13 +238,13 @@ class ColumnPreview(BaseModel): # For grouped file discovery, the file whose sample produced this preview; None for database columns. sampled_from: Optional[str] = None + statistics_common: CommonStatistics class StringPreview(ColumnPreview): """A preview of a string column.""" kind: Literal[ColumnKind.string] = ColumnKind.string - statistics_common: CommonStatistics statistics_kind: StringStatistics @@ -268,7 +252,6 @@ class NumericPreview(ColumnPreview): """A preview of a numeric column.""" kind: Literal[ColumnKind.numeric] = ColumnKind.numeric - statistics_common: CommonStatistics statistics_kind: NumericStatistics @@ -276,7 +259,6 @@ class TemporalPreview(ColumnPreview): """A preview of a date/time column.""" kind: Literal[ColumnKind.temporal] = ColumnKind.temporal - statistics_common: CommonStatistics statistics_kind: TemporalStatistics @@ -284,7 +266,6 @@ class BooleanPreview(ColumnPreview): """A preview of a boolean column.""" kind: Literal[ColumnKind.boolean] = ColumnKind.boolean - statistics_common: CommonStatistics statistics_kind: BooleanStatistics @@ -292,47 +273,10 @@ class UnsupportedPreview(ColumnPreview): """A column the preview ran against but declined; `statistics_kind.reason` says why.""" kind: Literal[ColumnKind.unsupported] = ColumnKind.unsupported - statistics_common: CommonStatistics statistics_kind: UnsupportedStatistics -class UnknownPreview(ColumnPreview): - """A Safe Data Preview whose `kind` or shape this client can't parse; raw fields stay on `model_extra`.""" - - kind: Optional[str] = None - statistics_common: Optional[CommonStatistics] = None - - -PREVIEW_BY_KIND: dict[str, type[ColumnPreview]] = { - ColumnKind.string.value: StringPreview, - ColumnKind.numeric.value: NumericPreview, - ColumnKind.temporal.value: TemporalPreview, - ColumnKind.boolean.value: BooleanPreview, - ColumnKind.unsupported.value: UnsupportedPreview, -} - - -def parse_safe_data_preview(value: Union[dict[str, JsonValue], ColumnPreview, None]) -> Optional[ColumnPreview]: - """Parse a raw preview into its typed variant, or `UnknownPreview` if the kind or shape is unrecognised.""" - - if value is None: - return None - if isinstance(value, ColumnPreview): - return value - if not isinstance(value, dict): - return None - kind = value.get("kind") - preview_type = PREVIEW_BY_KIND.get(kind) if isinstance(kind, str) else None - if preview_type is not None: - try: - return preview_type.model_validate(value) - except ValidationError: - pass - try: - return UnknownPreview.model_validate(value) - except ValidationError: - # Fallback to raw - return UnknownPreview.model_construct(**value) # type: ignore[arg-type] - - -SafeDataPreview = SerializeAsAny[ColumnPreview] +SafeDataPreview = Annotated[ + Union[StringPreview, NumericPreview, TemporalPreview, BooleanPreview, UnsupportedPreview], + Field(discriminator="kind"), +] diff --git a/tests/test_discovery.py b/tests/test_discovery.py index f23fc4d..3318b64 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -22,12 +22,10 @@ InDataDiscoveryConfig, RulesetGenerationRequest, RunId, - SafeDataPreviewOptions, SchemaDiscoveryFromConfigRequest, SchemaDiscoveryPage, SchemaDiscoveryRequest, SchemaDiscoveryResult, - StringDisclosureLevel, StringPreview, ) from datamasque.client.exceptions import ( @@ -651,18 +649,6 @@ def test_schema_discovery_request_model_dump_includes_set_fields(): } -def test_schema_discovery_request_model_dump_includes_safe_data_preview(): - req = SchemaDiscoveryRequest( - connection="conn-1", - in_data_discovery=InDataDiscoveryConfig( - enabled=True, - safe_data_preview=SafeDataPreviewOptions(enabled=False, string_level=StringDisclosureLevel.lengths), - ), - ) - dumped = req.model_dump(exclude_none=True, mode="json") - assert dumped["in_data_discovery"]["safe_data_preview"] == {"enabled": False, "string_level": "lengths"} - - def test_discovery_requests_accept_connection_config_objects(): """All three discovery request models accept a full `ConnectionConfig` and extract its `id`.""" connection = DatabaseConnectionConfig( @@ -802,11 +788,6 @@ def test_schema_discovery_result_parses_safe_data_preview(): assert preview.statistics_common.count_distinct == 988 -def test_schema_discovery_result_without_safe_data_preview(): - result = SchemaDiscoveryResult.model_validate(_schema_discovery_row(1, "email")) - assert result.data.safe_data_preview is None - - def test_list_schema_discovery_results_follows_pagination(client): run_id = RunId(42) page1 = { diff --git a/tests/test_safe_data_preview.py b/tests/test_safe_data_preview.py index 100f81e..f346a70 100644 --- a/tests/test_safe_data_preview.py +++ b/tests/test_safe_data_preview.py @@ -1,18 +1,12 @@ -"""Tests for the typed Safe Data Preview models and their tolerant parsing.""" +"""Tests for the typed Safe Data Preview models.""" -from pydantic import BaseModel +import pytest +from pydantic import BaseModel, ValidationError from datamasque.client import FileDiscoveryResult, SchemaDiscoveryColumn from datamasque.client.models.safe_data_preview import ( - BooleanPreview, - CommonStatistics, NumericPreview, - StringPreview, - TemporalPreview, - UnknownPreview, UnsupportedPreview, - UnsupportedPreviewReason, - parse_safe_data_preview, ) _STRING_PREVIEW = { @@ -89,6 +83,14 @@ "statistics_kind": {"reason": "Binary data isn't previewed"}, } +_KNOWN_PREVIEWS = [ + pytest.param(_STRING_PREVIEW, id="string"), + pytest.param(_NUMERIC_PREVIEW, id="numeric"), + pytest.param(_TEMPORAL_PREVIEW, id="temporal"), + pytest.param(_BOOLEAN_PREVIEW, id="boolean"), + pytest.param(_UNSUPPORTED_PREVIEW, id="unsupported"), +] + def _assert_no_unexpected_extras(value: object) -> None: """Recursively assert no field landed in a model's `model_extra` -- i.e. everything is typed.""" @@ -101,102 +103,7 @@ def _assert_no_unexpected_extras(value: object) -> None: _assert_no_unexpected_extras(item) -def test_each_preview_parses_to_its_typed_variant(): - assert type(parse_safe_data_preview(_STRING_PREVIEW)) is StringPreview - assert type(parse_safe_data_preview(_NUMERIC_PREVIEW)) is NumericPreview - assert type(parse_safe_data_preview(_TEMPORAL_PREVIEW)) is TemporalPreview - assert type(parse_safe_data_preview(_BOOLEAN_PREVIEW)) is BooleanPreview - assert type(parse_safe_data_preview(_UNSUPPORTED_PREVIEW)) is UnsupportedPreview - - -def test_each_preview_round_trips(): - assert parse_safe_data_preview(_STRING_PREVIEW).model_dump(mode="json") == _STRING_PREVIEW - assert parse_safe_data_preview(_NUMERIC_PREVIEW).model_dump(mode="json") == _NUMERIC_PREVIEW - assert parse_safe_data_preview(_TEMPORAL_PREVIEW).model_dump(mode="json") == _TEMPORAL_PREVIEW - assert parse_safe_data_preview(_BOOLEAN_PREVIEW).model_dump(mode="json") == _BOOLEAN_PREVIEW - assert parse_safe_data_preview(_UNSUPPORTED_PREVIEW).model_dump(mode="json") == _UNSUPPORTED_PREVIEW - - -def test_each_preview_has_no_untyped_fields(): - _assert_no_unexpected_extras(parse_safe_data_preview(_STRING_PREVIEW)) - _assert_no_unexpected_extras(parse_safe_data_preview(_NUMERIC_PREVIEW)) - _assert_no_unexpected_extras(parse_safe_data_preview(_TEMPORAL_PREVIEW)) - _assert_no_unexpected_extras(parse_safe_data_preview(_BOOLEAN_PREVIEW)) - _assert_no_unexpected_extras(parse_safe_data_preview(_UNSUPPORTED_PREVIEW)) - - -def test_numeric_preview_typed_access(): - preview = parse_safe_data_preview(_NUMERIC_PREVIEW) - assert isinstance(preview, NumericPreview) - assert preview.statistics_common.count_row == 1000 - assert preview.statistics_kind.summaries.mean == 41.3 - assert preview.statistics_kind.histograms.bins[0].count == 210 - - -def test_string_preview_typed_access(): - preview = parse_safe_data_preview(_STRING_PREVIEW) - assert isinstance(preview, StringPreview) - assert preview.statistics_kind.lengths.most_common[0].length == 20 - assert preview.statistics_kind.patterns.composition.letter == 0.78 - assert preview.statistics_kind.first_chars.top[0].masked == "j*******" - - -def test_unsupported_preview_reason_reads_as_str_and_matches_known_constant(): - preview = parse_safe_data_preview(_UNSUPPORTED_PREVIEW) - assert isinstance(preview, UnsupportedPreview) - assert UnsupportedPreviewReason(preview.statistics_kind.reason) is UnsupportedPreviewReason.is_binary - assert preview.statistics_kind.reason == "Binary data isn't previewed" - - -def test_none_passes_through(): - assert parse_safe_data_preview(None) is None - - -def test_already_parsed_model_passes_through(): - preview = parse_safe_data_preview(_BOOLEAN_PREVIEW) - assert parse_safe_data_preview(preview) is preview - - -def test_unknown_kind_degrades(): - payload = { - "kind": "geospatial", - "sampled_from": None, - "statistics_common": {"count_row": 5, "count_null": 0, "count_distinct": 5}, - "statistics_kind": {"crs": "EPSG:4326"}, - } - preview = parse_safe_data_preview(payload) - assert isinstance(preview, UnknownPreview) - assert preview.kind == "geospatial" - assert preview.model_extra["statistics_kind"] == {"crs": "EPSG:4326"} - - -def test_known_kind_with_broken_shape_degrades(): - # `numeric`, but missing the required `statistics_kind`. - payload = { - "kind": "numeric", - "sampled_from": None, - "statistics_common": {"count_row": 5, "count_null": 0, "count_distinct": 5}, - } - preview = parse_safe_data_preview(payload) - assert isinstance(preview, UnknownPreview) - assert preview.kind == "numeric" - assert isinstance(preview.statistics_common, CommonStatistics) - assert preview.statistics_common.count_null == 0 - - -def test_unparseable_payload_falls_back_without_raising(): - payload = {"kind": 123, "sampled_from": "f.csv", "statistics_common": {"count_row": "oops"}} - preview = parse_safe_data_preview(payload) - assert isinstance(preview, UnknownPreview) - assert preview.kind == 123 - - -def test_non_dict_preview_is_none(): - assert parse_safe_data_preview("not a preview") is None - assert parse_safe_data_preview(42) is None - - -def _column_data(**overrides: object) -> dict[str, object]: +def _build_column_data(**overrides: object) -> dict[str, object]: data: dict[str, object] = { "data_type": "text", "foreign_keys": [], @@ -210,36 +117,42 @@ def _column_data(**overrides: object) -> dict[str, object]: return data -def test_column_null_preview_is_none(): - column = SchemaDiscoveryColumn.model_validate(_column_data(safe_data_preview=None)) - assert column.safe_data_preview is None +@pytest.mark.parametrize("payload", _KNOWN_PREVIEWS) +def test_each_preview_has_no_untyped_fields(payload): + column = SchemaDiscoveryColumn.model_validate(_build_column_data(safe_data_preview=payload)) + _assert_no_unexpected_extras(column.safe_data_preview) -def test_column_absent_preview_is_none(): - column = SchemaDiscoveryColumn.model_validate(_column_data()) +@pytest.mark.parametrize( + "overrides", + [ + pytest.param({"safe_data_preview": None}, id="explicit-null"), + pytest.param({}, id="absent"), + ], +) +def test_column_without_preview_is_none(overrides): + column = SchemaDiscoveryColumn.model_validate(_build_column_data(**overrides)) assert column.safe_data_preview is None def test_column_unsupported_preview_is_distinct_from_none(): - column = SchemaDiscoveryColumn.model_validate(_column_data(safe_data_preview=_UNSUPPORTED_PREVIEW)) + column = SchemaDiscoveryColumn.model_validate(_build_column_data(safe_data_preview=_UNSUPPORTED_PREVIEW)) assert column.safe_data_preview is not None assert isinstance(column.safe_data_preview, UnsupportedPreview) -def test_column_unknown_preview_does_not_break_the_result(): - column = SchemaDiscoveryColumn.model_validate( - _column_data(safe_data_preview={"kind": "geospatial", "statistics_common": {"count_row": 1, "count_null": 0}}) - ) - assert isinstance(column.safe_data_preview, UnknownPreview) - assert column.data_type == "text" # the rest of the column still parsed - - -def test_column_model_dump_preserves_preview_fields(): - # SerializeAsAny guard: dumping through the ColumnPreview-typed field must keep the concrete fields. - column = SchemaDiscoveryColumn.model_validate(_column_data(safe_data_preview=_NUMERIC_PREVIEW)) - dumped = column.model_dump(mode="json")["safe_data_preview"] - assert dumped["kind"] == "numeric" - assert dumped["statistics_kind"]["summaries"]["mean"] == 41.3 +@pytest.mark.parametrize( + "safe_data_preview", + [ + pytest.param({"kind": "geospatial", "statistics_common": {"count_row": 1, "count_null": 0}}, id="unknown-kind"), + pytest.param({"kind": "numeric", "statistics_common": {"count_row": 5, "count_null": 0}}, id="broken-shape"), + pytest.param({"kind": 123, "statistics_common": {"count_row": "oops"}}, id="malformed"), + pytest.param("not a preview", id="non-dict"), + ], +) +def test_column_rejects_unparseable_preview(safe_data_preview): + with pytest.raises(ValidationError): + SchemaDiscoveryColumn.model_validate(_build_column_data(safe_data_preview=safe_data_preview)) def _file_discovery_payload(locator_extra: dict[str, object]) -> dict[str, object]: