diff --git a/.claude/skills/slayer-overview.md b/.claude/skills/slayer-overview.md index 353955ab..7a076b04 100644 --- a/.claude/skills/slayer-overview.md +++ b/.claude/skills/slayer-overview.md @@ -12,7 +12,7 @@ SLayer is a lightweight, agent-first semantic layer. Instead of writing raw SQL, - **SQLGenerator** — takes an EnrichedQuery (not SlayerQuery) and converts it to SQL via sqlglot (dialect-aware: postgres, mysql, bigquery, etc.) - **SlayerSQLClient** — executes SQL via SQLAlchemy with retry logic and statement timeouts - **Storage** — YAML or SQLite backends for model and datasource configs -- **Ingestion** — auto-generates models from DB schema with rollup-style FK joins (denormalized LEFT JOINs). It can be triggered manually (`slayer ingest`, `ingest_datasource_models`, `POST /ingest`) or **on every server boot** via `slayer serve --ingest-on-startup` / `slayer mcp --ingest-on-startup` (also `SLAYER_INGEST_ON_STARTUP=1`, or `create_app/create_mcp_server(ingest_on_startup=True)` programmatically). It is idempotent and continues on per-datasource failures. +- **Ingestion** — auto-generates models from DB schema with rollup-style FK joins (denormalized LEFT JOINs), importing DB column/table comments into `description`s (fill-if-empty, never overwriting; BigQuery also fills the datasource description from the dataset description; SQLite has no comments). It can be triggered manually (`slayer ingest`, `ingest_datasource_models`, `POST /ingest`) or **on every server boot** via `slayer serve --ingest-on-startup` / `slayer mcp --ingest-on-startup` (also `SLAYER_INGEST_ON_STARTUP=1`, or `create_app/create_mcp_server(ingest_on_startup=True)` programmatically). It is idempotent and continues on per-datasource failures. - **Interfaces** — MCP server (stdio via `slayer mcp`, SSE via `slayer serve` at `/mcp/sse`), REST API (FastAPI on port 5143), Python SDK, and two read-only wire-protocol facades for BI tools: Arrow Flight SQL (`slayer flight-serve`, port 5144) and Postgres (`slayer pg-serve`, port 5145; the connection `database` selects the SLayer datasource) ## Key Models diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 26af9d5c..18e47b08 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -141,6 +141,13 @@ jobs: timeout-minutes: 5 run: poetry run python examples/bigquery/verify.py + - name: Run BigQuery integration tests + if: steps.gate.outputs.skip != 'true' + timeout-minutes: 10 + env: + GCP_PROJECT_ID: ${{ secrets.GCP_PROJECT_ID }} + run: poetry run pytest tests/integration/test_integration_bigquery.py -m integration -v + - name: Dump server logs on failure if: failure() && steps.gate.outputs.skip != 'true' run: cat /tmp/slayer-bq.log diff --git a/DECISIONS.md b/DECISIONS.md index 9dbf0a97..f32d0fda 100644 --- a/DECISIONS.md +++ b/DECISIONS.md @@ -78,3 +78,4 @@ implementation detail. Include issue refs when known. - 2026-08-12 — Dotted dimension join-path binding (DEV-1780): a dotted dimension/time-dimension path resolves only when every hop is a direct join. Previously a hop that was not a direct join fell through leniently — the enriched dim kept its `A__B` alias in SELECT/GROUP BY but `_resolve_joins` emitted no join, shipping invalid SQL (unbound table alias). Filters and cross-model measures already rejected such paths; only dimensions/time-dimensions had the hole (the shared `_resolve_dotted_dim_with_stage_fallback` lenient branch). Fix is an engine routing pre-pass (`SlayerQueryEngine._route_dotted_dimension_refs`, run in `_enrich` before `enrich_query`, gated on `enforce_join_binding and source_model_origin is None`): it normalizes root-prefixes via `strip_source_model_prefix`, then for each dotted ref tries the explicit direct-join walk and, on `_NoJoinError`, routes via a datasource-scoped `JoinGraph`. A SHORT FORM (one model segment, e.g. `Consumer.name`) with exactly ONE route to the target auto-resolves — the ref is rewritten to the full routed path, so the result key is the full path (`root.Subscription.Customer.Consumer.name`), consistent with "joined dims keep the full path". Ambiguous (≥2 routes), unreachable (0), and explicit multi-hop chains with a broken hop are REJECTED with `UnresolvableDimensionJoinError(SlayerError, ValueError)` (mirrors the DEV-1645 `UnresolvableOrderColumnError` reject-don't-emit-invalid-SQL doctrine); the message suggests the short form when the target is uniquely reachable, else the shortest deterministic full path (`JoinGraph.shortest_path`), else nothing. `JoinGraph.count_simple_paths(root, target, cap=2)` classifies routes — it counts ALL simple paths (a 2-hop + 3-hop route is genuinely ambiguous; auto-picking the shorter would silently change join semantics), reverse-reachability-pruned and cycle-guarded. The rewrite map is also applied to matching `OrderItem.column` refs and `main_time_dimension` so dependent references stay consistent. Deliberate limits (conservative, prefer reject over a wrong route): routing runs only within a single datasource (`model.data_source` truthy; the graph is datasource-scoped) and is deferred when named-query stages are in scope (their virtual models aren't in the stored graph) — those refs fall through to the guard. A post-`_resolve_joins` safety-net guard in `enrich_query` (same gate) raises `UnresolvableDimensionJoinError` for any dim/time-dim whose alias is absent from `resolved_joins`, guaranteeing the invariant even for direct `enrich_query` callers; the re-rooted cross-model CTE enrichment passes `enforce_join_binding=False` (it legitimately carries source-local shared dims like `orders.status` that never bind to a base-table join). Out of scope: the multi-stage lenient cross-stage fall-through (`test_unresolvable_dotted_ref_falls_through`, where distinguishing a genuine error from a re-rooting artifact is unsolved) and leaf-column-missing-on-a-valid-path (the alias IS bound there — a different failure class). - 2026-08-16 — FK-derived joins name the MODEL, not the live object (DEV-1688 / DEV-1741 / #279). Model names strip `__` (reserved for join paths), so an FK to `reports__patient__drug` used to persist a join targeting a model that cannot exist; `_generate_joins` now takes the live→model map, and a target whose object was skipped on a sanitization collision drops its join rather than dangling. Stores written before the fix self-heal on the normal re-ingest path rather than via a schema migration — no version bump, and the repair demands the sanitized target AND identical `join_pairs` to match a freshly-generated join, so it can only rename the join the bug produced (name-only matching would collapse `a__b` and `a___b` onto one target and trip the duplicate-target guard, turning a merely-dangling store into a failed re-ingest). A store that never re-ingests keeps a join that was already broken. - 2026-08-18 — Two BigQuery-only failures fixed. (1) The DEV-1444 outer wrap re-parses already-emitted SQL, and BigQuery parses a quoted dotted alias (`` `orders.created_at` ``) into one part per segment, so the old qualifier strip (`col.set("table", None)`) both emitted an empty backtick pair — `400 Syntax error: Invalid empty identifier` on any computed measure plus ORDER BY — and, once that was fixed by replacing the node, silently dropped the model prefix, turning a syntax error into an unresolved name. `SqlDialect._outer_order_column` therefore re-resolves the column by picking the LONGEST part-suffix the inner SELECT actually projects (checked as a quoted identifier against `inner_sql`, which at that point still carries canonical dotted aliases — `rewrite_emitted_sql` mangles them afterwards). That subsumes the `_base.`-qualified form `_assemble_combined_sql` emits and keeps Postgres/DuckDB/MySQL output byte-identical, including the untouched bare-column and no-match fallbacks; hidden ORDER-BY hoists resolve too, which a `public`-list match would have missed. (2) `_get_columns_fallback` queried the bare `information_schema.columns`, which BigQuery resolves as `.information_schema.columns` — a project-level view a dataset-scoped service account (the normal least-privilege setup) cannot read, so every per-table introspection 403'd. It now qualifies by dataset (`` ``.INFORMATION_SCHEMA.COLUMNS ``), taken from `schema` or from the dotted table name, and drops the now-redundant `table_schema` predicate. Related hardening: `_live_schema_for_datasource` raises `IntrospectionUnavailable` when EVERY table in a datasource failed rather than returning an empty map, because empty is indistinguishable from "every table was dropped" — drift then reported a `WholeModelDelete` for every model in the tenant off one credential error, and `--force-clean` would act on it. `_collect_sql_table_diffs` catches it and returns no verdict; type refinement catches it and keeps the persisted types. +- 2026-08-20 — DB comments imported at ingestion (DEV-1809 / #317): column comments → `Column.description`, table comments → `SlayerModel.description`, BigQuery dataset description → `DatasourceConfig.description` — strictly fill-if-empty (DEV-1356 additive-only; no provenance tracking, so a DB comment edit never propagates over an existing description — clear it and re-ingest). Inspector path covers Postgres/MySQL/SQL Server/Snowflake/ClickHouse/BigQuery; per-dialect fallback comment SQL for MySQL/Snowflake/ClickHouse/DuckDB/Postgres (DuckDB reachable only via the fallback; SQL Server/BigQuery fallbacks stay types-only — disproportionate SQL for a path their Inspector already covers). Schema-level comments have no generic API, hence BigQuery-only datasource descriptions; comment fetching is always best-effort, never a failed ingest; dbt/OSI curated descriptions win over DB comments at creation time; drift detection ignores comment text. diff --git a/docs/concepts/ingestion.md b/docs/concepts/ingestion.md index 89fa505b..62c26fb4 100644 --- a/docs/concepts/ingestion.md +++ b/docs/concepts/ingestion.md @@ -42,6 +42,7 @@ Tables with no FK references use their plain table name with no joins. SLayer introspects each table's column types and generates a model: - **One `Column`** per non-joined column on the source table — name, `type` inferred from the database (`string` / `number` / `boolean` / `time` / `date`), `primary_key=True` for PKs. Whether each column is used as a group-by dimension or as an aggregation source is decided per query. +- **Column and table comments** stored in the database become `Column.description` and `SlayerModel.description` (see [Comments and descriptions](#comments-and-descriptions) below). - **`unique=True`** for columns that alone form a `UNIQUE` constraint or unique index. PK columns are not stamped redundantly — `primary_key` already implies uniqueness. Composite uniqueness is evaluated per key-set during join-cardinality inference rather than being flattened onto individual columns. - **A column literally named `count`** is renamed to `count_col` to avoid clashing with the always-available `*:count`. - **No auto-generated `measures`** — `SlayerModel.measures` is the named-formula library and stays empty after ingestion. You can add named formulas later via the API/MCP if you want bare-name shortcuts (`{"formula": "aov"}`). @@ -54,6 +55,36 @@ All models use `sql_table` (the source table) plus `joins` (direct FK joins only Each FK join also gets a structural [`cardinality`](models.md#join-cardinality) guess from the key constraints alone (no data is read): `many_to_one` by default, upgrading to `one_to_one` when the source key is itself unique. A side counts as unique only when some PK/unique key-set is a subset of the join key — if `(a)` is unique then `(a, b)` is too, but a constraint on `(a, b)` does not make `(a)` unique. When the target key cannot be *verified* unique from its constraints, cardinality is left unset rather than guessed. To infer it from the data instead, run `slayer validate-models --cardinality`. +### Comments and descriptions + +Textual metadata stored in the database is imported during ingestion, so +agents see the same documentation the DBA wrote: + +- **Column comments** → `Column.description` +- **Table comments** → `SlayerModel.description` +- **BigQuery dataset description** → the datasource's `description` + (BigQuery only — no cross-database API exists for schema-level comments) + +The import is strictly **fill-if-empty**: a comment lands only where the +existing `description` is empty. Hand-written descriptions are never +overwritten, on first ingest or any re-ingest — consistent with the +additive-only re-ingestion contract. To re-import a comment, clear the +description and re-run `slayer ingest`. Imported descriptions feed the +search corpus like any other description, so the next embedding refresh +picks them up. + +Per-database support: + +| Database | Column comments | Table comments | +|---|---|---| +| BigQuery | ✓ (field descriptions) | ✓ (table description) | +| Postgres, MySQL, SQL Server, Snowflake, ClickHouse | ✓ | ✓ | +| DuckDB | ✓ (`COMMENT ON`) | ✓ | +| SQLite | — (no such feature) | — | + +Comment fetching is best-effort: a database or driver that cannot surface +comments simply yields models without descriptions, never a failed ingest. + ### SQLite affinity probing SQLite's declared column types are affinity hints, not strict constraints: a column declared `INTEGER` can store `INTEGER`, `REAL`, `TEXT`, or `BLOB` values per row. To prevent silent truncation downstream (a column declared `INTEGER` but actually storing `0.99` would cast to `0` and break `AVG`/`SUM` results), SLayer runs an additional value-level probe on SQLite ingestion for every column the inspector reports as `INTEGER`-affinity. @@ -276,7 +307,7 @@ Ingest-on-startup: N/M datasources ingested (K failed: name1, name2) `slayer ingest` (and the equivalent MCP / REST entry points) is idempotent by default — re-runs are safe. For each in-scope live table: - **No persisted model with that name** → ingest from scratch via the path above. -- **Existing `sql_table`-mode model** → append new columns and joins from the live schema. Existing user metadata is **never** overwritten — `description`, `label`, `format`, `meta`, and `allowed_aggregations` are preserved verbatim. The only in-place updates are strictly additive gap-fills: a join's `cardinality` is set only when it is currently unset (a value you chose is never replaced), and a column's `unique` is only ever turned on, never off. Filling either one is enough to trigger a save, so a re-ingest that adds no columns or joins still persists newly-discovered constraint metadata. +- **Existing `sql_table`-mode model** → append new columns and joins from the live schema. Existing user metadata is **never** overwritten — `description`, `label`, `format`, `meta`, and `allowed_aggregations` are preserved verbatim. The only in-place updates are strictly additive gap-fills: a join's `cardinality` is set only when it is currently unset (a value you chose is never replaced), a column's `unique` is only ever turned on, never off, and an **empty** `description` (column and model level) is filled from the live DB comment. Filling any of these is enough to trigger a save, so a re-ingest that adds no columns or joins still persists newly-discovered metadata. Filled descriptions are reported per model (`ModelAddition.described_columns`, `model_described`; the CLI prints `+descriptions: …`). - **Existing `sql`-mode or query-backed model with the matching name** → skipped silently; those are user-authored. With the default YAML storage, two live tables whose quoted names differ only by letter case (`"Orders"` vs `orders`) cannot both be persisted — model names collide as filenames on macOS / Windows, so the save is rejected (`IdCollisionError`). The first table wins; the second surfaces as a per-model entry in `IdempotentIngestResult.errors` (or a per-model message on the CLI / MCP paths) without aborting the rest of the ingest. SQLite storage persists both. diff --git a/docs/database-support.md b/docs/database-support.md index f97b711c..067a47a0 100644 --- a/docs/database-support.md +++ b/docs/database-support.md @@ -19,14 +19,17 @@ against the live service in CI when they are. | **ClickHouse** | `tests/integration/test_integration_clickhouse.py` (`testcontainers[clickhouse]`) | `examples/clickhouse/` | | **SQL Server** | `tests/integration/test_integration_sqlserver.py` (`testcontainers`, `msodbcsql18` + `unixodbc-dev` on the runner) | `examples/sqlserver/` | | **Snowflake** | `tests/integration/test_integration_snowflake.py` (skips without `~/.snowflake/connections.toml`; profile name overridable via `$SLAYER_SNOWFLAKE_CONNECTION`) | `examples/snowflake/` (no Docker) | -| **BigQuery** | `examples/bigquery/verify.py` driven by CI against `bigquery-public-data.thelook_ecommerce` (gated on `GCP_PROJECT_ID` / `GCP_SA_KEY_B64` repo secrets) | `examples/bigquery/` (no Docker — managed service) | - -BigQuery does not yet have a pytest-style integration suite; its CI coverage -runs the example's `verify.py` directly via `.github/workflows/ci.yml`. That -exercises auto-ingestion, basic projection, joins, time-grain dimensions, and -the cardinality / sum-of-grouped-equals-total invariants — enough to catch -emitted-SQL regressions, but the verify-script tier is shallower than the -testcontainers suites. +| **BigQuery** | `tests/integration/test_integration_bigquery.py` (live ingestion against a temp dataset in the billing project; skips without `GCP_PROJECT_ID` + ADC) plus `examples/bigquery/verify.py` driven by CI against `bigquery-public-data.thelook_ecommerce` (gated on `GCP_PROJECT_ID` / `GCP_SA_KEY_B64` repo secrets) | `examples/bigquery/` (no Docker — managed service) | + +BigQuery CI coverage has two layers, both in the `bigquery-example` job of +`.github/workflows/ci.yml`: the example's `verify.py` exercises query +execution against the public dataset (basic projection, joins, time-grain +dimensions, and the cardinality / sum-of-grouped-equals-total invariants), +and `test_integration_bigquery.py` exercises live schema ingestion — creating +a temporary dataset with table/column/dataset descriptions in the billing +project, asserting comment import and idempotent re-ingest semantics, then +deleting it. The service account therefore needs dataset create/delete rights +in the billing project (e.g. `roles/bigquery.user`), not just `jobUser`. ## Tier 2 — code-covered @@ -87,6 +90,10 @@ mapping lives in `SqlDialect.build_approx_count_distinct`. ### SQLite caveats +SQLite has no table or column comments, so ingestion imports no descriptions +there (every other Tier 1 database's comments are imported — see +[Ingestion — Comments and descriptions](concepts/ingestion.md#comments-and-descriptions)). + SQLite has a much smaller built-in math/stat catalog than the other supported engines. SLayer registers Python aggregate and scalar UDFs on every new SQLite connection via SQLAlchemy's `connect` event (see @@ -225,6 +232,11 @@ plus `$GCP_PROJECT_ID` for billing). The `bigquery://` driver requires the - **No FK introspection.** BigQuery exposes no foreign-key metadata via `INFORMATION_SCHEMA`, so auto-ingestion cannot discover joins. Hand-declare `ModelJoin`s on the model. +- **Descriptions ARE imported.** Auto-ingestion maps BigQuery column + descriptions → `Column.description`, table descriptions → + `SlayerModel.description`, and the dataset description → the datasource's + `description` (fill-if-empty; see + [Ingestion — Comments and descriptions](concepts/ingestion.md#comments-and-descriptions)). - **Dotted alias mangling.** BigQuery rejects column names containing `.` (output schema names must match `[A-Za-z_][A-Za-z0-9_]*`), so SLayer rewrites `.` aliases (`orders._count`, diff --git a/slayer/cli.py b/slayer/cli.py index 835c07f3..0a2ad36d 100644 --- a/slayer/cli.py +++ b/slayer/cli.py @@ -2237,6 +2237,15 @@ def _run_datasources_create(args, storage): sys.exit(1) _persist_ingested_models(report.models, storage, assume_yes=args.yes) + + if report.schema_description and not ds.description: + ds = ds.model_copy(update={"description": report.schema_description}) + try: + run_sync(storage.save_datasource(ds)) + print("Datasource description imported.") + except Exception as e: + print(f"Could not save datasource description: {e}") + # After persistence, so the sections comment on what was just written. Exit # stays 0 — creating the datasource succeeded, whatever was skipped/hidden. _print_ingest_drift_and_errors(report, data_source=ds.name) diff --git a/slayer/dbt/converter.py b/slayer/dbt/converter.py index f6a57ca1..9e022f9e 100644 --- a/slayer/dbt/converter.py +++ b/slayer/dbt/converter.py @@ -323,11 +323,14 @@ def _convert_regular_model( if rm.description: model.description = rm.description + # Curated dbt descriptions win over freshly introspected DB comments; + # DB comments fill only the gaps (DEV-1809 — creation-time overlay, + # no persisted user edits exist yet). col_descriptions = {c.name: c.description for c in rm.columns if c.description} if col_descriptions: for c in model.columns: desc = col_descriptions.get(c.name) - if desc and not c.description: + if desc: c.description = desc return model diff --git a/slayer/engine/ingestion.py b/slayer/engine/ingestion.py index 394cfecf..946b0a6a 100644 --- a/slayer/engine/ingestion.py +++ b/slayer/engine/ingestion.py @@ -35,6 +35,7 @@ from slayer.engine.introspect_utils import ( # noqa: F401 (re-exported for back-compat) _FLOAT_LIKE_INFO_SCHEMA_TYPES, _INFO_SCHEMA_TYPE_MAP, + _clean_comment, _get_columns_fallback, _parse_info_schema_is_float, _safe_get_columns, @@ -56,6 +57,18 @@ logger = logging.getLogger(__name__) + +class IntrospectedColumn(BaseModel): + """One column as read from the live database during introspection.""" + + name: str + type: DataType + primary_key: bool = False + is_float: bool = False + db_type: str | None = None + comment: str | None = None + + # Module-level dedup set for unrecognized SA type warnings (see # _sa_type_to_data_type). Keyed by upper-cased class name. _logged_unmapped_sa_types: set[str] = set() @@ -778,6 +791,20 @@ def _safe_get_pk_constraint( return {"constrained_columns": []} +def _safe_get_table_comment( + inspector: sa.engine.Inspector, + table_name: str, + schema: str | None, +) -> str | None: + """Table comment via Inspector; None when unsupported or failing.""" + try: + return _clean_comment( + inspector.get_table_comment(table_name, schema=schema).get("text") + ) + except Exception: + return None + + def _introspect_query_columns_via_inspector( sa_engine: sa.Engine, inspector: sa.engine.Inspector, @@ -788,14 +815,15 @@ def _introspect_query_columns_via_inspector( fk_columns_by_table: dict[str, set[str]], joins: list[ModelJoin] | None = None, live_name_by_model: dict[str, str] | None = None, -) -> list[tuple]: +) -> list[IntrospectedColumn]: """Introspect columns from a rollup query or plain table. - Returns list of ``(column_name, DataType, is_primary_key, is_float, - db_type)`` tuples. ``db_type`` is the raw database type string and is only - populated when ``DataType`` came out opaque (``UNKNOWN``) — for mapped - types the declared ``DataType`` already carries everything, so leaving it - ``None`` keeps stored models and golden tests clean. + Returns a list of :class:`IntrospectedColumn`. ``db_type`` is the raw + database type string and is only populated when ``DataType`` came out + opaque (``UNKNOWN``) — for mapped types the declared ``DataType`` already + carries everything, so leaving it ``None`` keeps stored models and golden + tests clean. ``comment`` is the column's DB comment when the driver + surfaces one. For rollup queries, uses per-table inspector data since LIMIT 0 type inference can be unreliable across databases. @@ -819,8 +847,14 @@ def _introspect_query_columns_via_inspector( is_float = _sa_type_is_float(col_type) if data_type.is_opaque: db_type = _raw_db_type_str(col_type) - is_pk = col_name in pk_columns - results.append((col_name, data_type, is_pk, is_float, db_type)) + results.append(IntrospectedColumn( + name=col_name, + type=data_type, + primary_key=col_name in pk_columns, + is_float=is_float, + db_type=db_type, + comment=_clean_comment(col.get("comment")), + )) # Build list of (ref_table, dotted_path) from joins — supports diamond joins # where the same table appears via multiple paths @@ -866,8 +900,14 @@ def _introspect_query_columns_via_inspector( is_float = _sa_type_is_float(col_type) if data_type.is_opaque: ref_db_type = _raw_db_type_str(col_type) - is_pk = col["name"] in ref_pk_cols - results.append((alias, data_type, is_pk, is_float, ref_db_type)) + results.append(IntrospectedColumn( + name=alias, + type=data_type, + primary_key=col["name"] in ref_pk_cols, + is_float=is_float, + db_type=ref_db_type, + comment=_clean_comment(col.get("comment")), + )) return results @@ -879,7 +919,7 @@ def _introspect_query_columns_via_inspector( def _columns_to_model( name: str, - columns: list[tuple], + columns: list[IntrospectedColumn], data_source: str, sql_table: str | None = None, joins: list[ModelJoin] | None = None, @@ -887,14 +927,15 @@ def _columns_to_model( source_kind: ObjectKind | None = None, hidden: bool = False, meta: dict[str, Any] | None = None, + description: str | None = None, ) -> SlayerModel: - """Generate a SlayerModel from introspected ``(column_name, DataType, - is_pk, is_float, db_type)`` tuples. + """Generate a SlayerModel from :class:`IntrospectedColumn` entries. In v2 every Column is potentially both a dimension and a measure — what it's used as is decided per query. This function emits one Column per non-joined - column, with format inferred from the column's data type. ``db_type`` is - carried through verbatim (set only for opaque ``UNKNOWN`` columns). + column, with format inferred from the column's data type. ``db_type`` and + ``comment`` are carried through verbatim; ``description`` is the table + comment, when the database has one. """ cols: list[Column] = [] unique_set = unique_columns or set() @@ -902,19 +943,19 @@ def _columns_to_model( _INT_FORMAT = NumberFormat(type=NumberFormatType.INTEGER) _FLOAT_FORMAT = NumberFormat(type=NumberFormatType.FLOAT) - for col_name, data_type, is_pk, is_float, db_type in columns: + for col in columns: # Skip joined columns — they live on the target model and are # resolved via the join graph at query time. - if "." in col_name: + if "." in col.name: continue # Avoid name collision with the magic "*:count" / "_count" alias used # for COUNT(*) by renaming a literal "_count" column. - column_name = "count_col" if col_name == "_count" else col_name + column_name = "count_col" if col.name == "_count" else col.name - if is_float: + if col.is_float: fmt = _FLOAT_FORMAT - elif data_type in _NUMERIC_TYPES: + elif col.type in _NUMERIC_TYPES: fmt = _INT_FORMAT else: fmt = None @@ -922,12 +963,13 @@ def _columns_to_model( cols.append( Column( name=column_name, - sql=col_name, - type=data_type, - db_type=db_type, - primary_key=is_pk, - unique=(col_name in unique_set), + sql=col.name, + type=col.type, + db_type=col.db_type, + primary_key=col.primary_key, + unique=(col.name in unique_set), format=fmt, + description=col.comment, ) ) @@ -940,6 +982,7 @@ def _columns_to_model( source_kind=source_kind, hidden=hidden, meta=meta, + description=description, ) @@ -947,16 +990,16 @@ def _sqlite_probe_integer_columns( *, sa_engine: sa.Engine, sql_table: str, - columns: list[tuple], -) -> list[tuple]: + columns: list[IntrospectedColumn], +) -> list[IntrospectedColumn]: """DEV-1538: per-column SQLite affinity probe. - Walks the tuples ``(col_name, DataType, is_pk, is_float, db_type)`` produced by + Walks the :class:`IntrospectedColumn` entries produced by :func:`_introspect_query_columns_via_inspector` and, for every base column (alias without ``.``) that the SA inspector reported as :class:`DataType.INT`, runs :func:`slayer.sql.sqlite_introspect.probe_sqlite_integer_column` against - the actual storage classes. Mutates the tuple to the widened + the actual storage classes. Rewrites the entry to the widened :class:`DataType` whenever the probe disagrees with the declared affinity. @@ -975,17 +1018,17 @@ def _sqlite_probe_integer_columns( from slayer.sql.sqlite_introspect import probe_sqlite_integer_column schema, table = _parse_qualified_sql_table(sql_table) - out: list[tuple] = [] + out: list[IntrospectedColumn] = [] with sa_engine.connect() as conn: - for col_name, data_type, is_pk, is_float, db_type in columns: - if data_type is not DataType.INT or "." in col_name: - out.append((col_name, data_type, is_pk, is_float, db_type)) + for col in columns: + if col.type is not DataType.INT or "." in col.name: + out.append(col) continue try: verdict = probe_sqlite_integer_column( conn=conn, table=table, - column=col_name, + column=col.name, schema=schema, ) except Exception as exc: @@ -996,15 +1039,17 @@ def _sqlite_probe_integer_columns( logger.warning( "probe call raised for %s.%s; keeping declared INT: %s", sql_table, - col_name, + col.name, exc, ) verdict = None if verdict is None or verdict is DataType.INT: - out.append((col_name, data_type, is_pk, is_float, db_type)) + out.append(col) continue - new_is_float = verdict is DataType.DOUBLE - out.append((col_name, verdict, is_pk, new_is_float, db_type)) + out.append(col.model_copy(update={ + "type": verdict, + "is_float": verdict is DataType.DOUBLE, + })) return out @@ -1065,6 +1110,7 @@ def introspect_table_to_model( sql_table=sql_table, unique_columns=unique_columns, source_kind=source_kind, + description=_safe_get_table_comment(inspector, table_name, schema), ) @@ -1122,6 +1168,9 @@ class IngestionScanReport(BaseModel): # Every object discovered, modelled or not — lets the CLI tell an empty # schema apart from one whose objects were all skipped / already in sync. objects: list[IngestableObject] = Field(default_factory=list) + # BigQuery only: the dataset description (None elsewhere, or when the + # datasource already carries a description). See DEV-1809 fill-if-empty. + schema_description: str | None = None @property def hidden_internals(self) -> list[InternalTable]: @@ -1318,6 +1367,7 @@ def _build_one_model( source_kind=obj.kind, hidden=internal_tool is not None, meta=meta, + description=_safe_get_table_comment(inspector, obj.name, schema), ) @@ -1360,6 +1410,37 @@ def _collect_fk_columns( return out +def _fetch_bigquery_dataset_description( + *, + sa_engine: sa.Engine, + datasource: DatasourceConfig, + schema: str | None, +) -> str | None: + """BigQuery only: the dataset description, or None. + + Dataset resolution: explicit ``schema`` arg → the dialect's configured + default dataset → ``datasource.schema_name``. No generic Inspector API + exists for schema-level comments, so other dialects return None. + """ + try: + if getattr(sa_engine.dialect, "name", None) != "bigquery": + return None + dataset = ( + schema + or getattr(sa_engine.dialect, "dataset_id", None) + or datasource.schema_name + ) + if not dataset: + return None + with sa_engine.connect() as conn: + # The same private client handle sqlalchemy-bigquery itself uses + # for table metadata; there is no documented accessor. + client = conn.connection._client + return _clean_comment(client.get_dataset(dataset).description) + except Exception: + return None + + def ingest_datasource_report( datasource: DatasourceConfig, include_tables: list[str] | None = None, @@ -1459,11 +1540,20 @@ def ingest_datasource_report( ) ) + # Fetch while the engine is alive; skipped entirely when the + # datasource already carries a description (fill-if-empty). + schema_description = None + if not datasource.description: + schema_description = _fetch_bigquery_dataset_description( + sa_engine=sa_engine, datasource=datasource, schema=schema, + ) + return IngestionScanReport( models=models, skipped=skipped, objects=objects, internal_tables=internal_tables, + schema_description=schema_description, ) finally: # In a ``finally`` because discovery and the FK passes can raise a @@ -1665,6 +1755,46 @@ class AdditiveMergeResult(BaseModel): #: A metadata-only fill (join cardinality / column unique) that still #: has to be saved even when no column or join was added. metadata_changed: bool = False + #: DEV-1809: existing columns whose empty description was filled from a + #: DB comment, and whether the model-level description was filled. + described_columns: list[str] = Field(default_factory=list) + model_described: bool = False + + +def _merge_one_column( + *, + persisted_col: Column, + fresh_col: Column | None, + model_name: str, + sqlite_widen_enabled: bool, +) -> tuple[Column, bool, bool, bool]: + """Merge one persisted column against its fresh counterpart. + + Returns ``(merged, did_widen, unique_filled, described)`` — the probe + widening plus the two additive gap-fills (`unique` on, empty description + filled from the DB comment). + """ + merged_col, did_widen = _merge_persisted_column_with_probe( + persisted_col=persisted_col, + fresh_col=fresh_col, + model_name=model_name, + sqlite_widen_enabled=sqlite_widen_enabled, + ) + unique_filled = False + described = False + if fresh_col is not None: + # Set `unique` additively — never downgrade a user-set flag. + if fresh_col.unique and not merged_col.unique: + merged_col = merged_col.model_copy(update={"unique": True}) + unique_filled = True + # Fill an EMPTY description from the DB comment; existing + # descriptions are never overwritten. + if fresh_col.description and not merged_col.description: + merged_col = merged_col.model_copy( + update={"description": fresh_col.description} + ) + described = True + return merged_col, did_widen, unique_filled, described def _additive_merge_existing( @@ -1697,30 +1827,26 @@ def _additive_merge_existing( fresh_by_name: dict[str, Column] = {c.name: c for c in fresh.columns} widened_column_names: list[str] = [] + described_column_names: list[str] = [] merged_columns: list[Column] = [] metadata_changed = False for persisted_col in persisted.columns: - merged_col, did_widen = _merge_persisted_column_with_probe( + merged_col, did_widen, unique_filled, described = _merge_one_column( persisted_col=persisted_col, fresh_col=fresh_by_name.get(persisted_col.name), model_name=persisted.name, sqlite_widen_enabled=sqlite_widen_enabled, ) - # Set `unique` additively — never downgrade a user-set flag. - fresh_col = fresh_by_name.get(persisted_col.name) - if fresh_col is not None and fresh_col.unique and not merged_col.unique: - merged_col = merged_col.model_copy(update={"unique": True}) - metadata_changed = True + metadata_changed = metadata_changed or unique_filled + if described: + described_column_names.append(persisted_col.name) merged_columns.append(merged_col) if did_widen: widened_column_names.append(persisted_col.name) - new_column_names: list[str] = [] - for fresh_col in fresh.columns: - if fresh_col.name in existing_by_name: - continue - merged_columns.append(fresh_col) - new_column_names.append(fresh_col.name) + new_cols = [c for c in fresh.columns if c.name not in existing_by_name] + merged_columns.extend(new_cols) + new_column_names = [c.name for c in new_cols] new_joins, new_join_targets, joins_metadata_changed = _merge_joins_strict( persisted, fresh @@ -1734,18 +1860,24 @@ def _additive_merge_existing( and fresh.source_kind != persisted.source_kind ) + model_described = bool(fresh.description) and not persisted.description + if not ( new_column_names or new_join_targets or widened_column_names or kind_changed or metadata_changed + or described_column_names + or model_described ): return AdditiveMergeResult(merged=persisted) update: dict[str, Any] = {"columns": merged_columns, "joins": new_joins} if kind_changed: update["source_kind"] = fresh.source_kind + if model_described: + update["description"] = fresh.description return AdditiveMergeResult( merged=persisted.model_copy(update=update), @@ -1754,6 +1886,8 @@ def _additive_merge_existing( widened_columns=widened_column_names, kind_changed=kind_changed, metadata_changed=metadata_changed, + described_columns=described_column_names, + model_described=model_described, ) @@ -1780,6 +1914,8 @@ async def _process_one_table( new_columns=[c.name for c in fresh.columns], new_joins=[j.target_model for j in fresh.joins], source_kind=fresh.source_kind, + described_columns=[c.name for c in fresh.columns if c.description], + model_described=bool(fresh.description), ) if persisted.sql or persisted.source_queries: # User-authored sql / query-backed model with the matching name — @@ -1799,6 +1935,8 @@ async def _process_one_table( or outcome.widened_columns or outcome.kind_changed or outcome.metadata_changed + or outcome.described_columns + or outcome.model_described ): await storage.save_model(outcome.merged) kind_change = None @@ -1814,6 +1952,8 @@ async def _process_one_table( widened_columns=outcome.widened_columns, source_kind=outcome.merged.source_kind, kind_change=kind_change, + described_columns=outcome.described_columns, + model_described=outcome.model_described, ) @@ -1957,6 +2097,28 @@ async def ingest_datasource_idempotent( ) ) + # DEV-1809 fill-if-empty for the datasource description (BigQuery dataset + # description). The check runs against the freshly-loaded STORED config, + # not the caller's object — a stale caller copy must never clobber a + # description persisted since it was loaded. Best-effort: a save failure + # must not abort the pass (model additions above are already persisted); + # ``model_name=""`` is the established datasource-level tag. + datasource_described = False + if scan.schema_description and not datasource.description: + try: + stored = await storage.get_datasource(datasource.name) or datasource + if not stored.description: + await storage.save_datasource( + stored.model_copy(update={"description": scan.schema_description}) + ) + datasource_described = True + except Exception as exc: # noqa: BLE001 — best-effort isolation + errors.append(IngestionError( + model_name="", + data_source=datasource.name, + error=f"datasource description save: {exc}", + )) + scoped_models = await _scoped_models_for_validation( storage=storage, datasource=datasource, @@ -2010,6 +2172,7 @@ async def ingest_datasource_idempotent( skipped=scan.skipped, objects=scan.objects, hidden_internals=hidden_internals, + datasource_described=datasource_described, ) @@ -2092,16 +2255,20 @@ def _print_ingest_addition( out = file if file is not None else sys.stdout # Label non-table objects — a view-backed model has no PK and no joins. label = _KIND_LABELS.get(getattr(addition, "source_kind", None) or "", "") + described = getattr(addition, "described_columns", []) or [] + model_described = getattr(addition, "model_described", False) if addition.created: + suffix = f", {len(described)} described" if described else "" print( f"Created: {addition.model_name} " - f"({len(addition.new_columns)} columns){label}", + f"({len(addition.new_columns)} columns{suffix}){label}", file=out, ) return widened = getattr(addition, "widened_columns", []) or [] kind_change = getattr(addition, "kind_change", None) - if not (addition.new_columns or addition.new_joins or widened or kind_change): + if not (addition.new_columns or addition.new_joins or widened or kind_change + or described or model_described): return details = [] if addition.new_columns: @@ -2112,6 +2279,10 @@ def _print_ingest_addition( details.append(f"widened: {', '.join(widened)}") if kind_change: details.append(f"source_kind: {kind_change}") + if described: + details.append(f"+descriptions: {', '.join(described)}") + if model_described: + details.append("+model description") print(f"Updated: {addition.model_name} ({'; '.join(details)})", file=out) @@ -2169,6 +2340,8 @@ def _print_ingest_drift_and_errors( ``_unhide_hint``); it comes from the caller since neither result carries it. """ out = file if file is not None else sys.stdout + if getattr(result, "datasource_described", False): + print("Datasource description imported.", file=out) _print_report_section( entries=getattr(result, "to_delete", None) or [], header="\nPending drift (run `slayer validate-models` to inspect):", diff --git a/slayer/engine/introspect_utils.py b/slayer/engine/introspect_utils.py index fe30306a..7c3572aa 100644 --- a/slayer/engine/introspect_utils.py +++ b/slayer/engine/introspect_utils.py @@ -55,6 +55,93 @@ } +def _clean_comment(value: Optional[str]) -> Optional[str]: + """Normalize a DB comment: strip whitespace, empty → None.""" + if value is None: + return None + cleaned = value.strip() + return cleaned or None + + +# Per-dialect column-comment queries for the fallback path. INFORMATION_SCHEMA +# has no standard comment column, so each dialect needs its own source +# (DEV-1809). SQL Server / BigQuery are omitted: their Inspector paths already +# surface comments, and their fallback equivalents are disproportionately +# complex (sys.extended_properties / region-qualified INFORMATION_SCHEMA). +_COMMENT_FALLBACK_SQL = { + "mysql": ( + "SELECT column_name, column_comment FROM information_schema.columns " + "WHERE table_name = :table_name{schema_clause}", + " AND table_schema = :schema", + " AND table_schema = DATABASE()", + ), + "snowflake": ( + "SELECT column_name, comment FROM information_schema.columns " + "WHERE table_name = :table_name{schema_clause}", + " AND table_schema = :schema", + " AND table_schema = CURRENT_SCHEMA()", + ), + "clickhouse": ( + "SELECT name, comment FROM system.columns " + "WHERE table = :table_name AND database = {schema_clause}", + ":schema", + "currentDatabase()", + ), + "duckdb": ( + "SELECT column_name, comment FROM duckdb_columns() " + "WHERE table_name = :table_name{schema_clause}", + " AND schema_name = :schema", + " AND schema_name = current_schema()", + ), + "postgresql": ( + "SELECT a.attname, col_description(c.oid, a.attnum) " + "FROM pg_catalog.pg_class c " + "JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace " + "JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid " + "WHERE c.relname = :table_name AND a.attnum > 0 " + "AND NOT a.attisdropped{schema_clause}", + " AND n.nspname = :schema", + " AND pg_catalog.pg_table_is_visible(c.oid)", + ), +} +_COMMENT_FALLBACK_SQL["mariadb"] = _COMMENT_FALLBACK_SQL["mysql"] + + +def _get_column_comments_fallback( + sa_engine: sa.Engine, + table_name: str, + schema: Optional[str], +) -> Dict[str, str]: + """Column comments for the INFORMATION_SCHEMA fallback path. + + Without an explicit ``schema`` the query is scoped to the connection's + default/current schema, so a same-named table in another schema can't + cross-assign its comments. Best-effort: unknown dialects and any query + failure return ``{}``. + """ + try: + dialect_name = getattr(sa_engine.dialect, "name", None) + entry = _COMMENT_FALLBACK_SQL.get(dialect_name) + if entry is None: + return {} + template, schema_clause, default_clause = entry + clause = schema_clause if schema else default_clause + sql = template.format(schema_clause=clause) + params = {"table_name": table_name} + if schema: + params["schema"] = schema + with sa_engine.connect() as conn: + rows = conn.execute(sa.text(sql), params).fetchall() + out: Dict[str, str] = {} + for name, comment in rows: + cleaned = _clean_comment(comment) + if cleaned: + out[name] = cleaned + return out + except Exception: + return {} + + def _parse_info_schema_is_float(data_type_str: str) -> bool: """Determine if a NUMERIC/DECIMAL info-schema type string is float-like. @@ -136,6 +223,11 @@ def _get_columns_fallback( elif sa_type is None and ("CHAR" in base_type or "TEXT" in base_type): sa_type = DataType.TEXT result.append({"name": col_name, "type": sa_type or DataType.TEXT, "is_float": is_float}) + comments = _get_column_comments_fallback( + sa_engine=sa_engine, table_name=table_name, schema=schema, + ) + for col in result: + col["comment"] = comments.get(col["name"]) return result diff --git a/slayer/engine/schema_drift.py b/slayer/engine/schema_drift.py index e5d31fed..42bd69d4 100644 --- a/slayer/engine/schema_drift.py +++ b/slayer/engine/schema_drift.py @@ -117,6 +117,10 @@ class ModelAddition(BaseModel): # Human-readable transition (e.g. "view → table") when a re-ingest found # the live object changed kind; None when nothing changed. kind_change: str | None = None + # DEV-1809: columns whose empty description was filled from a DB comment + # (on created models: all columns that arrived with a description). + described_columns: list[str] = Field(default_factory=list) + model_described: bool = False class IngestionError(BaseModel): @@ -145,6 +149,9 @@ class IdempotentIngestResult(BaseModel): # post-merge state, not the scan's verdict (the merge preserves a persisted # ``hidden``). ``Any`` avoids a circular import; entries are ``InternalTable``. hidden_internals: list[Any] = Field(default_factory=list) + # DEV-1809: True when the datasource description was filled from the + # BigQuery dataset description during this pass. + datasource_described: bool = False class AppliedEntry(BaseModel): diff --git a/slayer/mcp/server.py b/slayer/mcp/server.py index fe94ae48..0e423d6c 100644 --- a/slayer/mcp/server.py +++ b/slayer/mcp/server.py @@ -194,28 +194,59 @@ def _empty_ingest_message(*, schema_name: str, ds: DatasourceConfig) -> str: ) +def _addition_has_changes(a: Any) -> bool: + """True when a non-created addition carries any change worth rendering.""" + return bool( + a.new_columns + or a.new_joins + or getattr(a, "widened_columns", None) + or getattr(a, "kind_change", None) + or getattr(a, "described_columns", None) + or getattr(a, "model_described", False) + ) + + def _render_new_models_section(new_models: list[Any]) -> list[str]: if not new_models: return [] lines = [f"Created {len(new_models)} new model(s):"] for a in new_models: + described = getattr(a, "described_columns", []) or [] + suffix = f", {len(described)} described" if described else "" lines.append( - f"- {a.model_name} ({len(a.new_columns)} columns, {len(a.new_joins)} joins)" + f"- {a.model_name} ({len(a.new_columns)} columns, " + f"{len(a.new_joins)} joins{suffix})" ) return lines +def _addition_update_details(a: Any) -> list[str]: + """Detail fragments for one non-created addition, in CLI-renderer order.""" + details = [] + if a.new_columns: + details.append(f"+columns: {', '.join(a.new_columns)}") + if a.new_joins: + details.append(f"+joins: {', '.join(a.new_joins)}") + widened = getattr(a, "widened_columns", []) or [] + if widened: + details.append(f"widened: {', '.join(widened)}") + kind_change = getattr(a, "kind_change", None) + if kind_change: + details.append(f"source_kind: {kind_change}") + described = getattr(a, "described_columns", []) or [] + if described: + details.append(f"+descriptions: {', '.join(described)}") + if getattr(a, "model_described", False): + details.append("+model description") + return details + + def _render_updated_section(updated: list[Any]) -> list[str]: if not updated: return [] lines = [f"Updated {len(updated)} existing model(s):"] for a in updated: - details = [] - if a.new_columns: - details.append(f"+columns: {', '.join(a.new_columns)}") - if a.new_joins: - details.append(f"+joins: {', '.join(a.new_joins)}") - lines.append(f"- {a.model_name} ({'; '.join(details)})") + lines.append(f"- {a.model_name} ({'; '.join(_addition_update_details(a))})") return lines @@ -291,12 +322,14 @@ def _render_ingest_result( # may carry neither attribute. skipped = list(getattr(result, "skipped", None) or []) hidden_internals = list(getattr(result, "hidden_internals", None) or []) + datasource_described = bool(getattr(result, "datasource_described", False)) if ( not additions and not result.to_delete and not result.errors and not skipped and not hidden_internals + and not datasource_described ): # Two distinct cases produce an empty result: # 1. The schema actually has no tables (the agent should look @@ -315,16 +348,17 @@ def _render_ingest_result( return "Datasource already in sync — no additive changes." new_models = [a for a in additions if a.created] - updated = [a for a in additions if not a.created and (a.new_columns or a.new_joins)] + updated = [a for a in additions if not a.created and _addition_has_changes(a)] unchanged = [ - a for a in additions - if not a.created and not a.new_columns and not a.new_joins + a for a in additions if not a.created and not _addition_has_changes(a) ] lines: list[str] = [] lines.extend(_render_new_models_section(new_models)) lines.extend(_render_updated_section(updated)) lines.extend(_render_unchanged_section(unchanged)) + if datasource_described: + lines.append("Datasource description imported.") lines.extend(_render_drift_section(list(result.to_delete))) # Same order as the CLI renderer, so the two surfaces read alike. lines.extend(_render_skipped_section(skipped)) @@ -1461,7 +1495,7 @@ async def create_datasource( Example: create_datasource(name="mydb", type="postgres", host="localhost", port=5432, database="app", username="user", password="pass") """ - from slayer.engine.ingestion import ingest_datasource as _ingest + from slayer.engine.ingestion import ingest_datasource_report as _ingest data = _build_dict( name=name, @@ -1493,12 +1527,23 @@ async def create_datasource( # Auto-ingest models try: - models = _ingest(datasource=ds, schema=schema_name or None) + ingest_output = _ingest(datasource=ds, schema=schema_name or None) except Exception as e: if isinstance(e, (sa.exc.OperationalError, sa.exc.DatabaseError)): lines.append(f"Auto-ingestion failed: {_friendly_db_error(e)}") return "\n".join(lines) raise + models = ingest_output.models + + if ingest_output.schema_description and not ds.description: + try: + ds = ds.model_copy( + update={"description": ingest_output.schema_description} + ) + await storage.save_datasource(ds) + lines.append("Datasource description imported.") + except Exception as exc: # noqa: BLE001 — best-effort + lines.append(f"Could not save datasource description: {exc}") save_errors: list[str] = [] saved_models = [] diff --git a/tests/integration/test_integration.py b/tests/integration/test_integration.py index 8e4c47f7..b594bfe2 100644 --- a/tests/integration/test_integration.py +++ b/tests/integration/test_integration.py @@ -1884,6 +1884,14 @@ async def test_diamond_joins_both_paths(diamond_env): assert "warehouses__regions" in result.sql +async def test_sqlite_ingest_has_no_descriptions(diamond_env): + """SQLite has no table/column comments — ingested descriptions stay None.""" + _, storage = diamond_env + shipments = await storage.get_model("shipments") + assert shipments.description is None + assert all(c.description is None for c in shipments.columns) + + async def test_query_filter_on_joined_dimension(diamond_env): """Query-level filter on a joined dimension resolves through the model.""" engine, _ = diamond_env diff --git a/tests/integration/test_integration_bigquery.py b/tests/integration/test_integration_bigquery.py new file mode 100644 index 00000000..33795a02 --- /dev/null +++ b/tests/integration/test_integration_bigquery.py @@ -0,0 +1,237 @@ +"""Live BigQuery integration tests for comment/description ingestion. + +Runs against a real GCP project: requires ``GCP_PROJECT_ID`` plus Application +Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS`` or ambient ADC), and a +service account allowed to create/delete datasets in that project. Skips +cleanly when the env var or credentials are absent (local dev machines); +fails loudly when credentials exist but lack dataset-create rights, so a CI +permissions regression is visible rather than silently skipped. + +A uniquely-named temporary dataset (with a description) is created, populated +with commented tables, ingested, and deleted on teardown. +""" + +import os +import uuid +from pathlib import Path + +import pytest + +pytest.importorskip("sqlalchemy_bigquery") + +from google.api_core.exceptions import Forbidden # noqa: E402 +from google.auth.exceptions import DefaultCredentialsError # noqa: E402 +from google.cloud import bigquery # noqa: E402 + +from slayer.core.enums import DataType # noqa: E402 +from slayer.core.models import DatasourceConfig # noqa: E402 +from slayer.engine.ingestion import ( # noqa: E402 + ingest_datasource, + ingest_datasource_idempotent, +) +from slayer.storage.yaml_storage import YAMLStorage # noqa: E402 + +pytestmark = pytest.mark.integration + +_PROJECT = os.environ.get("GCP_PROJECT_ID") +_DATASET_DESCRIPTION = "Dataset for SLayer ingestion tests" + + +@pytest.fixture(scope="module") +def bq_dataset(): + """Temp dataset with a description + commented tables; dropped on teardown.""" + if not _PROJECT: + pytest.skip("GCP_PROJECT_ID not set — live BigQuery tests need a billing project") + try: + client = bigquery.Client(project=_PROJECT) + except DefaultCredentialsError as exc: + pytest.skip(f"No Google Application Default Credentials: {exc}") + dataset_id = f"slayer_test_{uuid.uuid4().hex[:12]}" + dataset = bigquery.Dataset(f"{_PROJECT}.{dataset_id}") + dataset.description = _DATASET_DESCRIPTION + try: + client.create_dataset(dataset) + except Forbidden as exc: + client.close() + pytest.fail( + f"Service account cannot create datasets in {_PROJECT} " + f"(grant roles/bigquery.user or dataEditor): {exc}" + ) + except Exception: + client.close() + raise + try: + orders = bigquery.Table( + f"{_PROJECT}.{dataset_id}.orders", + schema=[ + bigquery.SchemaField("id", "INTEGER", description="Order id"), + bigquery.SchemaField( + "amount", "FLOAT", description="Order amount in USD" + ), + bigquery.SchemaField("status", "STRING"), + bigquery.SchemaField( + "payload", + "RECORD", + fields=[ + bigquery.SchemaField( + "city", "STRING", description="City in the payload" + ), + ], + description="Structured payload", + ), + ], + ) + orders.description = "All orders" + client.create_table(orders) + client.create_table(bigquery.Table( + f"{_PROJECT}.{dataset_id}.plain", + schema=[bigquery.SchemaField("x", "INTEGER")], + )) + yield client, dataset_id + finally: + client.delete_dataset( + f"{_PROJECT}.{dataset_id}", delete_contents=True, not_found_ok=True + ) + client.close() + + +def _ds_config(dataset_id: str, *, name: str = "bqtest", description: str | None = None): + # Dataset-qualified URL so table names reflect bare (not "dataset.table"). + return DatasourceConfig( + name=name, + type="bigquery", + connection_string=f"bigquery://{_PROJECT}/{dataset_id}", + description=description, + ) + + +@pytest.fixture(scope="module") +def bq_models(bq_dataset): + _, dataset_id = bq_dataset + models = ingest_datasource( + datasource=_ds_config(dataset_id), schema=dataset_id + ) + return {m.name: m for m in models}, dataset_id + + +class TestBigQueryCommentIngestion: + def test_column_descriptions_imported(self, bq_models) -> None: + models, _ = bq_models + cols = {c.name: c for c in models["orders"].columns} + assert cols["id"].description == "Order id" + assert cols["amount"].description == "Order amount in USD" + assert cols["status"].description is None + + def test_table_description_imported(self, bq_models) -> None: + models, _ = bq_models + assert models["orders"].description == "All orders" + assert models["plain"].description is None + + def test_record_column_description(self, bq_models) -> None: + models, _ = bq_models + cols = {c.name: c for c in models["orders"].columns} + assert cols["payload"].description == "Structured payload" + # Flattened RECORD subfields (dotted names) never become columns. + assert not any("." in name for name in cols) + + def test_types_sanity(self, bq_models) -> None: + models, _ = bq_models + cols = {c.name: c for c in models["orders"].columns} + assert cols["id"].type is DataType.INT + assert cols["amount"].type is DataType.DOUBLE + + +class TestBigQueryDatasetDescription: + async def test_dataset_description_fills_datasource( + self, bq_dataset, tmp_path: Path + ) -> None: + _, dataset_id = bq_dataset + storage = YAMLStorage(base_dir=str(tmp_path / "storage")) + ds = _ds_config(dataset_id) + await storage.save_datasource(ds) + + result = await ingest_datasource_idempotent( + datasource=ds, storage=storage, schema=dataset_id + ) + assert result.datasource_described is True + loaded = await storage.get_datasource("bqtest") + assert loaded.description == _DATASET_DESCRIPTION + + # Re-ingest with the reloaded datasource: no-op, nothing re-reported. + result2 = await ingest_datasource_idempotent( + datasource=loaded, storage=storage, schema=dataset_id + ) + assert result2.datasource_described is False + + async def test_preset_description_untouched( + self, bq_dataset, tmp_path: Path + ) -> None: + _, dataset_id = bq_dataset + storage = YAMLStorage(base_dir=str(tmp_path / "storage")) + ds = _ds_config(dataset_id, name="bqtest2", description="user text") + await storage.save_datasource(ds) + + result = await ingest_datasource_idempotent( + datasource=ds, storage=storage, schema=dataset_id + ) + assert result.datasource_described is False + loaded = await storage.get_datasource("bqtest2") + assert loaded.description == "user text" + + +class TestBigQueryReingest: + async def test_fill_if_empty_and_preservation( + self, bq_dataset, tmp_path: Path + ) -> None: + _, dataset_id = bq_dataset + storage = YAMLStorage(base_dir=str(tmp_path / "storage")) + ds = _ds_config(dataset_id) + await storage.save_datasource(ds) + await ingest_datasource_idempotent( + datasource=ds, storage=storage, schema=dataset_id + ) + + loaded = await storage.get_model("orders", data_source="bqtest") + for c in loaded.columns: + if c.name == "amount": + c.description = None + if c.name == "status": + c.description = "hand-authored" + await storage.save_model(loaded) + + result = await ingest_datasource_idempotent( + datasource=ds, storage=storage, schema=dataset_id + ) + addition = next(a for a in result.additions if a.model_name == "orders") + assert addition.described_columns == ["amount"] + loaded2 = await storage.get_model("orders", data_source="bqtest") + cols = {c.name: c for c in loaded2.columns} + assert cols["amount"].description == "Order amount in USD" + assert cols["status"].description == "hand-authored" + + async def test_new_commented_column_arrives_with_description( + self, bq_dataset, tmp_path: Path + ) -> None: + client, dataset_id = bq_dataset + storage = YAMLStorage(base_dir=str(tmp_path / "storage")) + ds = _ds_config(dataset_id) + await storage.save_datasource(ds) + await ingest_datasource_idempotent( + datasource=ds, storage=storage, schema=dataset_id + ) + + table = client.get_table(f"{_PROJECT}.{dataset_id}.orders") + table.schema = list(table.schema) + [ + bigquery.SchemaField("discount", "FLOAT", description="Discount applied"), + ] + client.update_table(table, ["schema"]) + + result = await ingest_datasource_idempotent( + datasource=ds, storage=storage, schema=dataset_id + ) + addition = next(a for a in result.additions if a.model_name == "orders") + assert "discount" in addition.new_columns + assert "discount" not in addition.described_columns + loaded = await storage.get_model("orders", data_source="bqtest") + discount = next(c for c in loaded.columns if c.name == "discount") + assert discount.description == "Discount applied" diff --git a/tests/integration/test_integration_clickhouse.py b/tests/integration/test_integration_clickhouse.py index 479ea1f2..fe0dc21d 100644 --- a/tests/integration/test_integration_clickhouse.py +++ b/tests/integration/test_integration_clickhouse.py @@ -993,10 +993,10 @@ def clickhouse_ingest_for_types_env(clickhouse_container): id Int32, customer_id Int32, quantity Int32, - amount Float64, + amount Float64 COMMENT 'Order amount', status String, created_at DateTime - ) ENGINE = MergeTree() ORDER BY id + ) ENGINE = MergeTree() ORDER BY id COMMENT 'All orders' """)) conn.execute(sa.text(""" INSERT INTO orders VALUES @@ -1030,6 +1030,17 @@ def test_clickhouse_datetime_typed_correctly(clickhouse_ingest_for_types_env) -> ) +@pytest.mark.integration +def test_clickhouse_comments_imported(clickhouse_ingest_for_types_env) -> None: + ds = clickhouse_ingest_for_types_env + models = ingest_datasource(datasource=ds, schema=None) + orders = next(m for m in models if m.name == "orders") + assert orders.description == "All orders" + by_name = {c.name: c for c in orders.columns} + assert by_name["amount"].description == "Order amount" + assert by_name["status"].description is None + + # --------------------------------------------------------------------------- # DEV-1727 — dialect-aware Mode-A {var} escaping (ClickHouse is a Tier-1 # backslash dialect with C-style string literals: the naive '' quote-doubling diff --git a/tests/integration/test_integration_duckdb.py b/tests/integration/test_integration_duckdb.py index 110255c2..1c049a86 100644 --- a/tests/integration/test_integration_duckdb.py +++ b/tests/integration/test_integration_duckdb.py @@ -366,6 +366,8 @@ def duckdb_ingest_env(tmp_path_factory): customer_id INTEGER REFERENCES customers(id) ) """) + conn.execute("COMMENT ON TABLE orders IS 'All orders'") + conn.execute("COMMENT ON COLUMN orders.amount IS 'Order amount'") conn.executemany("INSERT INTO regions VALUES (?, ?)", [(1, "US"), (2, "EU")]) conn.executemany( "INSERT INTO customers VALUES (?, ?, ?)", @@ -802,3 +804,19 @@ async def test_integration_duckdb_cross_model_derived_columnsql( assert response.row_count == 2 assert float(response.data[0]["a_tbl.ratio_using_derived"]) == pytest.approx(2.0) assert float(response.data[1]["a_tbl.ratio_using_derived"]) == pytest.approx(2.0) + + +@pytest.mark.integration +class TestDuckDBIngestComments: + def test_table_and_column_comments_imported(self, duckdb_ingest_env) -> None: + models, _ = duckdb_ingest_env + orders = next(m for m in models if m.name == "orders") + assert orders.description == "All orders" + by_name = {c.name: c for c in orders.columns} + assert by_name["amount"].description == "Order amount" + assert by_name["id"].description is None + + def test_uncommented_table_has_no_description(self, duckdb_ingest_env) -> None: + models, _ = duckdb_ingest_env + regions = next(m for m in models if m.name == "regions") + assert regions.description is None diff --git a/tests/integration/test_integration_mysql.py b/tests/integration/test_integration_mysql.py index 97521472..ee01791c 100644 --- a/tests/integration/test_integration_mysql.py +++ b/tests/integration/test_integration_mysql.py @@ -607,10 +607,10 @@ def mysql_ingest_env(mysql_container): cur.execute(""" CREATE TABLE orders ( id INTEGER PRIMARY KEY, - amount DECIMAL(10,2) NOT NULL, + amount DECIMAL(10,2) NOT NULL COMMENT 'Order amount', customer_id INTEGER, FOREIGN KEY (customer_id) REFERENCES customers(id) - ) ENGINE=InnoDB + ) ENGINE=InnoDB COMMENT='All orders' """) cur.executemany("INSERT INTO regions VALUES (%s, %s)", [(1, "US"), (2, "EU")]) cur.executemany( @@ -1330,3 +1330,19 @@ async def test_breakout_attempt_matches_nothing( variables={"v": "x\\' OR '1'='1"}, )) assert int(resp.data[0]["esc._count"]) == 0 + + +@pytest.mark.integration +class TestMySQLIngestComments: + def test_table_and_column_comments_imported(self, mysql_ingest_env) -> None: + models, _, _ = mysql_ingest_env + orders = next(m for m in models if m.name == "orders") + assert orders.description == "All orders" + by_name = {c.name: c for c in orders.columns} + assert by_name["amount"].description == "Order amount" + assert by_name["id"].description is None + + def test_uncommented_table_has_no_description(self, mysql_ingest_env) -> None: + models, _, _ = mysql_ingest_env + regions = next(m for m in models if m.name == "regions") + assert regions.description is None diff --git a/tests/integration/test_integration_postgres.py b/tests/integration/test_integration_postgres.py index 5538b459..51d9368a 100644 --- a/tests/integration/test_integration_postgres.py +++ b/tests/integration/test_integration_postgres.py @@ -563,6 +563,8 @@ def pg_ingest_env(postgresql_proc): customer_id INTEGER REFERENCES customers(id) ) """) + cur.execute("COMMENT ON TABLE orders IS 'All orders'") + cur.execute("COMMENT ON COLUMN orders.amount IS 'Order amount'") cur.executemany("INSERT INTO regions VALUES (%s, %s)", [(1, "US"), (2, "EU")]) cur.executemany( "INSERT INTO customers VALUES (%s, %s, %s)", @@ -1543,3 +1545,19 @@ async def test_derived_column_referencing_reserved_join_executes( # grant amounts 100 (usage 1) and 50 (usage 3) → bumped 101 and 51 by_bumped = {float(r["usage.bumped"]): r["usage._count"] for r in result.data} assert by_bumped == {101.0: 2, 51.0: 1} + + +@pytest.mark.integration +class TestPostgresIngestComments: + def test_table_and_column_comments_imported(self, pg_ingest_env) -> None: + models, _, _ = pg_ingest_env + orders = next(m for m in models if m.name == "orders") + assert orders.description == "All orders" + by_name = {c.name: c for c in orders.columns} + assert by_name["amount"].description == "Order amount" + assert by_name["id"].description is None + + def test_uncommented_table_has_no_description(self, pg_ingest_env) -> None: + models, _, _ = pg_ingest_env + regions = next(m for m in models if m.name == "regions") + assert regions.description is None diff --git a/tests/integration/test_integration_snowflake.py b/tests/integration/test_integration_snowflake.py index 8ba5fa91..99bbce53 100644 --- a/tests/integration/test_integration_snowflake.py +++ b/tests/integration/test_integration_snowflake.py @@ -142,6 +142,8 @@ def sf_transient_schema(): created_at TIMESTAMP_NTZ NOT NULL ) """) + cur.execute("COMMENT ON TABLE orders IS 'All orders'") + cur.execute("COMMENT ON COLUMN orders.quantity IS 'Units ordered'") cur.executemany( "INSERT INTO regions VALUES (%s, %s)", [(1, "US"), (2, "EU"), (3, "APAC")], @@ -622,6 +624,18 @@ def _first_pair(j) -> tuple[str, str]: assert customers_joins["regions"] == ("region_id", "id") +def test_auto_ingest_imports_comments(sf_datasource, sf_transient_schema) -> None: + """Snowflake table/column comments must land on model/column descriptions.""" + models = ingest_datasource(datasource=sf_datasource, schema=sf_transient_schema) + by_name = {m.name.lower(): m for m in models} + orders = by_name["orders"] + assert orders.description == "All orders" + cols = {c.name.lower(): c for c in orders.columns} + assert cols["quantity"].description == "Units ordered" + assert cols["status"].description is None + assert by_name["regions"].description is None + + # --------------------------------------------------------------------------- # EXPLAIN # --------------------------------------------------------------------------- diff --git a/tests/integration/test_integration_sqlserver.py b/tests/integration/test_integration_sqlserver.py index 7d9b59d1..064232b9 100644 --- a/tests/integration/test_integration_sqlserver.py +++ b/tests/integration/test_integration_sqlserver.py @@ -670,6 +670,14 @@ def sqlserver_ingest_env(sqlserver_container): "INSERT INTO orders (id, amount, customer_id) VALUES " "(1, 100, 1), (2, 200, 1), (3, 50, 2), (4, 150, 3)" )) + conn.execute(sa.text( + "EXEC sp_addextendedproperty 'MS_Description', 'All orders', " + "'SCHEMA', 'dbo', 'TABLE', 'orders'" + )) + conn.execute(sa.text( + "EXEC sp_addextendedproperty 'MS_Description', 'Order amount', " + "'SCHEMA', 'dbo', 'TABLE', 'orders', 'COLUMN', 'amount'" + )) finally: engine.dispose() @@ -680,6 +688,22 @@ def sqlserver_ingest_env(sqlserver_container): _drop_module_db(sqlserver_container, db_name) +@pytest.mark.integration +class TestSQLServerIngestComments: + def test_table_and_column_comments_imported(self, sqlserver_ingest_env) -> None: + models, _, _ = sqlserver_ingest_env + orders = next(m for m in models if m.name == "orders") + assert orders.description == "All orders" + by_name = {c.name: c for c in orders.columns} + assert by_name["amount"].description == "Order amount" + assert by_name["id"].description is None + + def test_uncommented_table_has_no_description(self, sqlserver_ingest_env) -> None: + models, _, _ = sqlserver_ingest_env + regions = next(m for m in models if m.name == "regions") + assert regions.description is None + + @pytest.mark.integration class TestRollupIngestionSQLServer: def test_orders_has_own_columns_only(self, sqlserver_ingest_env) -> None: diff --git a/tests/test_dbt_converter.py b/tests/test_dbt_converter.py index 1ffd9d76..52d47ebc 100644 --- a/tests/test_dbt_converter.py +++ b/tests/test_dbt_converter.py @@ -860,6 +860,52 @@ def test_name_collision_prefers_semantic_model(self) -> None: assert result.models[0].name == "orders" +class TestHiddenModelDescriptionPrecedence: + """Curated dbt descriptions win over DB comments; DB comments fill gaps.""" + + def _convert_with_db_comments(self, project) -> SlayerModel: + fake_model = _sample_slayer_model(name="raw_events") + fake_model.description = "db table comment" + for c in fake_model.columns: + if c.name == "event_id": + c.description = "db comment on event_id" + fake_model.columns.append( + Column(name="db_only", sql="db_only", type=DataType.TEXT, + description="db comment only") + ) + engine = MagicMock(spec=sa.Engine) + with patch.object(sa, "inspect", return_value=MagicMock()), \ + patch.object(converter_module, "introspect_table_to_model", return_value=fake_model): + result = DbtToSlayerConverter( + project=project, + data_source="test_db", + include_hidden_models=True, + sa_engine=engine, + ).convert() + return next(m for m in result.models if m.hidden) + + def test_curated_wins_and_db_comments_fill_gaps(self) -> None: + raw = self._convert_with_db_comments(_project_with_orphan()) + assert raw.description == "Raw event log" + cols = {c.name: c for c in raw.columns} + assert cols["event_id"].description == "Unique event identifier" + assert cols["db_only"].description == "db comment only" + + def test_db_comments_survive_without_curated_metadata(self) -> None: + project = DbtProject( + semantic_models=[], + metrics=[], + regular_models=[ + DbtRegularModel(name="raw_events", schema_name="staging", + alias="raw_events"), + ], + ) + raw = self._convert_with_db_comments(project) + assert raw.description == "db table comment" + cols = {c.name: c for c in raw.columns} + assert cols["event_id"].description == "db comment on event_id" + + class TestForeignEntityJoinsAllPrimaries: """Foreign entities must produce joins to ALL matching primary models, not just the first.""" diff --git a/tests/test_ingest_internal_tables.py b/tests/test_ingest_internal_tables.py index 3ff8147a..7bf72b25 100644 --- a/tests/test_ingest_internal_tables.py +++ b/tests/test_ingest_internal_tables.py @@ -672,11 +672,11 @@ async def test_user_edits_to_an_internal_model_survive_re_ingest( class TestColumnsToModelKwargs: def test_meta_is_propagated_verbatim(self) -> None: """`_columns_to_model` passes `meta` through untouched, since the caller merges the breadcrumb.""" - from slayer.engine.ingestion import _columns_to_model + from slayer.engine.ingestion import IntrospectedColumn, _columns_to_model model = _columns_to_model( name="t", - columns=[("id", DataType.INT, True, False, None)], + columns=[IntrospectedColumn(name="id", type=DataType.INT, primary_key=True)], data_source="ds", sql_table="t", hidden=True, @@ -687,11 +687,11 @@ def test_meta_is_propagated_verbatim(self) -> None: def test_defaults_leave_the_model_untouched(self) -> None: """The dbt hidden-import path calls this without the new kwargs.""" - from slayer.engine.ingestion import _columns_to_model + from slayer.engine.ingestion import IntrospectedColumn, _columns_to_model model = _columns_to_model( name="t", - columns=[("id", DataType.INT, True, False, None)], + columns=[IntrospectedColumn(name="id", type=DataType.INT, primary_key=True)], data_source="ds", sql_table="t", ) diff --git a/tests/test_ingestion.py b/tests/test_ingestion.py index 4c151847..d1559ff2 100644 --- a/tests/test_ingestion.py +++ b/tests/test_ingestion.py @@ -846,7 +846,10 @@ def test_dotted_alias_not_passed_to_probe(self) -> None: the target model's own probe pass, not the source table's.""" from unittest.mock import patch - from slayer.engine.ingestion import _sqlite_probe_integer_columns + from slayer.engine.ingestion import ( + IntrospectedColumn, + _sqlite_probe_integer_columns, + ) # Build a dummy SA engine just so the helper's dialect check passes. sa_engine = sa.create_engine("sqlite:///:memory:") @@ -866,8 +869,8 @@ def _capture(*, conn, table, column, schema=None): ): # Mixed bag: one base column (no '.') and one dotted alias. columns = [ - ("qty", DataType.INT, False, False, None), - ("customers.region_id", DataType.INT, False, False, None), + IntrospectedColumn(name="qty", type=DataType.INT), + IntrospectedColumn(name="customers.region_id", type=DataType.INT), ] _sqlite_probe_integer_columns( sa_engine=sa_engine, diff --git a/tests/test_ingestion_comments.py b/tests/test_ingestion_comments.py new file mode 100644 index 00000000..ba18f9c3 --- /dev/null +++ b/tests/test_ingestion_comments.py @@ -0,0 +1,918 @@ +"""Tests for importing DB column/table/dataset comments during ingestion. + +Column comments land on ``Column.description``, table comments on +``SlayerModel.description``, and (BigQuery only) the dataset description on +``DatasourceConfig.description`` — always fill-if-empty, never overwriting. +""" + +from __future__ import annotations + +import argparse +import io +import sqlite3 +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import duckdb +import pytest +import sqlalchemy as sa + +import slayer.engine.ingestion as ingestion_mod +from slayer.async_utils import run_sync +from slayer.cli import _run_datasources_create +from slayer.core.enums import DataType +from slayer.core.models import Column, DatasourceConfig, SlayerModel +from slayer.engine.ingestion import ( + IngestionScanReport, + IntrospectedColumn, + _additive_merge_existing, + _fetch_bigquery_dataset_description, + _print_ingest_addition, + _safe_get_table_comment, + _sqlite_probe_integer_columns, + ingest_datasource, + ingest_datasource_idempotent, + ingest_datasource_report, + introspect_table_to_model, +) +from slayer.engine.introspect_utils import ( + _clean_comment, + _get_column_comments_fallback, + _get_columns_fallback, +) +from slayer.engine.schema_drift import IdempotentIngestResult, ModelAddition +from slayer.mcp.server import _render_ingest_result +from slayer.storage.yaml_storage import YAMLStorage + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _mock_engine(dialect_name: str = "postgresql") -> MagicMock: + # No spec: sa.Engine sets ``dialect`` in __init__, so a spec'd mock + # rejects it. + engine = MagicMock() + engine.dialect.name = dialect_name + return engine + + +def _mock_conn_engine(dialect_name: str, execute_side_effect) -> tuple[MagicMock, MagicMock]: + """Engine whose ``connect()`` context yields a conn with stubbed execute.""" + engine = _mock_engine(dialect_name) + conn = MagicMock() + engine.connect.return_value.__enter__ = MagicMock(return_value=conn) + engine.connect.return_value.__exit__ = MagicMock(return_value=False) + if callable(execute_side_effect): + conn.execute.side_effect = execute_side_effect + else: + conn.execute.return_value.fetchall.return_value = execute_side_effect + return engine, conn + + +def _mock_inspector( + columns: list[dict], + *, + pk: list[str] | None = None, + table_comment: object = None, +) -> MagicMock: + inspector = MagicMock(spec=sa.engine.Inspector) + inspector.get_columns.return_value = columns + inspector.get_pk_constraint.return_value = {"constrained_columns": pk or []} + if isinstance(table_comment, Exception): + inspector.get_table_comment.side_effect = table_comment + else: + inspector.get_table_comment.return_value = {"text": table_comment} + return inspector + + +def _commented_duckdb(path: Path) -> None: + conn = duckdb.connect(str(path)) + conn.execute("CREATE TABLE customers (id INTEGER PRIMARY KEY, region VARCHAR)") + conn.execute( + """ + CREATE TABLE orders ( + id INTEGER PRIMARY KEY, + amount DOUBLE, + status VARCHAR, + customer_id INTEGER REFERENCES customers(id) + ) + """ + ) + conn.execute("COMMENT ON TABLE orders IS 'All orders'") + conn.execute("COMMENT ON COLUMN orders.amount IS 'Order amount in USD'") + conn.execute("COMMENT ON COLUMN customers.region IS 'Sales region'") + conn.execute("INSERT INTO customers VALUES (1, 'US')") + conn.execute("INSERT INTO orders VALUES (1, 100.0, 'completed', 1)") + conn.close() + + +def _model_by_name(models: list[SlayerModel], name: str) -> SlayerModel: + return next(m for m in models if m.name == name) + + +def _col(model: SlayerModel, name: str) -> Column: + return next(c for c in model.columns if c.name == name) + + +# --------------------------------------------------------------------------- +# Normalization +# --------------------------------------------------------------------------- + + +class TestCleanComment: + def test_strips_whitespace(self) -> None: + assert _clean_comment(" order id ") == "order id" + + def test_empty_string_is_none(self) -> None: + assert _clean_comment("") is None + + def test_whitespace_only_is_none(self) -> None: + assert _clean_comment(" \n\t ") is None + + def test_none_passthrough(self) -> None: + assert _clean_comment(None) is None + + +# --------------------------------------------------------------------------- +# Inspector path: column + table comments +# --------------------------------------------------------------------------- + + +class TestInspectorPathComments: + def test_column_comments_land_on_descriptions(self) -> None: + inspector = _mock_inspector( + [ + {"name": "id", "type": sa.INTEGER(), "comment": "row id"}, + {"name": "amount", "type": sa.NUMERIC(10, 2), "comment": None}, + {"name": "note", "type": sa.TEXT()}, # no comment key at all + ], + pk=["id"], + table_comment="Orders table", + ) + model = introspect_table_to_model( + sa_engine=_mock_engine(), + inspector=inspector, + table_name="orders", + schema=None, + data_source="ds", + ) + assert model.description == "Orders table" + assert _col(model, "id").description == "row id" + assert _col(model, "amount").description is None + assert _col(model, "note").description is None + + def test_whitespace_comment_normalized_to_none(self) -> None: + inspector = _mock_inspector( + [{"name": "id", "type": sa.INTEGER(), "comment": " "}], + table_comment=" ", + ) + model = introspect_table_to_model( + sa_engine=_mock_engine(), + inspector=inspector, + table_name="t", + schema=None, + data_source="ds", + ) + assert model.description is None + assert _col(model, "id").description is None + + def test_count_rename_keeps_description(self) -> None: + inspector = _mock_inspector( + [{"name": "_count", "type": sa.INTEGER(), "comment": "collides"}], + ) + model = introspect_table_to_model( + sa_engine=_mock_engine(), + inspector=inspector, + table_name="t", + schema=None, + data_source="ds", + ) + assert _col(model, "count_col").description == "collides" + + +class TestSafeGetTableComment: + def test_returns_text(self) -> None: + inspector = _mock_inspector([], table_comment="hello") + assert _safe_get_table_comment(inspector, "t", None) == "hello" + + def test_none_text(self) -> None: + inspector = _mock_inspector([], table_comment=None) + assert _safe_get_table_comment(inspector, "t", None) is None + + def test_not_implemented_is_skipped(self) -> None: + inspector = _mock_inspector([], table_comment=NotImplementedError()) + assert _safe_get_table_comment(inspector, "t", None) is None + + def test_arbitrary_error_is_skipped(self) -> None: + inspector = _mock_inspector([], table_comment=RuntimeError("boom")) + assert _safe_get_table_comment(inspector, "t", None) is None + + def test_whitespace_normalized(self) -> None: + inspector = _mock_inspector([], table_comment=" x ") + assert _safe_get_table_comment(inspector, "t", None) == "x" + + +class TestSqliteProbePreservesComments: + def test_non_sqlite_engine_keeps_comments(self) -> None: + cols = [IntrospectedColumn(name="a", type=DataType.INT, comment="kept")] + out = _sqlite_probe_integer_columns( + sa_engine=_mock_engine("postgresql"), sql_table="t", columns=cols + ) + assert out[0].comment == "kept" + + def test_widened_column_keeps_comment(self) -> None: + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + conn.execute(sa.text("CREATE TABLE t (a INTEGER)")) + conn.execute(sa.text("INSERT INTO t VALUES (1.5)")) + conn.commit() + cols = [IntrospectedColumn(name="a", type=DataType.INT, comment="probed")] + out = _sqlite_probe_integer_columns(sa_engine=engine, sql_table="t", columns=cols) + assert out[0].type is DataType.DOUBLE + assert out[0].comment == "probed" + + +# --------------------------------------------------------------------------- +# information_schema fallback: per-dialect comment SQL +# --------------------------------------------------------------------------- + + +class TestColumnCommentsFallback: + @pytest.mark.parametrize( + ("dialect", "source_marker"), + [ + ("mysql", "column_comment"), + ("mariadb", "column_comment"), + ("snowflake", "information_schema"), + ("clickhouse", "system.columns"), + ("duckdb", "duckdb_columns"), + ("postgresql", "col_description"), + ], + ) + def test_dialect_query_shape(self, dialect: str, source_marker: str) -> None: + engine, conn = _mock_conn_engine(dialect_name=dialect, execute_side_effect=[("id", "row id"), ("x", None)]) + result = _get_column_comments_fallback( + sa_engine=engine, table_name="orders", schema=None + ) + assert result == {"id": "row id"} + args, kwargs = conn.execute.call_args + sql_str = str(args[0]).lower() + assert source_marker in sql_str + params = args[1] if len(args) > 1 else kwargs + assert "orders" in params.values() + assert "orders" not in sql_str + + @pytest.mark.parametrize( + "dialect", ["mysql", "snowflake", "clickhouse", "duckdb", "postgresql"] + ) + def test_schema_is_bound_not_interpolated(self, dialect: str) -> None: + engine, conn = _mock_conn_engine(dialect_name=dialect, execute_side_effect=[("id", "row id")]) + result = _get_column_comments_fallback( + sa_engine=engine, table_name="orders", schema="s1" + ) + assert result == {"id": "row id"} + args, kwargs = conn.execute.call_args + sql_str = str(args[0]) + params = args[1] if len(args) > 1 else kwargs + assert "s1" in params.values() + assert "s1" not in sql_str + + def test_clickhouse_defaults_to_current_database(self) -> None: + engine, conn = _mock_conn_engine(dialect_name="clickhouse", execute_side_effect=[("id", "c")]) + _get_column_comments_fallback(sa_engine=engine, table_name="t", schema=None) + sql_str = str(conn.execute.call_args[0][0]).lower() + assert "currentdatabase()" in sql_str + + @pytest.mark.parametrize( + ("dialect", "default_marker"), + [ + ("mysql", "database()"), + ("snowflake", "current_schema()"), + ("duckdb", "current_schema()"), + ("postgresql", "pg_table_is_visible"), + ], + ) + def test_no_schema_scopes_to_default(self, dialect: str, default_marker: str) -> None: + engine, conn = _mock_conn_engine(dialect_name=dialect, execute_side_effect=[("id", "c")]) + _get_column_comments_fallback(sa_engine=engine, table_name="t", schema=None) + sql_str = str(conn.execute.call_args[0][0]).lower() + assert default_marker in sql_str + + def test_no_literal_interpolation(self) -> None: + engine, conn = _mock_conn_engine(dialect_name="mysql", execute_side_effect=[]) + _get_column_comments_fallback( + sa_engine=engine, + table_name="'; DROP TABLE users;--", + schema="'; DROP TABLE users;--", + ) + args, _ = conn.execute.call_args + assert "DROP TABLE" not in str(args[0]) + + def test_unknown_dialect_returns_empty(self) -> None: + engine, conn = _mock_conn_engine(dialect_name="mssql", execute_side_effect=[("id", "x")]) + assert _get_column_comments_fallback( + sa_engine=engine, table_name="t", schema=None + ) == {} + conn.execute.assert_not_called() + + def test_query_error_returns_empty(self) -> None: + def _boom(*args, **kwargs): + raise RuntimeError("no such view") + + engine, _ = _mock_conn_engine(dialect_name="duckdb", execute_side_effect=_boom) + assert _get_column_comments_fallback( + sa_engine=engine, table_name="t", schema=None + ) == {} + + def test_blank_comments_filtered(self) -> None: + engine, _ = _mock_conn_engine(dialect_name="mysql", execute_side_effect=[("a", ""), ("b", " "), ("c", "ok")]) + assert _get_column_comments_fallback( + sa_engine=engine, table_name="t", schema=None + ) == {"c": "ok"} + + def test_generic_fallback_merges_comments(self) -> None: + def _dispatch(clause, params=None): + res = MagicMock() + if "duckdb_columns" in str(clause): + res.fetchall.return_value = [("id", "row id")] + else: + res.fetchall.return_value = [("id", "INTEGER"), ("x", "VARCHAR")] + return res + + engine, _ = _mock_conn_engine(dialect_name="duckdb", execute_side_effect=_dispatch) + cols = _get_columns_fallback(sa_engine=engine, table_name="t", schema=None) + by_name = {c["name"]: c for c in cols} + assert by_name["id"]["comment"] == "row id" + assert by_name["x"].get("comment") is None + + +# --------------------------------------------------------------------------- +# DuckDB end-to-end (real engine — exercises the genuine fallback path) +# --------------------------------------------------------------------------- + + +class TestDuckDBEndToEnd: + @pytest.fixture + def duckdb_models(self, tmp_path: Path) -> list[SlayerModel]: + db_path = tmp_path / "commented.duckdb" + _commented_duckdb(db_path) + ds = DatasourceConfig(name="ds", type="duckdb", database=str(db_path)) + return ingest_datasource(datasource=ds) + + def test_table_comment_becomes_model_description(self, duckdb_models) -> None: + assert _model_by_name(duckdb_models, "orders").description == "All orders" + + def test_column_comments_become_descriptions(self, duckdb_models) -> None: + orders = _model_by_name(duckdb_models, "orders") + assert _col(orders, "amount").description == "Order amount in USD" + customers = _model_by_name(duckdb_models, "customers") + assert _col(customers, "region").description == "Sales region" + + def test_uncommented_stays_none(self, duckdb_models) -> None: + orders = _model_by_name(duckdb_models, "orders") + assert _col(orders, "status").description is None + assert _model_by_name(duckdb_models, "customers").description is None + + +# --------------------------------------------------------------------------- +# Idempotent merge: fill-if-empty + reporting +# --------------------------------------------------------------------------- + + +async def _duckdb_idempotent_setup(tmp_path: Path) -> tuple[YAMLStorage, DatasourceConfig]: + db_path = tmp_path / "live.duckdb" + _commented_duckdb(db_path) + storage = YAMLStorage(base_dir=str(tmp_path / "storage")) + ds = DatasourceConfig(name="ds", type="duckdb", database=str(db_path)) + await storage.save_datasource(ds) + return storage, ds + + +def _addition_for(name: str, additions) -> ModelAddition | None: + return next((a for a in additions if a.model_name == name), None) + + +class TestIdempotentDescriptionFill: + async def test_created_model_reports_descriptions(self, tmp_path: Path) -> None: + storage, ds = await _duckdb_idempotent_setup(tmp_path) + result = await ingest_datasource_idempotent(datasource=ds, storage=storage) + addition = _addition_for("orders", result.additions) + assert addition is not None + assert addition.created + assert "amount" in addition.described_columns + assert addition.model_described is True + loaded = await storage.get_model("orders", data_source="ds") + assert loaded.description == "All orders" + assert _col(loaded, "amount").description == "Order amount in USD" + + async def test_fill_if_empty_and_preserve_user_text(self, tmp_path: Path) -> None: + storage, ds = await _duckdb_idempotent_setup(tmp_path) + await ingest_datasource_idempotent(datasource=ds, storage=storage) + loaded = await storage.get_model("orders", data_source="ds") + loaded.description = None + for c in loaded.columns: + if c.name == "amount": + c.description = None + if c.name == "status": + c.description = "hand-authored" + await storage.save_model(loaded) + + result = await ingest_datasource_idempotent(datasource=ds, storage=storage) + addition = _addition_for("orders", result.additions) + assert addition is not None + assert not addition.created + assert addition.described_columns == ["amount"] + assert addition.model_described is True + assert addition.new_columns == [] + + loaded2 = await storage.get_model("orders", data_source="ds") + assert loaded2.description == "All orders" + assert _col(loaded2, "amount").description == "Order amount in USD" + assert _col(loaded2, "status").description == "hand-authored" + + async def test_reingest_is_noop_when_descriptions_present(self, tmp_path: Path) -> None: + storage, ds = await _duckdb_idempotent_setup(tmp_path) + await ingest_datasource_idempotent(datasource=ds, storage=storage) + result = await ingest_datasource_idempotent(datasource=ds, storage=storage) + for addition in result.additions: + assert addition.described_columns == [] + assert addition.model_described is False + + async def test_new_commented_column_arrives_with_description( + self, tmp_path: Path + ) -> None: + storage, ds = await _duckdb_idempotent_setup(tmp_path) + await ingest_datasource_idempotent(datasource=ds, storage=storage) + conn = duckdb.connect(ds.database) + conn.execute("ALTER TABLE orders ADD COLUMN discount DOUBLE") + conn.execute("COMMENT ON COLUMN orders.discount IS 'Discount applied'") + conn.close() + + result = await ingest_datasource_idempotent(datasource=ds, storage=storage) + addition = _addition_for("orders", result.additions) + assert addition is not None + assert "discount" in addition.new_columns + # New columns carry their description implicitly — not double-counted. + assert "discount" not in addition.described_columns + loaded = await storage.get_model("orders", data_source="ds") + assert _col(loaded, "discount").description == "Discount applied" + + +class TestAdditiveMergeDescriptions: + def _persisted(self, description=None, col_description=None) -> SlayerModel: + return SlayerModel( + name="t", + sql_table="t", + data_source="ds", + description=description, + columns=[ + Column(name="a", sql="a", type=DataType.INT, description=col_description) + ], + ) + + def _fresh(self) -> SlayerModel: + return SlayerModel( + name="t", + sql_table="t", + data_source="ds", + description="fresh model desc", + columns=[ + Column(name="a", sql="a", type=DataType.INT, description="fresh col desc") + ], + ) + + def test_fills_empty_descriptions(self) -> None: + outcome = _additive_merge_existing( + persisted=self._persisted(), fresh=self._fresh() + ) + assert outcome.described_columns == ["a"] + assert outcome.model_described is True + assert outcome.merged.description == "fresh model desc" + assert _col(outcome.merged, "a").description == "fresh col desc" + assert outcome.new_columns == [] + assert outcome.new_joins == [] + assert outcome.widened_columns == [] + + def test_existing_descriptions_untouched(self) -> None: + outcome = _additive_merge_existing( + persisted=self._persisted(description="mine", col_description="my col"), + fresh=self._fresh(), + ) + assert outcome.described_columns == [] + assert outcome.model_described is False + assert outcome.merged.description == "mine" + assert _col(outcome.merged, "a").description == "my col" + + +# --------------------------------------------------------------------------- +# BigQuery dataset description +# --------------------------------------------------------------------------- + + +def _bq_engine(dataset_id: str | None, description: str | None = "ds desc"): + """Mock BigQuery engine exposing the client the way sqlalchemy-bigquery does.""" + engine = _mock_engine("bigquery") + engine.dialect.dataset_id = dataset_id + client = MagicMock() + client.get_dataset.return_value = SimpleNamespace(description=description) + conn = MagicMock() + conn.connection._client = client + engine.connect.return_value.__enter__ = MagicMock(return_value=conn) + engine.connect.return_value.__exit__ = MagicMock(return_value=False) + return engine, client + + +class TestFetchBigQueryDatasetDescription: + def test_non_bigquery_dialect_returns_none(self) -> None: + ds = DatasourceConfig(name="d", type="duckdb", database="x.db") + assert ( + _fetch_bigquery_dataset_description( + sa_engine=_mock_engine("duckdb"), datasource=ds, schema=None + ) + is None + ) + + def test_explicit_schema_wins(self) -> None: + engine, client = _bq_engine(dataset_id="dialect_ds") + ds = DatasourceConfig(name="d", type="bigquery", schema_name="cfg_ds") + out = _fetch_bigquery_dataset_description( + sa_engine=engine, datasource=ds, schema="explicit_ds" + ) + assert out == "ds desc" + assert client.get_dataset.call_args[0][0] == "explicit_ds" + + def test_dialect_default_beats_schema_name(self) -> None: + engine, client = _bq_engine(dataset_id="dialect_ds") + ds = DatasourceConfig(name="d", type="bigquery", schema_name="cfg_ds") + _fetch_bigquery_dataset_description(sa_engine=engine, datasource=ds, schema=None) + assert client.get_dataset.call_args[0][0] == "dialect_ds" + + def test_schema_name_is_last_resort(self) -> None: + engine, client = _bq_engine(dataset_id=None) + ds = DatasourceConfig(name="d", type="bigquery", schema_name="cfg_ds") + _fetch_bigquery_dataset_description(sa_engine=engine, datasource=ds, schema=None) + assert client.get_dataset.call_args[0][0] == "cfg_ds" + + def test_unresolvable_dataset_returns_none(self) -> None: + engine, client = _bq_engine(dataset_id=None) + ds = DatasourceConfig(name="d", type="bigquery") + assert ( + _fetch_bigquery_dataset_description( + sa_engine=engine, datasource=ds, schema=None + ) + is None + ) + client.get_dataset.assert_not_called() + + def test_client_error_returns_none(self) -> None: + engine, client = _bq_engine(dataset_id="d1") + client.get_dataset.side_effect = RuntimeError("403") + ds = DatasourceConfig(name="d", type="bigquery") + assert ( + _fetch_bigquery_dataset_description( + sa_engine=engine, datasource=ds, schema=None + ) + is None + ) + + def test_whitespace_description_is_none(self) -> None: + engine, _ = _bq_engine(dataset_id="d1", description=" ") + ds = DatasourceConfig(name="d", type="bigquery") + assert ( + _fetch_bigquery_dataset_description( + sa_engine=engine, datasource=ds, schema=None + ) + is None + ) + + +class TestIngestDatasourceReport: + def test_returns_models_and_no_description_for_duckdb(self, tmp_path: Path) -> None: + db_path = tmp_path / "x.duckdb" + _commented_duckdb(db_path) + ds = DatasourceConfig(name="ds", type="duckdb", database=str(db_path)) + out = ingest_datasource_report(datasource=ds) + assert isinstance(out, IngestionScanReport) + assert {m.name for m in out.models} == {"customers", "orders"} + assert out.schema_description is None + + def test_wrapper_returns_plain_model_list(self, tmp_path: Path) -> None: + db_path = tmp_path / "x.duckdb" + _commented_duckdb(db_path) + ds = DatasourceConfig(name="ds", type="duckdb", database=str(db_path)) + models = ingest_datasource(datasource=ds) + assert isinstance(models, list) + assert all(isinstance(m, SlayerModel) for m in models) + + def test_fetch_skipped_when_description_set( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + fetch = MagicMock(return_value="never used") + monkeypatch.setattr( + ingestion_mod, "_fetch_bigquery_dataset_description", fetch + ) + db_path = tmp_path / "x.duckdb" + _commented_duckdb(db_path) + ds = DatasourceConfig( + name="ds", type="duckdb", database=str(db_path), description="already set" + ) + out = ingest_datasource_report(datasource=ds) + fetch.assert_not_called() + assert out.schema_description is None + + def test_fetch_called_when_description_empty( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + fetch = MagicMock(return_value="dataset says hi") + monkeypatch.setattr( + ingestion_mod, "_fetch_bigquery_dataset_description", fetch + ) + db_path = tmp_path / "x.duckdb" + _commented_duckdb(db_path) + ds = DatasourceConfig(name="ds", type="duckdb", database=str(db_path)) + out = ingest_datasource_report(datasource=ds) + fetch.assert_called_once() + assert out.schema_description == "dataset says hi" + + +# --------------------------------------------------------------------------- +# Idempotent path: datasource description persistence + failure isolation +# --------------------------------------------------------------------------- + + +def _sqlite_live_db(tmp_path: Path) -> str: + db_path = str(tmp_path / "live.db") + conn = sqlite3.connect(db_path) + conn.execute("CREATE TABLE t (id INTEGER PRIMARY KEY)") + conn.commit() + conn.close() + return db_path + + +def _patch_full_ingest(monkeypatch: pytest.MonkeyPatch, description: str | None): + real = ingestion_mod.ingest_datasource_report + + def _fake(*args, **kwargs): + out = real(*args, **kwargs) + return out.model_copy(update={"schema_description": description}) + + monkeypatch.setattr(ingestion_mod, "ingest_datasource_report", _fake) + + +class TestIdempotentDatasourceDescription: + async def test_description_persisted_and_reported( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + db_path = _sqlite_live_db(tmp_path) + storage = YAMLStorage(base_dir=str(tmp_path / "storage")) + ds = DatasourceConfig(name="ds", type="sqlite", database=db_path) + await storage.save_datasource(ds) + _patch_full_ingest(monkeypatch, "From the dataset") + + result = await ingest_datasource_idempotent(datasource=ds, storage=storage) + assert result.datasource_described is True + loaded = await storage.get_datasource("ds") + assert loaded.description == "From the dataset" + + # Re-ingest with the reloaded datasource (as the CLI does): the + # persisted description must make the second pass a no-op. + result2 = await ingest_datasource_idempotent(datasource=loaded, storage=storage) + assert result2.datasource_described is False + reloaded = await storage.get_datasource("ds") + assert reloaded.description == "From the dataset" + + async def test_existing_description_never_overwritten( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + db_path = _sqlite_live_db(tmp_path) + storage = YAMLStorage(base_dir=str(tmp_path / "storage")) + ds = DatasourceConfig( + name="ds", type="sqlite", database=db_path, description="user text" + ) + await storage.save_datasource(ds) + _patch_full_ingest(monkeypatch, "From the dataset") + + result = await ingest_datasource_idempotent(datasource=ds, storage=storage) + assert result.datasource_described is False + loaded = await storage.get_datasource("ds") + assert loaded.description == "user text" + + async def test_save_failure_is_isolated( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + db_path = _sqlite_live_db(tmp_path) + + class _FailingSave(YAMLStorage): + async def save_datasource(self, datasource: DatasourceConfig) -> None: + if getattr(self, "_armed", False): + raise RuntimeError("disk full") + await super().save_datasource(datasource) + + storage = _FailingSave(base_dir=str(tmp_path / "storage")) + ds = DatasourceConfig(name="ds", type="sqlite", database=db_path) + await storage.save_datasource(ds) + storage._armed = True + _patch_full_ingest(monkeypatch, "From the dataset") + + result = await ingest_datasource_idempotent(datasource=ds, storage=storage) + assert result.datasource_described is False + assert any( + e.model_name == "" and "datasource" in e.error.lower() + for e in result.errors + ) + # Model ingestion itself still succeeded. + assert _addition_for("t", result.additions) is not None + + +# --------------------------------------------------------------------------- +# Report shape + rendering +# --------------------------------------------------------------------------- + + +class TestReportShape: + def test_model_addition_defaults(self) -> None: + addition = ModelAddition(model_name="m", data_source="ds") + assert addition.described_columns == [] + assert addition.model_described is False + + def test_result_default(self) -> None: + assert IdempotentIngestResult().datasource_described is False + + def test_print_updated_with_descriptions(self) -> None: + addition = ModelAddition( + model_name="orders", + data_source="ds", + described_columns=["amount", "status"], + model_described=True, + ) + buf = io.StringIO() + _print_ingest_addition(addition, file=buf) + out = buf.getvalue() + assert "+descriptions: amount, status" in out + assert "+model description" in out + + def test_print_created_with_descriptions(self) -> None: + addition = ModelAddition( + model_name="orders", + data_source="ds", + created=True, + new_columns=["a", "b", "c"], + described_columns=["a", "b"], + ) + buf = io.StringIO() + _print_ingest_addition(addition, file=buf) + assert "Created: orders (3 columns, 2 described)" in buf.getvalue() + + def test_print_created_without_descriptions_unchanged(self) -> None: + addition = ModelAddition( + model_name="orders", data_source="ds", created=True, new_columns=["a"] + ) + buf = io.StringIO() + _print_ingest_addition(addition, file=buf) + assert "Created: orders (1 columns)" in buf.getvalue() + + def test_description_only_update_is_printed(self) -> None: + addition = ModelAddition( + model_name="orders", data_source="ds", described_columns=["amount"] + ) + buf = io.StringIO() + _print_ingest_addition(addition, file=buf) + assert "Updated: orders" in buf.getvalue() + + +class TestMcpIngestRender: + def test_description_only_update_rendered(self) -> None: + result = IdempotentIngestResult( + additions=[ModelAddition( + model_name="orders", data_source="ds", + described_columns=["amount"], model_described=True, + )], + ) + out = _render_ingest_result( + result, + schema_name="", + ds=DatasourceConfig(name="ds", type="sqlite", database=":memory:"), + ) + assert "+descriptions: amount" in out + assert "+model description" in out + + def test_datasource_only_description_not_swallowed(self) -> None: + result = IdempotentIngestResult(datasource_described=True) + out = _render_ingest_result( + result, + schema_name="", + ds=DatasourceConfig(name="ds", type="sqlite", database=":memory:"), + ) + assert "Datasource description imported." in out + assert "already in sync" not in out + + +# --------------------------------------------------------------------------- +# CLI `datasources create --ingest` wiring +# --------------------------------------------------------------------------- + + +class TestCliDatasourcesCreateIngest: + def _args(self, db_path: str, description: str | None = None) -> argparse.Namespace: + return argparse.Namespace( + connection_string=f"sqlite:///{db_path}", + name="ds", + description=description, + yes=True, + ingest=True, + include=None, + exclude=None, + schema=None, + ) + + def test_dataset_description_imported( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + db_path = _sqlite_live_db(tmp_path) + storage = YAMLStorage(base_dir=str(tmp_path / "storage")) + _patch_full_ingest(monkeypatch, "Imported dataset description") + + _run_datasources_create(self._args(db_path), storage) + loaded = run_sync(storage.get_datasource("ds")) + assert loaded.description == "Imported dataset description" + assert run_sync(storage.get_model("t", data_source="ds")) is not None + + def test_user_description_wins( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + db_path = _sqlite_live_db(tmp_path) + storage = YAMLStorage(base_dir=str(tmp_path / "storage")) + _patch_full_ingest(monkeypatch, "Imported dataset description") + + _run_datasources_create(self._args(db_path, description="user says"), storage) + loaded = run_sync(storage.get_datasource("ds")) + assert loaded.description == "user says" + + +# --------------------------------------------------------------------------- +# BigQuery driver contract (credential-free — pins the installed package) +# --------------------------------------------------------------------------- + + +class TestBigQueryDriverContract: + def test_get_columns_carries_field_description(self) -> None: + pytest.importorskip("sqlalchemy_bigquery") + from google.cloud.bigquery import SchemaField + from sqlalchemy_bigquery._types import get_columns as bq_get_columns + + cols = bq_get_columns([SchemaField("id", "INTEGER", description="row id")]) + assert cols[0]["comment"] == "row id" + + def test_get_table_comment_returns_table_description(self) -> None: + pytest.importorskip("sqlalchemy_bigquery") + from sqlalchemy_bigquery import BigQueryDialect + + dialect = BigQueryDialect() + dialect._get_table = MagicMock( + return_value=SimpleNamespace(description="tbl desc") + ) + out = dialect.get_table_comment(MagicMock(), "t") + assert out == {"text": "tbl desc"} + + def test_real_dialect_reflection_flows_into_slayer_model(self) -> None: + """Run the REAL BigQueryDialect reflection over a locally built + Table — only the network call is mocked — and feed its genuine + output through introspect_table_to_model.""" + pytest.importorskip("sqlalchemy_bigquery") + from google.cloud.bigquery import SchemaField, Table + from sqlalchemy_bigquery import BigQueryDialect + + table = Table("proj.dset.orders", schema=[ + SchemaField("id", "INTEGER", description="row id"), + SchemaField("amount", "FLOAT", description="order amount"), + SchemaField("status", "STRING"), + ]) + table.description = "All orders" + + dialect = BigQueryDialect() + with patch.object(BigQueryDialect, "_get_table", return_value=table): + real_cols = dialect.get_columns(MagicMock(), "orders") + real_comment = dialect.get_table_comment(MagicMock(), "orders") + + inspector = MagicMock(spec=sa.engine.Inspector) + inspector.get_columns.return_value = real_cols + inspector.get_pk_constraint.return_value = {"constrained_columns": []} + inspector.get_table_comment.return_value = real_comment + + model = introspect_table_to_model( + sa_engine=_mock_engine("bigquery"), + inspector=inspector, + table_name="orders", + schema="dset", + data_source="bq", + ) + assert model.description == "All orders" + cols = {c.name: c for c in model.columns} + assert cols["id"].description == "row id" + assert cols["id"].type is DataType.INT + assert cols["amount"].description == "order amount" + assert cols["amount"].type is DataType.DOUBLE + assert cols["status"].description is None diff --git a/tests/test_ingestion_name_sanitize.py b/tests/test_ingestion_name_sanitize.py index a15bf49d..8eaa16d4 100644 --- a/tests/test_ingestion_name_sanitize.py +++ b/tests/test_ingestion_name_sanitize.py @@ -544,7 +544,7 @@ def _introspect(self, workspace: Path, joins): sa_engine, inspector = self._fixture(workspace) return [ - c[0] + c.name for c in _introspect_query_columns_via_inspector( sa_engine=sa_engine, inspector=inspector, diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 09b471ca..ca82200f 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -3,12 +3,14 @@ import json import os import shutil +import sqlite3 import tempfile from typing import Any from collections.abc import Generator import pytest +import slayer.engine.ingestion as ingestion_mod from slayer.core.enums import DataType from slayer.core.models import ( Aggregation, @@ -2948,6 +2950,30 @@ async def test_create_reports_replaced(self, mcp_server, storage: YAMLStorage) - result = await _call(mcp_server, name="create_datasource", arguments={"name": "ds", "type": "sqlite", "database": ":memory:"}) assert "replaced" in result + async def test_create_auto_ingest_imports_dataset_description( + self, mcp_server, storage: YAMLStorage, tmp_path, monkeypatch + ) -> None: + db_path = str(tmp_path / "live.db") + conn = sqlite3.connect(db_path) + conn.execute("CREATE TABLE t (id INTEGER PRIMARY KEY)") + conn.commit() + conn.close() + + real = ingestion_mod.ingest_datasource_report + + def _fake(*args, **kwargs): + out = real(*args, **kwargs) + return out.model_copy(update={"schema_description": "From the dataset"}) + + monkeypatch.setattr(ingestion_mod, "ingest_datasource_report", _fake) + result = await _call(mcp_server, name="create_datasource", arguments={ + "name": "ds", "type": "sqlite", "database": db_path, + }) + assert "ds" in result + loaded = await storage.get_datasource("ds") + assert loaded.description == "From the dataset" + assert await storage.get_model("t", data_source="ds") is not None + async def test_list_with_malformed_datasource(self, mcp_server, storage: YAMLStorage) -> None: # A valid datasource alongside a malformed one await storage.save_datasource(DatasourceConfig(name="good", type="sqlite", database=":memory:")) diff --git a/tests/test_osi_converter.py b/tests/test_osi_converter.py index 2610463f..ef6aab4a 100644 --- a/tests/test_osi_converter.py +++ b/tests/test_osi_converter.py @@ -118,6 +118,33 @@ def test_model_and_semantic_model_ai_context(shop_engine): assert "osi_semantic_model" in orders.meta +def test_db_comments_fill_gaps_curated_wins(shop_engine, monkeypatch): + """Introspected DB comments survive only where OSI has no curated text.""" + import slayer.osi.converter as osi_converter_module + + real = osi_converter_module.introspect_table_to_model + + def _with_db_comments(**kwargs): + model = real(**kwargs) + model.description = "db table comment" + for c in model.columns: + if c.name in ("order_id", "customer_id"): + c.description = f"db comment {c.name}" + return model + + monkeypatch.setattr( + osi_converter_module, "introspect_table_to_model", _with_db_comments + ) + orders = _by_name(_shop_result(shop_engine))["orders"] + cols = {c.name: c for c in orders.columns} + # Curated OSI description wins over the DB comment. + assert "Order line items" in orders.description + assert "db table comment" not in orders.description + assert cols["order_id"].description == "Order id" + # No curated metadata on customer_id — the DB comment survives. + assert cols["customer_id"].description == "db comment customer_id" + + # ─────────────────────────── relationships -> joins ───────────────────────── def test_joins_from_relationships(shop_engine):