From 114c463c2b8adcd7bbbe6eeb7c40437fa7ddad81 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 11:35:49 +0100 Subject: [PATCH 1/7] perf(adapter): add an invocation-scoped warehouse inspection catalog Add a shared InspectionCatalog that serves list_relations and get_columns_for_relations requests for one planning invocation. Adapters implementing SchemaScopedMetadataReader (Snowflake) read each schema's table list and column list at most once with unfiltered INFORMATION_SCHEMA queries and filter client-side using the stored identifier case, so no metadata query carries a large IN list. Sparse column requests keep exact SHOW COLUMNS reads. Independent schema reads run with bounded concurrency on per-thread cursors. Other adapters keep their request shapes; identical requests are memoized. Snowflake source freshness reuses the catalog's schema listing when a catalog is open, and every inspection read can be recorded for verbose output. Refs CHI-455 --- .../adapter/contract/classes/base_adapter.py | 1 + .../relations/_helpers/inspection_context.py | 74 ++ .../relations/classes/inspection_catalog.py | 342 +++++++++ src/sqlbuild/adapter/relations/constants.py | 2 + .../main/active_inspection_catalog.py | 14 + .../main/get_columns_for_inspection.py | 20 + .../relations/main/inspection_query_sink.py | 24 + .../main/list_relations_for_inspection.py | 27 + .../relations/main/open_inspection_catalog.py | 31 + .../relations/main/record_inspection_query.py | 12 + .../relations/main/render_inspection_sql.py | 18 + .../relations/main/run_bounded_inspections.py | 55 ++ .../main/run_recorded_inspection_query.py | 34 + src/sqlbuild/adapter/relations/models.py | 44 ++ src/sqlbuild/adapter/relations/types.py | 47 +- .../snowflake/classes/snowflake_adapter.py | 249 +++++-- .../run_bounded_inspections/_test_types.py | 25 + .../main/run_bounded_inspections/helpers.py | 47 ++ .../test_run_bounded_inspections.py | 93 +++ .../adapters/snowflake/inspection/__init__.py | 0 .../snowflake/inspection/_test_types.py | 63 ++ .../adapters/snowflake/inspection/helpers.py | 648 ++++++++++++++++++ .../inspection/test_inspection_catalog.py | 310 +++++++++ 23 files changed, 2129 insertions(+), 51 deletions(-) create mode 100644 src/sqlbuild/adapter/relations/_helpers/inspection_context.py create mode 100644 src/sqlbuild/adapter/relations/classes/inspection_catalog.py create mode 100644 src/sqlbuild/adapter/relations/main/active_inspection_catalog.py create mode 100644 src/sqlbuild/adapter/relations/main/get_columns_for_inspection.py create mode 100644 src/sqlbuild/adapter/relations/main/inspection_query_sink.py create mode 100644 src/sqlbuild/adapter/relations/main/list_relations_for_inspection.py create mode 100644 src/sqlbuild/adapter/relations/main/open_inspection_catalog.py create mode 100644 src/sqlbuild/adapter/relations/main/record_inspection_query.py create mode 100644 src/sqlbuild/adapter/relations/main/render_inspection_sql.py create mode 100644 src/sqlbuild/adapter/relations/main/run_bounded_inspections.py create mode 100644 src/sqlbuild/adapter/relations/main/run_recorded_inspection_query.py create mode 100644 src/sqlbuild/adapter/relations/models.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/_test_types.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/helpers.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/test_run_bounded_inspections.py create mode 100644 tests/unit/src/sqlbuild/adapters/snowflake/inspection/__init__.py create mode 100644 tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py create mode 100644 tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py create mode 100644 tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py diff --git a/src/sqlbuild/adapter/contract/classes/base_adapter.py b/src/sqlbuild/adapter/contract/classes/base_adapter.py index e1de3f9f62..f0aff72de9 100644 --- a/src/sqlbuild/adapter/contract/classes/base_adapter.py +++ b/src/sqlbuild/adapter/contract/classes/base_adapter.py @@ -107,6 +107,7 @@ class BaseAdapter(RetentionAdapterMixin, StrictAdapter): views_read_with_reader_access: ClassVar[bool] = False allows_implicit_managed_write_schema: ClassVar[bool] = False execution_duration_limit_seconds: ClassVar[int | None] = None + metadata_inspection_concurrency: ClassVar[int] = 1 _snapshot_sql_dialect: ClassVar[SnapshotSqlDialect] = SnapshotSqlDialect( timestamp_type="TIMESTAMP", distinct_condition=render_is_distinct_from, diff --git a/src/sqlbuild/adapter/relations/_helpers/inspection_context.py b/src/sqlbuild/adapter/relations/_helpers/inspection_context.py new file mode 100644 index 0000000000..007b395852 --- /dev/null +++ b/src/sqlbuild/adapter/relations/_helpers/inspection_context.py @@ -0,0 +1,74 @@ +"""Invocation-local inspection catalog and verbose query sink state.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from contextvars import ContextVar, Token +from typing import TYPE_CHECKING, Any + +from sqlbuild.adapter.relations.models import InspectionQueryRecord +from sqlbuild.diagnostics.main.log_debug_event import log_debug_event +from sqlbuild.diagnostics.main.log_sql import log_sql + +if TYPE_CHECKING: + from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog + +_LOGGER: logging.Logger = logging.getLogger("sqlbuild.inspection") +_ACTIVE_CATALOG: ContextVar[InspectionCatalog | None] = ContextVar( + "sqlbuild_inspection_catalog", default=None +) +_QUERY_SINK: ContextVar[Callable[[InspectionQueryRecord], None] | None] = ContextVar( + "sqlbuild_inspection_query_sink", default=None +) + + +def current_catalog(*, adapter: Any, connection: Any) -> InspectionCatalog | None: + """Return the open catalog for exactly this adapter and connection, if any.""" + + catalog: InspectionCatalog | None = _ACTIVE_CATALOG.get() + if catalog is None or catalog.adapter is not adapter or catalog.connection is not connection: + return None + return catalog + + +def emit_inspection_record(*, record: InspectionQueryRecord) -> None: + """Log one inspection read and forward it to the active verbose sink.""" + + log_sql(logger=_LOGGER, sql=record.sql, action="inspection") + log_debug_event( + logger=_LOGGER, + message="warehouse inspection query", + sqlbuild_elapsed_seconds=f"{record.elapsed_seconds:.3f}", + sqlbuild_row_count=record.row_count, + sqlbuild_error=record.error, + ) + sink: Callable[[InspectionQueryRecord], None] | None = _QUERY_SINK.get() + if sink is not None: + sink(record) + + +def activate_catalog(catalog: InspectionCatalog) -> Token[InspectionCatalog | None]: + """Make ``catalog`` the open catalog for this context.""" + + return _ACTIVE_CATALOG.set(catalog) + + +def deactivate_catalog(token: Token[InspectionCatalog | None]) -> None: + """Restore the catalog that was open before ``activate_catalog``.""" + + _ACTIVE_CATALOG.reset(token) + + +def activate_query_sink( + sink: Callable[[InspectionQueryRecord], None], +) -> Token[Callable[[InspectionQueryRecord], None] | None]: + """Deliver recorded inspection reads in this context to ``sink``.""" + + return _QUERY_SINK.set(sink) + + +def deactivate_query_sink(token: Token[Callable[[InspectionQueryRecord], None] | None]) -> None: + """Restore the sink that was active before ``activate_query_sink``.""" + + _QUERY_SINK.reset(token) diff --git a/src/sqlbuild/adapter/relations/classes/inspection_catalog.py b/src/sqlbuild/adapter/relations/classes/inspection_catalog.py new file mode 100644 index 0000000000..6a80deb62a --- /dev/null +++ b/src/sqlbuild/adapter/relations/classes/inspection_catalog.py @@ -0,0 +1,342 @@ +"""Invocation-scoped warehouse metadata shared by every planning inspection phase.""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any + +from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo +from sqlbuild.adapter.relations._helpers.inspection_context import emit_inspection_record +from sqlbuild.adapter.relations.constants import ( + INSPECTION_DESCRIBED_NAME_LIMIT, + INSPECTION_IN_LIST_LIMIT, +) +from sqlbuild.adapter.relations.main.run_bounded_inspections import run_bounded_inspections +from sqlbuild.adapter.relations.models import ( + InspectionQueryRecord, + SchemaColumnListing, + SchemaRelationListing, +) +from sqlbuild.adapter.relations.types import SchemaScopedMetadataReader + +type _RelationIdentity = tuple[str | None, str | None, str] +type _SchemaKey = tuple[str | None, str] +type _RelationRequestKey = tuple[str | None, tuple[str, ...] | None, tuple[str, ...] | None] + + +class InspectionCatalog: + """Read each schema's metadata at most once per planning invocation, filtering client-side.""" + + def __init__(self, *, adapter: Any, connection: Any) -> None: + self.adapter: Any = adapter + self.connection: Any = connection + self._reader: SchemaScopedMetadataReader | None = ( + adapter if isinstance(adapter, SchemaScopedMetadataReader) else None + ) + self._concurrency: int = max(1, int(getattr(adapter, "metadata_inspection_concurrency", 1))) + self._relation_listings: dict[_SchemaKey, SchemaRelationListing] = {} + self._column_listings: dict[_SchemaKey, SchemaColumnListing] = {} + self._relation_requests: dict[_RelationRequestKey, tuple[RelationInfo, ...]] = {} + self._relation_columns: dict[_RelationIdentity, tuple[ColumnInfo, ...] | None] = {} + + @property + def concurrency(self) -> int: + """Return the bounded number of concurrent reads this adapter supports.""" + + return self._concurrency + + def cached_relation_listing( + self, *, database: str | None, schema: str + ) -> SchemaRelationListing | None: + """Return a schema listing already read in this invocation, without reading.""" + + return self._relation_listings.get(self._schema_key(database=database, schema=schema)) + + def list_relations( + self, + *, + database: str | None, + schemas: tuple[str, ...] | None, + names: tuple[str, ...] | None = None, + ) -> tuple[RelationInfo, ...]: + """Return the relations ``adapter.list_relations`` would return for this request.""" + + if self._reader is None or not schemas: + return self._list_relations_directly(database=database, schemas=schemas, names=names) + unique_schemas: tuple[str, ...] = tuple(dict.fromkeys(schemas)) + self.prefetch_relation_listings(database=database, schemas=unique_schemas) + wanted: frozenset[str] | None = ( + None if names is None else frozenset(self._reader.metadata_name_key(n) for n in names) + ) + relations: list[RelationInfo] = [] + schema: str + for schema in unique_schemas: + listing: SchemaRelationListing = self._relation_listings[ + self._schema_key(database=database, schema=schema) + ] + relations.extend( + entry.relation + for entry in listing.entries + if wanted is None or entry.stored_name in wanted + ) + return tuple(relations) + + def prefetch_relation_listings(self, *, database: str | None, schemas: tuple[str, ...]) -> None: + """Read every not-yet-listed schema of one database concurrently.""" + + self.prefetch_schema_listings(scopes=tuple((database, schema) for schema in schemas)) + + def prefetch_schema_listings( + self, *, scopes: tuple[tuple[str | None, str], ...], best_effort: bool = False + ) -> None: + """Read unlisted schemas concurrently; best-effort failures stay uncached for re-read.""" + + reader: SchemaScopedMetadataReader | None = self._reader + if reader is None: + return + missing: dict[_SchemaKey, tuple[str | None, str]] = {} + database: str | None + schema: str + for database, schema in scopes: + key: _SchemaKey = self._schema_key(database=database, schema=schema) + if key not in self._relation_listings: + missing.setdefault(key, (database, schema)) + listings: list[SchemaRelationListing | None] = run_bounded_inspections( + tasks=tuple( + self._bind_relation_listing( + reader=reader, database=database, schema=schema, best_effort=best_effort + ) + for database, schema in missing.values() + ), + concurrency=self._concurrency, + ) + missing_key: _SchemaKey + listing: SchemaRelationListing | None + for missing_key, listing in zip(missing, listings, strict=True): + if listing is not None: + self._relation_listings[missing_key] = listing + + def get_columns_for_relations( + self, *, relations: tuple[RelationInfo, ...] + ) -> dict[_RelationIdentity, tuple[ColumnInfo, ...]]: + """Return the columns ``adapter.get_columns_for_relations`` would return.""" + + unique: dict[_RelationIdentity, RelationInfo] = {} + relation: RelationInfo + for relation in relations: + unique.setdefault(relation.identity, relation) + missing: tuple[RelationInfo, ...] = tuple( + relation + for identity, relation in unique.items() + if identity not in self._relation_columns + ) + if missing: + if self._reader is None: + self._read_columns_directly(relations=missing) + else: + self._read_columns_by_schema(reader=self._reader, relations=missing) + result: dict[_RelationIdentity, tuple[ColumnInfo, ...]] = {} + identity: _RelationIdentity + for identity in unique: + columns: tuple[ColumnInfo, ...] | None = self._relation_columns.get(identity) + if columns is not None: + result[identity] = columns + return result + + def _list_relations_directly( + self, + *, + database: str | None, + schemas: tuple[str, ...] | None, + names: tuple[str, ...] | None, + ) -> tuple[RelationInfo, ...]: + key: _RelationRequestKey = ( + database, + None if schemas is None else tuple(schemas), + None if names is None else tuple(names), + ) + cached: tuple[RelationInfo, ...] | None = self._relation_requests.get(key) + if cached is not None: + return cached + relations: tuple[RelationInfo, ...] + if self._reader is not None and names is not None and len(names) > INSPECTION_IN_LIST_LIMIT: + unique_names: tuple[str, ...] = tuple(dict.fromkeys(names)) + chunks: tuple[tuple[str, ...], ...] = tuple( + unique_names[start : start + INSPECTION_IN_LIST_LIMIT] + for start in range(0, len(unique_names), INSPECTION_IN_LIST_LIMIT) + ) + chunk_results: list[tuple[RelationInfo, ...]] = run_bounded_inspections( + tasks=tuple( + self._bind_direct_listing(database=database, schemas=schemas, names=chunk) + for chunk in chunks + ), + concurrency=self._concurrency, + ) + relations = _flatten_relations(chunk_results) + else: + relations = self._bind_direct_listing(database=database, schemas=schemas, names=names)() + self._relation_requests[key] = relations + return relations + + def _bind_direct_listing( + self, + *, + database: str | None, + schemas: tuple[str, ...] | None, + names: tuple[str, ...] | None, + ) -> Callable[[], tuple[RelationInfo, ...]]: + def list_directly() -> tuple[RelationInfo, ...]: + started: float = time.monotonic() + relations: tuple[RelationInfo, ...] = self.adapter.list_relations( + connection=self.connection, database=database, schemas=schemas, names=names + ) + emit_inspection_record( + record=InspectionQueryRecord( + sql=( + f"list_relations(database={database}, " + f"schemas={_describe(schemas)}, names={_describe(names)})" + ), + elapsed_seconds=time.monotonic() - started, + row_count=len(relations), + ) + ) + return relations + + return list_directly + + def _bind_relation_listing( + self, + *, + reader: SchemaScopedMetadataReader, + database: str | None, + schema: str, + best_effort: bool = False, + ) -> Callable[[], SchemaRelationListing | None]: + def read_listing() -> SchemaRelationListing | None: + try: + return reader.read_schema_relation_listing( + connection=self.connection, database=database, schema=schema + ) + except Exception: + if best_effort: + return None + raise + + return read_listing + + def _read_columns_directly(self, *, relations: tuple[RelationInfo, ...]) -> None: + started: float = time.monotonic() + columns: dict[_RelationIdentity, tuple[ColumnInfo, ...]] = ( + self.adapter.get_columns_for_relations(connection=self.connection, relations=relations) + ) + emit_inspection_record( + record=InspectionQueryRecord( + sql=f"get_columns_for_relations(relations={len(relations)})", + elapsed_seconds=time.monotonic() - started, + row_count=sum(len(value) for value in columns.values()), + ) + ) + relation: RelationInfo + for relation in relations: + self._relation_columns[relation.identity] = columns.get(relation.identity) + + def _read_columns_by_schema( + self, *, reader: SchemaScopedMetadataReader, relations: tuple[RelationInfo, ...] + ) -> None: + unscoped: list[RelationInfo] = [] + by_schema: dict[_SchemaKey, list[RelationInfo]] = {} + relation: RelationInfo + for relation in relations: + if relation.schema is None: + unscoped.append(relation) + continue + by_schema.setdefault( + self._schema_key(database=relation.database, schema=relation.schema), [] + ).append(relation) + if unscoped: + self._read_columns_directly(relations=tuple(unscoped)) + listing_keys: list[_SchemaKey] = [] + listing_tasks: list[Callable[[], SchemaColumnListing]] = [] + exact_relations: list[RelationInfo] = [] + schema_key: _SchemaKey + scoped: list[RelationInfo] + for schema_key, scoped in by_schema.items(): + if schema_key in self._column_listings: + continue + if len(scoped) <= reader.exact_column_inspection_limit: + exact_relations.extend(scoped) + continue + listing_keys.append(schema_key) + listing_tasks.append( + self._bind_column_listing( + reader=reader, database=scoped[0].database, schema=scoped[0].schema or "" + ) + ) + exact_tasks: list[Callable[[], tuple[ColumnInfo, ...]]] = [ + self._bind_relation_columns(reader=reader, relation=exact) for exact in exact_relations + ] + results: list[SchemaColumnListing | tuple[ColumnInfo, ...]] = run_bounded_inspections( + tasks=(*listing_tasks, *exact_tasks), concurrency=self._concurrency + ) + listing: SchemaColumnListing | tuple[ColumnInfo, ...] + for schema_key, listing in zip(listing_keys, results[: len(listing_keys)], strict=True): + if isinstance(listing, SchemaColumnListing): + self._column_listings[schema_key] = listing + exact: RelationInfo + exact_columns: SchemaColumnListing | tuple[ColumnInfo, ...] + for exact, exact_columns in zip(exact_relations, results[len(listing_keys) :], strict=True): + if not isinstance(exact_columns, SchemaColumnListing): + self._relation_columns[exact.identity] = exact_columns + for schema_key, scoped in by_schema.items(): + column_listing: SchemaColumnListing | None = self._column_listings.get(schema_key) + if column_listing is None: + continue + for relation in scoped: + self._relation_columns[relation.identity] = ( + column_listing.columns_by_stored_name.get( + reader.metadata_name_key(relation.name) + ) + ) + + def _bind_column_listing( + self, *, reader: SchemaScopedMetadataReader, database: str | None, schema: str + ) -> Callable[[], SchemaColumnListing]: + def read_listing() -> SchemaColumnListing: + return reader.read_schema_column_listing( + connection=self.connection, database=database, schema=schema + ) + + return read_listing + + def _bind_relation_columns( + self, *, reader: SchemaScopedMetadataReader, relation: RelationInfo + ) -> Callable[[], tuple[ColumnInfo, ...]]: + def read_columns() -> tuple[ColumnInfo, ...]: + return reader.read_relation_columns(connection=self.connection, relation=relation) + + return read_columns + + def _schema_key(self, *, database: str | None, schema: str) -> _SchemaKey: + if self._reader is None: + return (database, schema) + return ( + None if database is None else self._reader.metadata_name_key(database), + self._reader.metadata_name_key(schema), + ) + + +def _flatten_relations(chunks: list[tuple[RelationInfo, ...]]) -> tuple[RelationInfo, ...]: + relations: list[RelationInfo] = [] + chunk: tuple[RelationInfo, ...] + for chunk in chunks: + relations.extend(chunk) + return tuple(relations) + + +def _describe(values: tuple[str, ...] | None) -> str: + if values is None: + return "*" + if len(values) <= INSPECTION_DESCRIBED_NAME_LIMIT: + return "[" + ", ".join(values) + "]" + return f"[{len(values)} names]" diff --git a/src/sqlbuild/adapter/relations/constants.py b/src/sqlbuild/adapter/relations/constants.py index e55b715814..c0d775459f 100644 --- a/src/sqlbuild/adapter/relations/constants.py +++ b/src/sqlbuild/adapter/relations/constants.py @@ -1,4 +1,6 @@ """Shared relation metadata limits.""" +INSPECTION_DESCRIBED_NAME_LIMIT: int = 3 +INSPECTION_IN_LIST_LIMIT: int = 50 METADATA_NAME_FILTER_LIMIT: int = 250 SHORTENED_LOGICAL_NAME_HASH_LENGTH: int = 8 diff --git a/src/sqlbuild/adapter/relations/main/active_inspection_catalog.py b/src/sqlbuild/adapter/relations/main/active_inspection_catalog.py new file mode 100644 index 0000000000..df8553843a --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/active_inspection_catalog.py @@ -0,0 +1,14 @@ +"""Look up the inspection catalog open for one adapter connection.""" + +from __future__ import annotations + +from typing import Any + +from sqlbuild.adapter.relations._helpers.inspection_context import current_catalog +from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog + + +def active_inspection_catalog(*, adapter: Any, connection: Any) -> InspectionCatalog | None: + """Return the open catalog for exactly this adapter and connection, if any.""" + + return current_catalog(adapter=adapter, connection=connection) diff --git a/src/sqlbuild/adapter/relations/main/get_columns_for_inspection.py b/src/sqlbuild/adapter/relations/main/get_columns_for_inspection.py new file mode 100644 index 0000000000..7e904dea4a --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/get_columns_for_inspection.py @@ -0,0 +1,20 @@ +"""Read relation columns through the planning inspection catalog when one is open.""" + +from __future__ import annotations + +from typing import Any + +from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo +from sqlbuild.adapter.relations._helpers.inspection_context import current_catalog +from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog + + +def get_columns_for_inspection( + *, adapter: Any, connection: Any, relations: tuple[RelationInfo, ...] +) -> dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]]: + """Read relation columns through the open catalog, or directly when none is open.""" + + catalog: InspectionCatalog | None = current_catalog(adapter=adapter, connection=connection) + if catalog is None: + return adapter.get_columns_for_relations(connection=connection, relations=relations) + return catalog.get_columns_for_relations(relations=relations) diff --git a/src/sqlbuild/adapter/relations/main/inspection_query_sink.py b/src/sqlbuild/adapter/relations/main/inspection_query_sink.py new file mode 100644 index 0000000000..39b16b83ae --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/inspection_query_sink.py @@ -0,0 +1,24 @@ +"""Deliver recorded inspection reads to a verbose output sink.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from contextvars import Token + +from sqlbuild.adapter.relations._helpers.inspection_context import ( + activate_query_sink, + deactivate_query_sink, +) +from sqlbuild.adapter.relations.models import InspectionQueryRecord + + +@contextmanager +def inspection_query_sink(sink: Callable[[InspectionQueryRecord], None]) -> Iterator[None]: + """Deliver every inspection read recorded in this context to ``sink``.""" + + token: Token[Callable[[InspectionQueryRecord], None] | None] = activate_query_sink(sink) + try: + yield + finally: + _ = deactivate_query_sink(token) diff --git a/src/sqlbuild/adapter/relations/main/list_relations_for_inspection.py b/src/sqlbuild/adapter/relations/main/list_relations_for_inspection.py new file mode 100644 index 0000000000..7ada506c6d --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/list_relations_for_inspection.py @@ -0,0 +1,27 @@ +"""List relations through the planning inspection catalog when one is open.""" + +from __future__ import annotations + +from typing import Any + +from sqlbuild.adapter.contract.models import RelationInfo +from sqlbuild.adapter.relations._helpers.inspection_context import current_catalog +from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog + + +def list_relations_for_inspection( + *, + adapter: Any, + connection: Any, + database: str | None, + schemas: tuple[str, ...] | None, + names: tuple[str, ...] | None = None, +) -> tuple[RelationInfo, ...]: + """List relations through the open catalog, or directly when none is open.""" + + catalog: InspectionCatalog | None = current_catalog(adapter=adapter, connection=connection) + if catalog is None: + return adapter.list_relations( + connection=connection, database=database, schemas=schemas, names=names + ) + return catalog.list_relations(database=database, schemas=schemas, names=names) diff --git a/src/sqlbuild/adapter/relations/main/open_inspection_catalog.py b/src/sqlbuild/adapter/relations/main/open_inspection_catalog.py new file mode 100644 index 0000000000..ea5d4db17c --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/open_inspection_catalog.py @@ -0,0 +1,31 @@ +"""Scope one warehouse inspection catalog to a planning invocation.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import Token +from typing import Any + +from sqlbuild.adapter.relations._helpers.inspection_context import ( + activate_catalog, + current_catalog, + deactivate_catalog, +) +from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog + + +@contextmanager +def open_inspection_catalog(*, adapter: Any, connection: Any) -> Iterator[InspectionCatalog]: + """Share metadata reads on this connection until the context exits; nested opens reuse it.""" + + existing: InspectionCatalog | None = current_catalog(adapter=adapter, connection=connection) + if existing is not None: + yield existing + return + catalog: InspectionCatalog = InspectionCatalog(adapter=adapter, connection=connection) + token: Token[InspectionCatalog | None] = activate_catalog(catalog) + try: + yield catalog + finally: + _ = deactivate_catalog(token) diff --git a/src/sqlbuild/adapter/relations/main/record_inspection_query.py b/src/sqlbuild/adapter/relations/main/record_inspection_query.py new file mode 100644 index 0000000000..3a39884fd7 --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/record_inspection_query.py @@ -0,0 +1,12 @@ +"""Record one planner warehouse inspection read.""" + +from __future__ import annotations + +from sqlbuild.adapter.relations._helpers.inspection_context import emit_inspection_record +from sqlbuild.adapter.relations.models import InspectionQueryRecord + + +def record_inspection_query(*, record: InspectionQueryRecord) -> None: + """Log one inspection read and forward it to the active verbose sink.""" + + return emit_inspection_record(record=record) diff --git a/src/sqlbuild/adapter/relations/main/render_inspection_sql.py b/src/sqlbuild/adapter/relations/main/render_inspection_sql.py new file mode 100644 index 0000000000..490dc644bd --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/render_inspection_sql.py @@ -0,0 +1,18 @@ +"""Render parameterized inspection SQL as readable diagnostic text.""" + +from __future__ import annotations + + +def render_inspection_sql(*, query: str, params: tuple[object, ...]) -> str: + """Inline ``%s`` parameters as quoted literals for display only; never execute the result.""" + + parts: list[str] = query.split("%s") + if len(parts) != len(params) + 1: + return query + rendered: list[str] = [parts[0]] + value: object + part: str + for value, part in zip(params, parts[1:], strict=True): + rendered.append("'" + str(value).replace("'", "''") + "'") + rendered.append(part) + return "".join(rendered) diff --git a/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py b/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py new file mode 100644 index 0000000000..7389ce64ad --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py @@ -0,0 +1,55 @@ +"""Run independent warehouse inspection reads with bounded concurrency.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from concurrent.futures import Future, ThreadPoolExecutor, as_completed +from contextvars import copy_context +from typing import cast + +from sqlbuild.adapter.relations.types import InspectionCompletion + + +def run_bounded_inspections[ResultT]( + *, + tasks: Sequence[Callable[[], ResultT]], + concurrency: int, + on_complete: InspectionCompletion[ResultT] | None = None, +) -> list[ResultT]: + """Return results in input order, report completions on this thread, raise lowest failure.""" + + if concurrency <= 1 or len(tasks) <= 1: + results: list[ResultT] = [] + index: int + task: Callable[[], ResultT] + for index, task in enumerate(tasks): + result: ResultT = task() + if on_complete is not None: + on_complete(index=index, result=result) + results.append(result) + return results + completed: dict[int, ResultT] = {} + failures: dict[int, BaseException] = {} + with ThreadPoolExecutor(max_workers=min(concurrency, len(tasks))) as pool: + futures: dict[Future[ResultT], int] = { + cast(Future[ResultT], pool.submit(copy_context().run, task)): index + for index, task in enumerate(tasks) + } + future: Future[ResultT] + for future in as_completed(futures): + if future.cancelled(): + continue + future_index: int = futures[future] + error: BaseException | None = future.exception() + if error is not None: + failures[future_index] = error + pending: Future[ResultT] + for pending in futures: + _ = pending.cancel() + continue + completed[future_index] = future.result() + if on_complete is not None and not failures: + on_complete(index=future_index, result=completed[future_index]) + if failures: + raise failures[min(failures)] + return [completed[index] for index in range(len(tasks))] diff --git a/src/sqlbuild/adapter/relations/main/run_recorded_inspection_query.py b/src/sqlbuild/adapter/relations/main/run_recorded_inspection_query.py new file mode 100644 index 0000000000..026d0a9c7b --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/run_recorded_inspection_query.py @@ -0,0 +1,34 @@ +"""Run and record one planner warehouse inspection read.""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any + +from sqlbuild.adapter.relations._helpers.inspection_context import emit_inspection_record +from sqlbuild.adapter.relations.models import InspectionQueryRecord + + +def run_recorded_inspection_query(*, sql: str, run: Callable[[], list[Any]]) -> list[Any]: + """Run one inspection read, then record its text, elapsed time, and row count.""" + + started: float = time.monotonic() + try: + rows: list[Any] = run() + except Exception as error: + _ = emit_inspection_record( + record=InspectionQueryRecord( + sql=sql, + elapsed_seconds=time.monotonic() - started, + row_count=None, + error=str(error), + ) + ) + raise + _ = emit_inspection_record( + record=InspectionQueryRecord( + sql=sql, elapsed_seconds=time.monotonic() - started, row_count=len(rows) + ) + ) + return rows diff --git a/src/sqlbuild/adapter/relations/models.py b/src/sqlbuild/adapter/relations/models.py new file mode 100644 index 0000000000..535f003c3c --- /dev/null +++ b/src/sqlbuild/adapter/relations/models.py @@ -0,0 +1,44 @@ +"""Invocation-scoped warehouse inspection records.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo + + +@dataclass(frozen=True) +class ListedRelation: + """One relation from a schema listing with the exact name stored by the warehouse.""" + + stored_name: str + relation: RelationInfo + + +@dataclass(frozen=True) +class SchemaRelationListing: + """Every relation one database schema exposes to the inspecting role.""" + + database: str | None + schema: str + entries: tuple[ListedRelation, ...] + rows: tuple[tuple[object, ...], ...] = () + + +@dataclass(frozen=True) +class SchemaColumnListing: + """Every column one database schema exposes, keyed by stored relation name.""" + + database: str | None + schema: str + columns_by_stored_name: dict[str, tuple[ColumnInfo, ...]] + + +@dataclass(frozen=True) +class InspectionQueryRecord: + """One warehouse metadata or cursor-bound read issued while planning.""" + + sql: str + elapsed_seconds: float + row_count: int | None + error: str | None = None diff --git a/src/sqlbuild/adapter/relations/types.py b/src/sqlbuild/adapter/relations/types.py index 28fd8713bf..bb9135356c 100644 --- a/src/sqlbuild/adapter/relations/types.py +++ b/src/sqlbuild/adapter/relations/types.py @@ -1,6 +1,12 @@ -"""Adapter relation naming contracts.""" +"""Adapter relation naming and inspection contracts.""" -from typing import Protocol +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +if TYPE_CHECKING: + from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo + from sqlbuild.adapter.relations.models import SchemaColumnListing, SchemaRelationListing class RelationLocation(Protocol): @@ -10,3 +16,40 @@ class RelationLocation(Protocol): schema: str | None name: str qualified_name: str | None + + +@runtime_checkable +class SchemaScopedMetadataReader(Protocol): + """Adapter able to list a whole schema's relations and columns in one read each.""" + + exact_column_inspection_limit: int + + def read_schema_relation_listing( + self, *, connection: Any, database: str | None, schema: str + ) -> SchemaRelationListing: + """Read every relation in one schema.""" + ... + + def read_schema_column_listing( + self, *, connection: Any, database: str | None, schema: str + ) -> SchemaColumnListing: + """Read every column in one schema.""" + ... + + def read_relation_columns( + self, *, connection: Any, relation: RelationInfo + ) -> tuple[ColumnInfo, ...]: + """Read one relation's columns.""" + ... + + def metadata_name_key(self, name: str) -> str: + """Return the stored metadata name a logical identifier filter matches.""" + ... + + +class InspectionCompletion[ResultT](Protocol): + """Callback receiving one finished inspection read on the calling thread.""" + + def __call__(self, *, index: int, result: ResultT) -> None: + """Handle the result of the task at ``index``.""" + ... diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index 7d1f9964d0..b0ccbf6f6d 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -80,7 +80,18 @@ StatementSizeLimit, TablePromotionMode, ) +from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog +from sqlbuild.adapter.relations.main.active_inspection_catalog import active_inspection_catalog from sqlbuild.adapter.relations.main.relation_age_timestamp import relation_age_timestamp_utc +from sqlbuild.adapter.relations.main.render_inspection_sql import render_inspection_sql +from sqlbuild.adapter.relations.main.run_recorded_inspection_query import ( + run_recorded_inspection_query, +) +from sqlbuild.adapter.relations.models import ( + ListedRelation, + SchemaColumnListing, + SchemaRelationListing, +) from sqlbuild.adapter.state_sql.main.render_insert_source_freshness_records_sql import ( render_insert_source_freshness_records_sql, ) @@ -124,6 +135,10 @@ from sqlbuild.sql_values.models import SqlValue _EXACT_COLUMN_INSPECTION_LIMIT: int = 32 +_METADATA_INSPECTION_CONCURRENCY: int = 8 +_RELATION_LISTING_COLUMNS: str = ( + "table_name, table_schema, table_type, is_transient, created, last_altered, retention_time" +) _BULK_COLUMN_RELATION_CHUNK_SIZE: int = 2_000 @@ -133,6 +148,8 @@ class SnowflakeAdapter(MicrobatchMixin, UnkeyedDiffMixin, BaseAdapter): adapter_name: ClassVar[str] = BuiltinAdapter.SNOWFLAKE.value sql_analysis_dialect_name: ClassVar[str | None] = "snowflake" execution_duration_limit_seconds: ClassVar[int | None] = MAX_STATEMENT_TIMEOUT_SECONDS + metadata_inspection_concurrency: ClassVar[int] = _METADATA_INSPECTION_CONCURRENCY + exact_column_inspection_limit: ClassVar[int] = _EXACT_COLUMN_INSPECTION_LIMIT max_identifier_length: ClassVar[int] = 255 state_tables_transient: ClassVar[bool] = True @@ -357,6 +374,64 @@ def _information_schema_relation(self, *, database: str | None, name: str) -> st return f"information_schema.{name}" return f"{self.render_identifier(database)}.information_schema.{name}" + def _freshness_rows( + self, + *, + connection: Any, + database: str | None, + schema: str | None, + scoped_requests: list[TableFreshnessRequest], + ) -> list[tuple[Any, ...]]: + """Reuse the planning catalog's schema listing, or read the requested tables.""" + + catalog: InspectionCatalog | None = active_inspection_catalog( + adapter=self, connection=connection + ) + if catalog is not None and schema is not None: + catalog.prefetch_relation_listings( + database=scoped_requests[0].database, schemas=(schema,) + ) + listing: SchemaRelationListing | None = catalog.cached_relation_listing( + database=scoped_requests[0].database, schema=schema + ) + if listing is not None: + wanted: frozenset[str] = frozenset( + self._information_schema_identifier(request.name) for request in scoped_requests + ) + return [ + (database, row[1], row[0], row[2], row[5]) + for row in listing.rows + if str(row[0]) in wanted + ] + clauses: list[str] = [] + params: list[str] = [] + if database is not None: + clauses.append("table_catalog = %s") + params.append(database) + if schema is not None: + clauses.append("table_schema = %s") + params.append(schema) + placeholders: str = ", ".join(["%s"] * len(scoped_requests)) + clauses.append(f"table_name IN ({placeholders})") + params.extend( + self._information_schema_identifier(request.name) for request in scoped_requests + ) + cursor: Any = connection.cursor() + try: + cursor.execute( + "SELECT table_catalog, table_schema, table_name, table_type, last_altered " + "FROM " + + self._information_schema_relation( + database=scoped_requests[0].database, name="tables" + ) + + " WHERE " + + " AND ".join(clauses), + tuple(params), + ) + return list(cursor.fetchall()) + finally: + cursor.close() + @staticmethod def _freshness_request_key( *, database: str | None, schema: str | None, name: str @@ -703,34 +778,12 @@ def get_tables_freshness_metadata( ) requests_by_scope.setdefault(scope, []).append(request) for (database, schema), scoped_requests in requests_by_scope.items(): - clauses: list[str] = [] - params: list[str] = [] - if database is not None: - clauses.append("table_catalog = %s") - params.append(database) - if schema is not None: - clauses.append("table_schema = %s") - params.append(schema) - placeholders: str = ", ".join(["%s"] * len(scoped_requests)) - clauses.append(f"table_name IN ({placeholders})") - params.extend( - self._information_schema_identifier(request.name) for request in scoped_requests + rows: list[tuple[Any, ...]] = self._freshness_rows( + connection=connection, + database=database, + schema=schema, + scoped_requests=scoped_requests, ) - cursor: Any = connection.cursor() - try: - cursor.execute( - "SELECT table_catalog, table_schema, table_name, table_type, last_altered " - "FROM " - + self._information_schema_relation( - database=scoped_requests[0].database, name="tables" - ) - + " WHERE " - + " AND ".join(clauses), - tuple(params), - ) - rows: list[tuple[Any, ...]] = list(cursor.fetchall()) - finally: - cursor.close() row: tuple[Any, ...] for row in rows: matched_request: TableFreshnessRequest | None = ( @@ -1701,8 +1754,7 @@ def list_relations( names: tuple[str, ...] | None = None, ) -> tuple[Any, ...]: query: str = ( - "SELECT table_name, table_schema, table_type, is_transient, created, last_altered, " - "retention_time " + f"SELECT {_RELATION_LISTING_COLUMNS} " f"FROM {self._information_schema_relation(database=database, name='tables')} WHERE 1=1" ) params: list[str] = [] @@ -1723,22 +1775,116 @@ def list_relations( rows: list[tuple[Any, ...]] = cursor.fetchall() finally: cursor.close() - from sqlbuild.adapter.contract.models import RelationInfo + return tuple(self._relation_info_from_row(row=row, database=database) for row in rows) - return tuple( - RelationInfo( - database=None if row[1] is None else database, - schema=None if row[1] is None else str(row[1]).lower(), - name=str(row[0]).lower(), - relation_type=str(row[2]).lower(), - is_transient=( - None if row[3] is None else str(row[3]).upper() == TRUE_METADATA_VALUE - ), - created_at=relation_age_timestamp_utc(row[4]), - last_altered_at=relation_age_timestamp_utc(row[5]), - retention_days=None if row[6] is None else int(row[6]), + @staticmethod + def _relation_info_from_row(*, row: tuple[Any, ...], database: str | None) -> RelationInfo: + return RelationInfo( + database=None if row[1] is None else database, + schema=None if row[1] is None else str(row[1]).lower(), + name=str(row[0]).lower(), + relation_type=str(row[2]).lower(), + is_transient=(None if row[3] is None else str(row[3]).upper() == TRUE_METADATA_VALUE), + created_at=relation_age_timestamp_utc(row[4]), + last_altered_at=relation_age_timestamp_utc(row[5]), + retention_days=None if row[6] is None else int(row[6]), + ) + + def read_schema_relation_listing( + self, *, connection: _SnowflakeConnection, database: str | None, schema: str + ) -> SchemaRelationListing: + """List every relation in one schema with a single unfiltered metadata read.""" + + query: str = ( + f"SELECT {_RELATION_LISTING_COLUMNS} " + f"FROM {self._information_schema_relation(database=database, name='tables')} " + "WHERE table_schema = %s" + ) + params: list[str] = [self._information_schema_identifier(schema)] + if database is not None: + query += " AND table_catalog = %s" + params.append(self._information_schema_identifier(database)) + rows: list[tuple[Any, ...]] = self._fetch_inspection_rows( + connection=connection, query=query, params=tuple(params) + ) + return SchemaRelationListing( + database=database, + schema=schema, + entries=tuple( + ListedRelation( + stored_name=str(row[0]), + relation=self._relation_info_from_row(row=row, database=database), + ) + for row in rows + ), + rows=tuple(tuple(row) for row in rows), + ) + + def read_schema_column_listing( + self, *, connection: _SnowflakeConnection, database: str | None, schema: str + ) -> SchemaColumnListing: + """List every column in one schema with a single unfiltered metadata read.""" + + query: str = ( + "SELECT table_name, column_name, data_type, numeric_precision, numeric_scale, " + "character_maximum_length FROM " + + self._information_schema_relation(database=database, name="columns") + + " WHERE table_schema = %s" + ) + params: list[str] = [self._information_schema_identifier(schema)] + if database is not None: + query += " AND table_catalog = %s" + params.append(self._information_schema_identifier(database)) + query += " ORDER BY table_name, ordinal_position" + rows: list[tuple[Any, ...]] = self._fetch_inspection_rows( + connection=connection, query=query, params=tuple(params) + ) + columns: dict[str, list[ColumnInfo]] = {} + row: tuple[Any, ...] + for row in rows: + columns.setdefault(str(row[0]), []).append( + ColumnInfo( + name=str(row[1]).lower(), + type=self._build_information_schema_type( + data_type=str(row[2]), + numeric_precision=row[3], + numeric_scale=row[4], + character_maximum_length=row[5], + ), + ) ) - for row in rows + return SchemaColumnListing( + database=database, + schema=schema, + columns_by_stored_name={name: tuple(value) for name, value in columns.items()}, + ) + + def read_relation_columns( + self, *, connection: _SnowflakeConnection, relation: RelationInfo + ) -> tuple[ColumnInfo, ...]: + """Read one relation's columns with an exact SHOW COLUMNS call.""" + + return self._show_columns_for_relation(connection=connection, relation=relation) + + def metadata_name_key(self, name: str) -> str: + """Unquoted logical identifiers match the uppercase name Snowflake stores.""" + + return self._information_schema_identifier(name) + + @staticmethod + def _fetch_inspection_rows( + *, connection: _SnowflakeConnection, query: str, params: tuple[str, ...] + ) -> list[tuple[Any, ...]]: + def fetch() -> list[tuple[Any, ...]]: + cursor: Any = connection.cursor() + try: + cursor.execute(query, params) + return list(cursor.fetchall()) + finally: + cursor.close() + + return run_recorded_inspection_query( + sql=render_inspection_sql(query=query, params=params), run=fetch ) def with_relation_age_metadata( @@ -1924,12 +2070,17 @@ def _show_columns_for_relation( relation_kind: str = ( "VIEW" if VIEW_RELATION_TYPE_TOKEN in relation.relation_type.lower() else "TABLE" ) - cursor: Any = connection.cursor() - try: - cursor.execute(f"SHOW COLUMNS IN {relation_kind} {qualified_name}") - rows: list[tuple[Any, ...]] = list(cursor.fetchall()) - finally: - cursor.close() + query: str = f"SHOW COLUMNS IN {relation_kind} {qualified_name}" + + def fetch() -> list[tuple[Any, ...]]: + cursor: Any = connection.cursor() + try: + cursor.execute(query) + return list(cursor.fetchall()) + finally: + cursor.close() + + rows: list[tuple[Any, ...]] = run_recorded_inspection_query(sql=query, run=fetch) return tuple( ColumnInfo( name=str(row[2]).lower(), diff --git a/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/_test_types.py b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/_test_types.py new file mode 100644 index 0000000000..5b128cf7b8 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/_test_types.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class BoundedInspectionTestCase: + """Independent reads run with a concurrency bound.""" + + description: str + task_count: int + concurrency: int + expected_results: tuple[int, ...] + expected_parallel: bool + + +@dataclass(frozen=True) +class BoundedInspectionFailureTestCase: + """Independent reads where some fail.""" + + description: str + task_count: int + concurrency: int + failing_indexes: frozenset[int] + expected_error: str diff --git a/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/helpers.py b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/helpers.py new file mode 100644 index 0000000000..4c28e099fb --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/helpers.py @@ -0,0 +1,47 @@ +"""Helpers for bounded inspection read tests.""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Callable + + +def _succeed(index: int) -> int: + return index * 10 + + +def _fail(index: int) -> int: + raise RuntimeError(f"read {index} failed") + + +_OUTCOMES: dict[bool, Callable[[int], int]] = {False: _succeed, True: _fail} +_LATENCIES: tuple[float, float] = (0.02, 0.005) + + +class ConcurrencyTracker: + """Build read tasks that record how many run at once.""" + + def __init__(self) -> None: + self._lock: threading.Lock = threading.Lock() + self._active: int = 0 + self.max_active: int = 0 + + def tasks(self, *, count: int, failing: frozenset[int]) -> tuple[Callable[[], int], ...]: + """Return ``count`` tasks; even tasks are slower so finish order differs from input.""" + + return tuple(self._task(index=index, failing=failing) for index in range(count)) + + def _task(self, *, index: int, failing: frozenset[int]) -> Callable[[], int]: + def run() -> int: + with self._lock: + self._active += 1 + self.max_active = max(self.max_active, self._active) + try: + time.sleep(_LATENCIES[index % 2]) + return _OUTCOMES[index in failing](index) + finally: + with self._lock: + self._active -= 1 + + return run diff --git a/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/test_run_bounded_inspections.py b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/test_run_bounded_inspections.py new file mode 100644 index 0000000000..3577b1ddc0 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/test_run_bounded_inspections.py @@ -0,0 +1,93 @@ +"""Bounded inspection reads keep input order and fail deterministically.""" + +from __future__ import annotations + +import threading + +import pytest + +from sqlbuild.adapter.relations.main.run_bounded_inspections import run_bounded_inspections +from tests.unit.src.sqlbuild.adapter.relations.main.run_bounded_inspections._test_types import ( + BoundedInspectionFailureTestCase, + BoundedInspectionTestCase, +) +from tests.unit.src.sqlbuild.adapter.relations.main.run_bounded_inspections.helpers import ( + ConcurrencyTracker, +) + + +@pytest.mark.parametrize( + "test_case", + [ + BoundedInspectionTestCase( + description="sequential", + task_count=5, + concurrency=1, + expected_results=(0, 10, 20, 30, 40), + expected_parallel=False, + ), + BoundedInspectionTestCase( + description="bounded parallel", + task_count=20, + concurrency=4, + expected_results=tuple(index * 10 for index in range(20)), + expected_parallel=True, + ), + ], + ids=lambda case: case.description, +) +def test_given_independent_reads_when_running_then_returns_results_in_input_order( + test_case: BoundedInspectionTestCase, +) -> None: + tracker: ConcurrencyTracker = ConcurrencyTracker() + completions: list[tuple[int, int, int]] = [] + + results: list[int] = run_bounded_inspections( + tasks=tracker.tasks(count=test_case.task_count, failing=frozenset()), + concurrency=test_case.concurrency, + on_complete=lambda index, result: completions.append( + (index, result, threading.get_ident()) + ), + ) + + assert tuple(results) == test_case.expected_results + assert tracker.max_active <= test_case.concurrency + assert (tracker.max_active > 1) is test_case.expected_parallel + assert sorted(index for index, _, _ in completions) == list(range(test_case.task_count)) + assert {thread for _, _, thread in completions} == {threading.get_ident()} + + +@pytest.mark.parametrize( + "test_case", + [ + BoundedInspectionFailureTestCase( + description="lowest failing index wins regardless of finish order", + task_count=12, + concurrency=6, + failing_indexes=frozenset({3, 4}), + expected_error="read 3 failed", + ), + BoundedInspectionFailureTestCase( + description="sequential failure stops at the failing read", + task_count=6, + concurrency=1, + failing_indexes=frozenset({2}), + expected_error="read 2 failed", + ), + ], + ids=lambda case: case.description, +) +def test_given_failing_read_when_running_then_raises_lowest_index_failure( + test_case: BoundedInspectionFailureTestCase, +) -> None: + tracker: ConcurrencyTracker = ConcurrencyTracker() + + with pytest.raises(RuntimeError, match=test_case.expected_error): + _ = run_bounded_inspections( + tasks=tracker.tasks(count=test_case.task_count, failing=test_case.failing_indexes), + concurrency=test_case.concurrency, + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/__init__.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py new file mode 100644 index 0000000000..d6e365e2ff --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class RelationRequestEquivalenceTestCase: + """One list_relations request answered both directly and through the catalog.""" + + description: str + database: str | None + schemas: tuple[str, ...] | None + names: tuple[str, ...] | None + expected_relation_count: int + + +@dataclass(frozen=True) +class ColumnRequestEquivalenceTestCase: + """One column request whose catalog answer must equal the direct adapter answer.""" + + description: str + schemas: tuple[str, ...] + names: tuple[str, ...] | None + expected_query_kinds: tuple[str, ...] + + +@dataclass(frozen=True) +class FreshnessReuseTestCase: + """Source freshness requests served from schema listings already read by the catalog.""" + + description: str + requests: tuple[tuple[str, str, str], ...] + listed_schemas: tuple[str, ...] + expected_metadata_reads: int + + +@dataclass(frozen=True) +class FreshnessErrorTestCase: + """A freshness request that must fail identically with and without the catalog.""" + + description: str + request: tuple[str, str, str] + expected_error_fragment: str + + +@dataclass(frozen=True) +class RepeatedLookupTestCase: + """Several filtered lookups of one schema answered from one listing.""" + + description: str + lookups: tuple[tuple[str, ...] | None, ...] + expected_relation_counts: tuple[int, ...] + expected_query_kinds: tuple[str, ...] + + +@dataclass(frozen=True) +class SpeculativePrefetchTestCase: + """Best-effort listing of a schema whose read fails.""" + + description: str + failing_schema: str + expected_error_fragment: str + expected_failed_reads: int diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py new file mode 100644 index 0000000000..e9f55799a3 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -0,0 +1,648 @@ +"""Recording Snowflake warehouse double for inspection round-trip tests. + +The double answers ``INFORMATION_SCHEMA`` reads by evaluating the issued SQL against an in-memory +DuckDB copy of synthetic metadata, so filtering, ordering, and case semantics come from the SQL the +adapter actually sends rather than from hand-written fakes. +""" + +from __future__ import annotations + +import json +import re +import threading +import time +from collections.abc import Callable, Iterator +from dataclasses import dataclass, field +from datetime import UTC, datetime +from itertools import compress +from pathlib import Path +from typing import Any + +import duckdb +import pyarrow + +from sqlbuild.adapter.contract.models import RelationInfo +from sqlbuild.adapters.snowflake.classes.snowflake_adapter import SnowflakeAdapter +from sqlbuild.adapters.snowflake.classes.snowflake_connection import _SnowflakeConnection + +_METADATA_RELATION_PATTERN: re.Pattern[str] = re.compile( + r'(?:"(?P[^"]*)"\.)?information_schema\.(?Ptables|columns)\b', + re.IGNORECASE, +) +_OTHER_INFORMATION_SCHEMA_PATTERN: re.Pattern[str] = re.compile( + r"information_schema\.", re.IGNORECASE +) +_SHOW_COLUMNS_PATTERN: re.Pattern[str] = re.compile( + r"^\s*SHOW\s+COLUMNS\s+IN\s+(?:TABLE|VIEW)\s+(?P\S+)\s*$", re.IGNORECASE +) +_SHOW_PATTERN: re.Pattern[str] = re.compile(r"^\s*SHOW\s", re.IGNORECASE) +_CURSOR_BOUND_PATTERN: re.Pattern[str] = re.compile( + r"^SELECT\s+CAST\((?:MIN|MAX)\(.*\)\s+AS\s+VARCHAR\).*\s+FROM\s+(?P\S+)\s*$", + re.IGNORECASE | re.DOTALL, +) +_ANY_PATTERN: re.Pattern[str] = re.compile(r"") +_IN_LIST_PATTERN: re.Pattern[str] = re.compile(r"\bIN\s*\(([^()]*)\)", re.IGNORECASE) +_TABLES_PATTERN: re.Pattern[str] = re.compile(r"information_schema\.tables\b", re.IGNORECASE) +_COLUMNS_PATTERN: re.Pattern[str] = re.compile(r"information_schema\.columns\b", re.IGNORECASE) +_QUERY_KINDS: tuple[tuple[re.Pattern[str], str], ...] = ( + (_TABLES_PATTERN, "tables"), + (_COLUMNS_PATTERN, "columns"), + (_SHOW_COLUMNS_PATTERN, "show_columns"), + (_CURSOR_BOUND_PATTERN, "cursor_bounds"), + (_OTHER_INFORMATION_SCHEMA_PATTERN, "other_metadata"), + (_SHOW_PATTERN, "other_metadata"), + (_ANY_PATTERN, "data"), +) +_FIXED_CREATED_AT: datetime = datetime(2026, 1, 1, tzinfo=UTC) +_SHOW_TYPE_NAMES: dict[str, str] = {"NUMBER": "FIXED", "FLOAT": "REAL"} +_STATUS_DESCRIPTION: tuple[tuple[str], ...] = (("status",),) +_SHOW_DESCRIPTION: tuple[tuple[str], ...] = ( + ("table_name",), + ("schema_name",), + ("column_name",), + ("data_type",), +) +_BOUND_MARKERS: tuple[str, str] = ("MIN(", "MAX(") + + +@dataclass(frozen=True) +class FakeColumn: + """One synthetic Snowflake column with INFORMATION_SCHEMA type fields.""" + + name: str + data_type: str = "TEXT" + numeric_precision: int | None = None + numeric_scale: int | None = None + character_maximum_length: int | None = 16_777_216 + + +@dataclass(frozen=True) +class FakeRelation: + """One synthetic relation with its stored (case-exact) Snowflake name.""" + + database: str + schema: str + name: str + table_type: str = "BASE TABLE" + is_transient: bool = False + columns: tuple[FakeColumn, ...] = (FakeColumn(name="ID", data_type="NUMBER"),) + retention_time: int = 1 + + +@dataclass(frozen=True) +class RecordedQuery: + """One statement issued against the double.""" + + sql: str + params: tuple[object, ...] + thread_id: int + started_at: float + finished_at: float + row_count: int + + @property + def kind(self) -> str: + """Classify the statement for round-trip budgets.""" + + return next(filter(lambda rule: rule[0].search(self.sql) is not None, _QUERY_KINDS))[1] + + @property + def largest_in_list(self) -> int: + """Return the largest literal or placeholder IN list in the statement.""" + + return max( + (len(match.group(1).split(",")) for match in _IN_LIST_PATTERN.finditer(self.sql)), + default=0, + ) + + +@dataclass +class RecordingSnowflakeWarehouse: + """Thread-safe in-memory Snowflake metadata and cursor-bound responder.""" + + relations: tuple[FakeRelation, ...] + statement_latency_seconds: float = 0.0 + failing_relations: frozenset[str] = frozenset() + failing_metadata_schemas: frozenset[str] = frozenset() + cursor_values: tuple[str, str] = ("2026-01-01 00:00:00.000", "2026-01-31 00:00:00.000") + queries: list[RecordedQuery] = field(default_factory=list) + attempted_sql: list[str] = field(default_factory=list) + _lock: threading.Lock = field(default_factory=threading.Lock) + _active: dict[bool, int] = field(default_factory=lambda: {True: 0, False: 0}) + _peak: dict[bool, int] = field(default_factory=lambda: {True: 0, False: 0}) + + def __post_init__(self) -> None: + self._database: duckdb.DuckDBPyConnection = duckdb.connect(":memory:") + self._database.execute( + "CREATE TABLE fake_tables (table_catalog VARCHAR, table_schema VARCHAR, " + "table_name VARCHAR, table_type VARCHAR, is_transient VARCHAR, " + "created TIMESTAMPTZ, last_altered TIMESTAMPTZ, retention_time INTEGER)" + ) + self._database.execute( + "CREATE TABLE fake_columns (table_catalog VARCHAR, table_schema VARCHAR, " + "table_name VARCHAR, column_name VARCHAR, ordinal_position INTEGER, " + "data_type VARCHAR, numeric_precision INTEGER, numeric_scale INTEGER, " + "character_maximum_length INTEGER)" + ) + self._load(table="fake_tables", rows=[_table_row(relation) for relation in self.relations]) + column_rows: list[dict[str, object]] = [] + relation: FakeRelation + for relation in self.relations: + column_rows.extend(_column_rows(relation)) + self._load(table="fake_columns", rows=column_rows) + + @property + def max_concurrent_cursor_bounds(self) -> int: + """Return the highest number of cursor-bound statements in flight at once.""" + + return self._peak[True] + + @property + def max_concurrent_metadata(self) -> int: + """Return the highest number of other statements in flight at once.""" + + return self._peak[False] + + def raw_connection(self) -> FakeRawSnowflakeConnection: + """Return a raw DB-API connection double bound to this warehouse.""" + + return FakeRawSnowflakeConnection(warehouse=self) + + def queries_of_kind(self, kind: str) -> tuple[RecordedQuery, ...]: + """Return recorded statements of one classification.""" + + with self._lock: + return tuple(filter(lambda query: query.kind == kind, self.queries)) + + def reset(self) -> None: + """Forget recorded statements and concurrency peaks.""" + + with self._lock: + self.queries.clear() + self.attempted_sql.clear() + self._peak.update({True: 0, False: 0}) + + def run(self, *, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[Any, ...]], Any]: + """Execute one statement and record it.""" + + started_at: float = time.monotonic() + is_cursor_bound: bool = _CURSOR_BOUND_PATTERN.match(sql) is not None + with self._lock: + self.attempted_sql.append(sql) + self._active[is_cursor_bound] += 1 + self._peak[is_cursor_bound] = max( + self._peak[is_cursor_bound], self._active[is_cursor_bound] + ) + try: + time.sleep(self.statement_latency_seconds) + rows, description = self._answer(sql=sql, params=params) + finally: + with self._lock: + self._active[is_cursor_bound] -= 1 + with self._lock: + self.queries.append( + RecordedQuery( + sql=sql, + params=params, + thread_id=threading.get_ident(), + started_at=started_at, + finished_at=time.monotonic(), + row_count=len(rows), + ) + ) + return rows, description + + def _load(self, *, table: str, rows: list[dict[str, object]]) -> None: + staged: pyarrow.Table = pyarrow.Table.from_pylist(rows) + self._database.register("staged_rows", staged) + self._database.execute(f"INSERT INTO {table} BY NAME SELECT * FROM staged_rows") + self._database.unregister("staged_rows") + + def _answer(self, *, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[Any, ...]], Any]: + routes: tuple[tuple[re.Pattern[str], Callable[..., Any]], ...] = ( + (_METADATA_RELATION_PATTERN, self._answer_information_schema), + (_OTHER_INFORMATION_SCHEMA_PATTERN, _answer_empty_metadata), + (_SHOW_COLUMNS_PATTERN, self._answer_show_columns), + (_CURSOR_BOUND_PATTERN, self._answer_cursor_bounds), + (_ANY_PATTERN, _answer_status), + ) + route: tuple[re.Pattern[str], Callable[..., Any]] = next( + filter(lambda candidate: candidate[0].search(sql) is not None, routes) + ) + return route[1](sql=sql, params=params) + + def _answer_information_schema( + self, *, sql: str, params: tuple[object, ...] + ) -> tuple[list[tuple[Any, ...]], Any]: + _MISSING_RELATION_ACTIONS[ + not self.failing_metadata_schemas.isdisjoint(str(param) for param in params) + ](sql) + duckdb_sql: str = _METADATA_RELATION_PATTERN.sub(_scope_metadata_relation, sql).replace( + "%s", "?" + ) + with self._lock: + cursor: duckdb.DuckDBPyConnection = self._database.cursor() + result: duckdb.DuckDBPyConnection = cursor.execute(duckdb_sql, list(params)) + rows: list[tuple[Any, ...]] = result.fetchall() + return rows, result.description + + def _answer_show_columns( + self, *, sql: str, params: tuple[object, ...] + ) -> tuple[list[tuple[Any, ...]], Any]: + del params + relation: str = next(_SHOW_COLUMNS_PATTERN.finditer(sql)).group("relation") + parts: list[str] = [part.strip('"').upper() for part in relation.split(".")] + matches: tuple[FakeRelation, ...] = tuple( + filter( + lambda fake: (fake.database, fake.schema, fake.name)[-len(parts) :] == tuple(parts), + self.relations, + ) + ) + rows: list[tuple[Any, ...]] = [] + fake: FakeRelation + for fake in matches: + rows.extend( + (fake.name, fake.schema, column.name, json.dumps(_show_columns_type(column))) + for column in fake.columns + ) + _MISSING_RELATION_ACTIONS[not rows](relation) + return rows, _SHOW_DESCRIPTION + + def _answer_cursor_bounds( + self, *, sql: str, params: tuple[object, ...] + ) -> tuple[list[tuple[Any, ...]], Any]: + del params + relation: str = next(_CURSOR_BOUND_PATTERN.finditer(sql)).group("relation") + _MISSING_RELATION_ACTIONS[relation.lower() in self.failing_relations](relation) + select_list: str = sql.split(" FROM ")[0] + values: tuple[str, ...] = tuple( + compress(self.cursor_values, (marker in select_list for marker in _BOUND_MARKERS)) + ) + return [values], tuple((f"_{index}",) for index in range(len(values))) + + +def _accept_relation(relation: str) -> None: + del relation + + +def _raise_unavailable_relation(relation: str) -> None: + raise RuntimeError(f"Object '{relation}' does not exist or not authorized.") + + +_MISSING_RELATION_ACTIONS: dict[bool, Callable[[str], None]] = { + False: _accept_relation, + True: _raise_unavailable_relation, +} + + +def _answer_empty_metadata( + *, sql: str, params: tuple[object, ...] +) -> tuple[list[tuple[Any, ...]], Any]: + del sql, params + return [], (("value",),) + + +def _answer_status(*, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[Any, ...]], Any]: + del sql, params + return [], _STATUS_DESCRIPTION + + +def _scope_metadata_relation(match: re.Match[str]) -> str: + table: str = f"fake_{match.group('view').lower()}" + database: str = (match.group("database") or "").replace("'", "''") + scoped: str = f"(SELECT * FROM {table} WHERE table_catalog = '{database}') AS {table}" + return (scoped, table)[match.group("database") is None] + + +def _table_row(relation: FakeRelation) -> dict[str, object]: + return { + "table_catalog": relation.database, + "table_schema": relation.schema, + "table_name": relation.name, + "table_type": relation.table_type, + "is_transient": ("NO", "YES")[relation.is_transient], + "created": _FIXED_CREATED_AT, + "last_altered": _FIXED_CREATED_AT, + "retention_time": relation.retention_time, + } + + +def _column_rows(relation: FakeRelation) -> list[dict[str, object]]: + rows: list[dict[str, object]] = [] + position: int + column: FakeColumn + for position, column in enumerate(relation.columns, start=1): + rows.append( + { + "table_catalog": relation.database, + "table_schema": relation.schema, + "table_name": relation.name, + "column_name": column.name, + "ordinal_position": position, + "data_type": column.data_type, + "numeric_precision": column.numeric_precision, + "numeric_scale": column.numeric_scale, + "character_maximum_length": column.character_maximum_length, + } + ) + return rows + + +def _show_columns_type(column: FakeColumn) -> dict[str, object]: + candidates: tuple[tuple[str, object], ...] = ( + ("type", _SHOW_TYPE_NAMES.get(column.data_type, column.data_type)), + ("nullable", True), + ("precision", column.numeric_precision), + ("scale", column.numeric_scale), + ("length", (None, column.character_maximum_length)[column.data_type == "TEXT"]), + ) + return dict(filter(lambda item: item[1] is not None, candidates)) + + +class FakeRawSnowflakeCursor: + """DB-API cursor double that delegates statements to the recording warehouse.""" + + def __init__(self, *, warehouse: RecordingSnowflakeWarehouse) -> None: + self._warehouse: RecordingSnowflakeWarehouse = warehouse + self._rows: list[tuple[Any, ...]] = [] + self.description: Any = None + self.rowcount: int = -1 + self.sfqid: str | None = None + + def execute(self, sql: str, params: Any = None, **kwargs: Any) -> FakeRawSnowflakeCursor: + del kwargs + rows, description = self._warehouse.run(sql=sql, params=tuple(params or ())) + self._rows = list(rows) + self.description = description + self.rowcount = len(rows) + return self + + def fetchall(self) -> list[tuple[Any, ...]]: + rows: list[tuple[Any, ...]] = self._rows + self._rows = [] + return rows + + def fetchone(self) -> tuple[Any, ...] | None: + first: list[tuple[Any, ...]] = self._rows[:1] + self._rows = self._rows[1:] + return next(iter(first), None) + + def fetchmany(self, size: int) -> list[tuple[Any, ...]]: + rows: list[tuple[Any, ...]] = self._rows[:size] + self._rows = self._rows[size:] + return rows + + def close(self) -> None: + self._rows = [] + + def __iter__(self) -> Iterator[tuple[Any, ...]]: + return iter(self.fetchall()) + + +class FakeRawSnowflakeConnection: + """Raw connection double handing out cursors bound to one recording warehouse.""" + + def __init__(self, *, warehouse: RecordingSnowflakeWarehouse) -> None: + self._warehouse: RecordingSnowflakeWarehouse = warehouse + + def cursor(self) -> FakeRawSnowflakeCursor: + return FakeRawSnowflakeCursor(warehouse=self._warehouse) + + def close(self) -> None: + return None + + +class OfflineSnowflakeAdapter(SnowflakeAdapter): + """Snowflake adapter whose connections reach the recording warehouse instead of the network.""" + + def __init__(self, *, warehouse: RecordingSnowflakeWarehouse) -> None: + super().__init__() + self._recording_warehouse: RecordingSnowflakeWarehouse = warehouse + + def connect(self, config: dict[str, Any]) -> _SnowflakeConnection: + del config + return _SnowflakeConnection(self._recording_warehouse.raw_connection()) + + +@dataclass(frozen=True) +class SyntheticSnowflakeProject: + """Synthetic multi-schema project files and the warehouse relations they inspect.""" + + project_dir: Path + relations: tuple[FakeRelation, ...] + model_schemas: dict[str, str] + incremental_model_names: tuple[str, ...] + + +_TABLE_MODEL_HEADER: str = "MODEL (\n materialized table,\n);\n\n" +_INCREMENTAL_MODEL_HEADER: str = ( + "MODEL (\n" + " materialized incremental,\n" + " incremental_strategy delete_insert,\n" + " cursor ordered_at,\n" + " cursor_type timestamp,\n" + " cursor_grain day,\n" + ");\n\n" +) +_ORDER_COLUMNS: tuple[FakeColumn, ...] = ( + FakeColumn(name="ORDER_ID", data_type="NUMBER", numeric_precision=38, numeric_scale=0), + FakeColumn(name="ORDERED_AT", data_type="TIMESTAMP_NTZ"), +) + + +def write_synthetic_snowflake_project( + *, + project_dir: Path, + database: str = "ANALYTICS", + model_schemas: tuple[str, ...] = ("staging", "intermediate", "marts"), + source_schema: str = "raw", + models_per_schema: int = 40, + incremental_every: int = 4, + sources: int = 30, + unmanaged_relations_per_schema: int = 600, + state_tables: tuple[str, ...] = ("_sqlbuild_fingerprints", "_sqlbuild_migrations"), +) -> SyntheticSnowflakeProject: + """Write a synthetic orders project and the matching warehouse relation set.""" + + (project_dir / "models").mkdir(parents=True, exist_ok=True) + (project_dir / "sources").mkdir(parents=True, exist_ok=True) + (project_dir / "sqlbuild_project.toml").write_text( + _project_toml(database=database, model_schemas=model_schemas), encoding="utf-8" + ) + (project_dir / "sources" / "raw.yml").write_text( + _sources_yml(database=database, source_schema=source_schema, sources=sources), + encoding="utf-8", + ) + relations: list[FakeRelation] = [ + FakeRelation( + database=database, + schema=source_schema.upper(), + name=f"ORDERS_{index}", + columns=_ORDER_COLUMNS, + ) + for index in range(sources) + ] + names_by_schema: dict[str, str] = {} + incremental_flags: list[bool] = [] + schema_index: int + schema: str + for schema_index, schema in enumerate(model_schemas): + (project_dir / "models" / schema).mkdir(parents=True, exist_ok=True) + model_index: int + for model_index in range(models_per_schema): + name: str = f"{schema}_orders_{model_index}" + upstream: str = ( + f'__ref("{model_schemas[schema_index - 1]}_orders_{model_index}")', + f'__source("raw_orders_{model_index % sources}")', + )[schema_index == 0] + incremental: bool = model_index % incremental_every == 0 + header: str = (_TABLE_MODEL_HEADER, _INCREMENTAL_MODEL_HEADER)[incremental] + (project_dir / "models" / schema / f"{name}.sql").write_text( + f"{header}SELECT order_id, ordered_at FROM {upstream}\n", encoding="utf-8" + ) + names_by_schema[name] = schema + incremental_flags.append(incremental) + relations.append( + FakeRelation( + database=database, + schema=schema.upper(), + name=name.upper(), + is_transient=model_index % 2 == 1, + columns=_ORDER_COLUMNS, + ) + ) + for schema in model_schemas: + relations.extend( + FakeRelation( + database=database, + schema=schema.upper(), + name=state_table.upper(), + is_transient=True, + ) + for state_table in state_tables + ) + for schema in (*model_schemas, source_schema): + relations.extend( + FakeRelation( + database=database, + schema=schema.upper(), + name=f"ADHOC_SCRATCH_{index}", + table_type=("BASE TABLE", "VIEW")[index % 5 == 0], + is_transient=index % 3 == 0, + ) + for index in range(unmanaged_relations_per_schema) + ) + return SyntheticSnowflakeProject( + project_dir=project_dir, + relations=tuple(relations), + model_schemas=names_by_schema, + incremental_model_names=tuple(compress(names_by_schema, incremental_flags)), + ) + + +def build_inspection_catalog_relations() -> tuple[FakeRelation, ...]: + """Return a metadata set with case variants, views, transient and external tables.""" + + return ( + FakeRelation(database="ANALYTICS", schema="STAGING", name="ORDERS", columns=_ORDER_COLUMNS), + FakeRelation( + database="ANALYTICS", + schema="STAGING", + name="orders", + columns=(FakeColumn(name="quoted_only", data_type="TEXT"),), + ), + FakeRelation(database="ANALYTICS", schema="STAGING", name="OrderLines"), + FakeRelation( + database="ANALYTICS", + schema="STAGING", + name="CUSTOMERS_V", + table_type="VIEW", + columns=(FakeColumn(name="CUSTOMER_ID", data_type="NUMBER", numeric_precision=38),), + ), + FakeRelation( + database="ANALYTICS", + schema="STAGING", + name="TRANSIENT_EVENTS", + is_transient=True, + retention_time=0, + columns=(FakeColumn(name="AMOUNT", data_type="FLOAT"),), + ), + FakeRelation( + database="ANALYTICS", + schema="STAGING", + name="PARTNER_FEED", + table_type="EXTERNAL TABLE", + ), + FakeRelation( + database="ANALYTICS", + schema="STAGING", + name="DAILY_TOTALS", + table_type="MATERIALIZED VIEW", + ), + *( + FakeRelation(database="ANALYTICS", schema="STAGING", name=f"INVENTORY_{index}") + for index in range(40) + ), + FakeRelation(database="ANALYTICS", schema="MARTS", name="REVENUE", columns=_ORDER_COLUMNS), + FakeRelation(database="ANALYTICS", schema="MARTS", name="revenue_v", table_type="VIEW"), + FakeRelation(database="ANALYTICS", schema="Mixed_Case", name="SUPPORT_TICKETS"), + FakeRelation(database="ARCHIVE", schema="STAGING", name="ORDERS", columns=_ORDER_COLUMNS), + ) + + +def build_offline_snowflake( + *, relations: tuple[FakeRelation, ...], statement_latency_seconds: float = 0.0 +) -> tuple[OfflineSnowflakeAdapter, _SnowflakeConnection, RecordingSnowflakeWarehouse]: + """Build an offline adapter, an open connection, and its recording warehouse.""" + + warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( + relations=relations, statement_latency_seconds=statement_latency_seconds + ) + adapter: OfflineSnowflakeAdapter = OfflineSnowflakeAdapter(warehouse=warehouse) + return adapter, adapter.connect({}), warehouse + + +def sorted_relation_reprs(relations: tuple[RelationInfo, ...]) -> list[str]: + """Return an order-independent comparable form of listed relations.""" + + return sorted(repr(relation) for relation in relations) + + +def _project_toml(*, database: str, model_schemas: tuple[str, ...]) -> str: + path_defaults: str = "".join( + f'[path_defaults.{schema}]\nschema = "{schema}"\n\n' for schema in model_schemas + ) + return ( + 'name = "orders_platform"\n' + 'adapter = "snowflake"\n\n' + "[connection]\n" + 'account = "example"\n' + 'user = "example"\n' + f'database = "{database.lower()}"\n' + f'schema = "{model_schemas[0]}"\n\n' + "[connection.session_parameters]\n" + "QUOTED_IDENTIFIERS_IGNORE_CASE = true\n\n" + "[defaults]\n" + f'database = "{database.lower()}"\n' + 'materialized = "table"\n\n' + path_defaults + ) + + +def _sources_yml(*, database: str, source_schema: str, sources: int) -> str: + lines: list[str] = ["sources:"] + index: int + for index in range(sources): + lines.extend( + ( + f" - name: raw_orders_{index}", + f" database: {database.lower()}", + f" schema: {source_schema}", + f" table: orders_{index}", + " columns:", + " - name: order_id", + " type: NUMBER(38, 0)", + " - name: ordered_at", + " type: TIMESTAMP_NTZ", + ) + ) + return "\n".join(lines) + "\n" diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py new file mode 100644 index 0000000000..3f69723535 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py @@ -0,0 +1,310 @@ +"""Catalog-served Snowflake metadata must equal what the adapter returns directly.""" + +from __future__ import annotations + +from collections import Counter + +import pytest + +from sqlbuild.adapter.contract.exceptions import AdapterUserError +from sqlbuild.adapter.contract.models import ( + ColumnInfo, + RelationInfo, + TableFreshnessMetadata, + TableFreshnessRequest, +) +from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog +from sqlbuild.adapter.relations.constants import INSPECTION_IN_LIST_LIMIT +from sqlbuild.adapter.relations.main.open_inspection_catalog import open_inspection_catalog +from tests.unit.src.sqlbuild.adapters.snowflake.inspection._test_types import ( + ColumnRequestEquivalenceTestCase, + FreshnessErrorTestCase, + FreshnessReuseTestCase, + RelationRequestEquivalenceTestCase, + RepeatedLookupTestCase, + SpeculativePrefetchTestCase, +) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + OfflineSnowflakeAdapter, + RecordingSnowflakeWarehouse, + build_inspection_catalog_relations, + build_offline_snowflake, + sorted_relation_reprs, +) + +_MANY_NAMES: tuple[str, ...] = ( + "orders", + "customers_v", + *(f"inventory_{index}" for index in range(INSPECTION_IN_LIST_LIMIT + 10)), +) + + +@pytest.mark.parametrize( + "test_case", + [ + RelationRequestEquivalenceTestCase( + description="whole schema with case variants, views, transient and external tables", + database="analytics", + schemas=("staging",), + names=None, + expected_relation_count=47, + ), + RelationRequestEquivalenceTestCase( + description="name filter across schemas ignores quoted lowercase twins", + database="analytics", + schemas=("staging", "marts"), + names=("orders", "orderlines", "customers_v", "revenue_v", "missing"), + expected_relation_count=2, + ), + RelationRequestEquivalenceTestCase( + description="uppercase request matches unquoted storage", + database="ANALYTICS", + schemas=("STAGING",), + names=("ORDERS",), + expected_relation_count=1, + ), + RelationRequestEquivalenceTestCase( + description="quoted mixed-case schema is not matched by a logical name", + database="analytics", + schemas=("mixed_case",), + names=None, + expected_relation_count=0, + ), + RelationRequestEquivalenceTestCase( + description="session database lists every catalog", + database=None, + schemas=("staging",), + names=("orders",), + expected_relation_count=2, + ), + RelationRequestEquivalenceTestCase( + description="database-wide name lookup", + database="analytics", + schemas=None, + names=("orders", "revenue"), + expected_relation_count=2, + ), + RelationRequestEquivalenceTestCase( + description="database-wide lookup beyond the IN-list cap is chunked", + database="analytics", + schemas=None, + names=_MANY_NAMES, + expected_relation_count=42, + ), + ], + ids=lambda case: case.description, +) +def test_given_metadata_request_when_served_by_catalog_then_matches_direct_adapter( + test_case: RelationRequestEquivalenceTestCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=build_inspection_catalog_relations() + ) + direct: tuple[RelationInfo, ...] = adapter.list_relations( + connection=connection, + database=test_case.database, + schemas=test_case.schemas, + names=test_case.names, + ) + warehouse.reset() + + served: tuple[RelationInfo, ...] = InspectionCatalog( + adapter=adapter, connection=connection + ).list_relations(database=test_case.database, schemas=test_case.schemas, names=test_case.names) + + assert len(served) == test_case.expected_relation_count + assert sorted_relation_reprs(served) == sorted_relation_reprs(direct) + assert all(query.largest_in_list <= INSPECTION_IN_LIST_LIMIT for query in warehouse.queries) + + +@pytest.mark.parametrize( + "test_case", + [ + ColumnRequestEquivalenceTestCase( + description="sparse request keeps exact SHOW COLUMNS reads", + schemas=("staging", "marts"), + names=("orders", "customers_v", "transient_events", "revenue"), + expected_query_kinds=("show_columns",) * 4, + ), + ColumnRequestEquivalenceTestCase( + description="broad request reads each schema's columns once without IN lists", + schemas=("staging",), + names=None, + expected_query_kinds=("columns",), + ), + ], + ids=lambda case: case.description, +) +def test_given_relations_when_reading_columns_through_catalog_then_matches_direct_adapter( + test_case: ColumnRequestEquivalenceTestCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=build_inspection_catalog_relations() + ) + relations: tuple[RelationInfo, ...] = adapter.list_relations( + connection=connection, + database="analytics", + schemas=test_case.schemas, + names=test_case.names, + ) + direct: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( + adapter.get_columns_for_relations(connection=connection, relations=relations) + ) + warehouse.reset() + catalog: InspectionCatalog = InspectionCatalog(adapter=adapter, connection=connection) + + served: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( + catalog.get_columns_for_relations(relations=relations) + ) + repeated: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( + catalog.get_columns_for_relations(relations=relations) + ) + + assert served == direct + assert repeated == direct + assert Counter(query.kind for query in warehouse.queries) == Counter( + test_case.expected_query_kinds + ) + assert all(query.largest_in_list == 0 for query in warehouse.queries) + + +@pytest.mark.parametrize( + "test_case", + [ + FreshnessReuseTestCase( + description="freshness reuses listings without another table read", + requests=( + ("analytics", "staging", "orders"), + ("analytics", "staging", "transient_events"), + ("analytics", "marts", "revenue"), + ), + listed_schemas=("staging", "marts"), + expected_metadata_reads=2, + ) + ], + ids=lambda case: case.description, +) +def test_given_open_catalog_when_reading_source_freshness_then_reuses_schema_listing( + test_case: FreshnessReuseTestCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=build_inspection_catalog_relations() + ) + requests: tuple[TableFreshnessRequest, ...] = tuple( + TableFreshnessRequest(database=database, schema=schema, name=name) + for database, schema, name in test_case.requests + ) + direct: dict[TableFreshnessRequest, TableFreshnessMetadata] = ( + adapter.get_tables_freshness_metadata(connection=connection, requests=requests) + ) + warehouse.reset() + + with open_inspection_catalog(adapter=adapter, connection=connection) as catalog: + _ = catalog.list_relations(database="analytics", schemas=test_case.listed_schemas) + served: dict[TableFreshnessRequest, TableFreshnessMetadata] = ( + adapter.get_tables_freshness_metadata(connection=connection, requests=requests) + ) + + assert served == direct + assert len(warehouse.queries) == test_case.expected_metadata_reads + + +@pytest.mark.parametrize( + "test_case", + [ + FreshnessErrorTestCase( + description="view freshness fails like the direct read", + request=("analytics", "staging", "customers_v"), + expected_error_fragment="only supports physical tables", + ) + ], + ids=lambda case: case.description, +) +def test_given_open_catalog_when_freshness_targets_a_view_then_raises_like_direct_read( + test_case: FreshnessErrorTestCase, +) -> None: + adapter, connection, _ = build_offline_snowflake(relations=build_inspection_catalog_relations()) + database, schema, name = test_case.request + requests: tuple[TableFreshnessRequest, ...] = ( + TableFreshnessRequest(database=database, schema=schema, name=name), + ) + with pytest.raises(AdapterUserError, match=test_case.expected_error_fragment): + _ = adapter.get_tables_freshness_metadata(connection=connection, requests=requests) + + with ( + open_inspection_catalog(adapter=adapter, connection=connection), + pytest.raises(AdapterUserError, match=test_case.expected_error_fragment), + ): + _ = adapter.get_tables_freshness_metadata(connection=connection, requests=requests) + + +@pytest.mark.parametrize( + "test_case", + [ + RepeatedLookupTestCase( + description="state table, model, and whole-schema lookups share one listing", + lookups=(("_sqlbuild_fingerprints",), ("orders",), None), + expected_relation_counts=(0, 1, 47), + expected_query_kinds=("tables",), + ) + ], + ids=lambda case: case.description, +) +def test_given_repeated_lookups_when_catalog_is_open_then_lists_schema_once( + test_case: RepeatedLookupTestCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=build_inspection_catalog_relations() + ) + catalog: InspectionCatalog = InspectionCatalog(adapter=adapter, connection=connection) + + counts: tuple[int, ...] = tuple( + len(catalog.list_relations(database="analytics", schemas=("staging",), names=names)) + for names in test_case.lookups + ) + + assert counts == test_case.expected_relation_counts + assert tuple(query.kind for query in warehouse.queries) == test_case.expected_query_kinds + + +@pytest.mark.parametrize( + "test_case", + [ + SpeculativePrefetchTestCase( + description="unreadable speculative schema fails only the request that needs it", + failing_schema="RESTRICTED", + expected_error_fragment="not authorized", + expected_failed_reads=2, + ) + ], + ids=lambda case: case.description, +) +def test_given_unreadable_schema_when_prefetching_best_effort_then_only_real_reads_fail( + test_case: SpeculativePrefetchTestCase, +) -> None: + warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( + relations=build_inspection_catalog_relations(), + failing_metadata_schemas=frozenset({test_case.failing_schema}), + ) + adapter: OfflineSnowflakeAdapter = OfflineSnowflakeAdapter(warehouse=warehouse) + catalog: InspectionCatalog = InspectionCatalog(adapter=adapter, connection=adapter.connect({})) + + catalog.prefetch_schema_listings( + scopes=(("analytics", "staging"), ("analytics", test_case.failing_schema.lower())), + best_effort=True, + ) + staging: tuple[RelationInfo, ...] = catalog.list_relations( + database="analytics", schemas=("staging",), names=("orders",) + ) + with pytest.raises(RuntimeError, match=test_case.expected_error_fragment): + _ = catalog.list_relations( + database="analytics", schemas=(test_case.failing_schema.lower(),) + ) + + assert len(staging) == 1 + assert len(warehouse.queries_of_kind("tables")) == 1 + assert len(warehouse.attempted_sql) - len(warehouse.queries) == test_case.expected_failed_reads + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) From bc767efc981ca57d066c12039c934ed8fbeb9d83 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 11:35:59 +0100 Subject: [PATCH 2/7] perf(planner): reuse schema metadata and read cursor bounds concurrently Open one inspection catalog per execution plan and route snapshot, migration-state, source-column, and semantic-source metadata lookups through it, so each schema is listed once and state-table and source lookups reuse the listing. Target and source schemas are prefetched together (best effort; a failed speculative read is re-read and reported by the lookup that needs it). Snapshot columns are read only for listed relations that planning looks up by node name. Cursor-bound MIN/MAX reads stay one statement per relation and run with the adapter's bounded inspection concurrency (8 on Snowflake, sequential elsewhere); results merge in query order and a failed read still leaves only that relation unavailable. Bounded eligible-max reads fail the plan naming the model and relation, as before. Adds a recording-warehouse budget guard over a synthetic 4-schema project, an equivalence check against the direct adapter path, and an architecture guard keeping catalog SQL in adapter modules. Refs CHI-455 --- .../planner/_helpers/output/plan_entry.py | 56 +-- .../_helpers/warehouse/semantic_sources.py | 65 +++- .../planner/_helpers/warehouse/snapshot.py | 326 ++++++++++++++---- .../classes/migration_state_inspection.py | 20 +- .../planner/main/execution/execution.py | 16 +- .../_helpers/test_warehouse_snapshot.py | 9 +- .../compiler/planner/main/_test_types.py | 28 ++ .../sqlbuild/compiler/planner/main/helpers.py | 63 ++++ .../main/test_snowflake_inspection_budget.py | 155 +++++++++ .../sqlbuild/adapter/relations/_test_types.py | 12 + .../src/sqlbuild/adapter/relations/helpers.py | 26 ++ .../relations/test_metadata_sql_ownership.py | 51 +++ .../planner/_helpers/warehouse/_test_types.py | 17 + .../planner/_helpers/warehouse/helpers.py | 145 ++++++++ .../test_warehouse_cursor_concurrency.py | 139 ++++++++ 15 files changed, 1010 insertions(+), 118 deletions(-) create mode 100644 tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/_test_types.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/helpers.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/test_metadata_sql_ownership.py create mode 100644 tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py create mode 100644 tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_warehouse_cursor_concurrency.py diff --git a/src/sqlbuild/compiler/planner/_helpers/output/plan_entry.py b/src/sqlbuild/compiler/planner/_helpers/output/plan_entry.py index ffeb2240a4..88c9a47452 100644 --- a/src/sqlbuild/compiler/planner/_helpers/output/plan_entry.py +++ b/src/sqlbuild/compiler/planner/_helpers/output/plan_entry.py @@ -8,7 +8,11 @@ from typing import Any from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter -from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo +from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo, RelationLookup +from sqlbuild.adapter.relations.main.get_columns_for_inspection import get_columns_for_inspection +from sqlbuild.adapter.relations.main.list_relations_for_inspection import ( + list_relations_for_inspection, +) from sqlbuild.compiler.compile.constants import CURSOR_INPUTS_CONFIG_KEY from sqlbuild.compiler.compile.exceptions import CompileInputError from sqlbuild.compiler.compile.main._cursor_roles import resolve_cursor_input_roles @@ -1056,41 +1060,37 @@ def gather_source_columns( schemas: set[str] for db_key_iter, schemas in source_schemas.items(): database: str | None = db_key_iter or None - names: tuple[str, ...] | None = _build_source_table_name_filter( - project=project, - database=database, - schemas=schemas, - source_entries=entries, - ) - relations: tuple[RelationInfo, ...] = adapter.list_relations( + scoped: dict[str, tuple[str | None, str | None, str]] = { + entry_iter.name: RelationLookup.key( + database=database, + schema=entry_iter.schema, + name=entry_iter.table if entry_iter.table is not None else entry_iter.name, + ) + for entry_iter in entries + if entry_iter.expression is None + and (entry_iter.database or None) == database + and entry_iter.schema in schemas + } + relations: tuple[RelationInfo, ...] = list_relations_for_inspection( + adapter=adapter, connection=connection, database=database, schemas=tuple(sorted(schemas)), - names=names, + names=_build_source_table_name_filter( + project=project, database=database, schemas=schemas, source_entries=entries + ), ) + wanted: frozenset[tuple[str | None, str | None, str]] = frozenset(scoped.values()) all_columns: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( - adapter.get_columns_for_relations( + get_columns_for_inspection( + adapter=adapter, connection=connection, - relations=relations, + relations=tuple(relation for relation in relations if relation.identity in wanted), ) ) - entry_iter: SourceEntry - for entry_iter in entries: - if ( - entry_iter.expression is not None - or (entry_iter.database or None) != database - or entry_iter.schema not in schemas - ): - continue - table_name: str = entry_iter.table if entry_iter.table is not None else entry_iter.name - identity: tuple[str | None, str | None, str] = ( - None if database is None else database.lower(), - None if entry_iter.schema is None else entry_iter.schema.lower(), - table_name.lower(), - ) - cols: tuple[ColumnInfo, ...] | None = all_columns.get(identity) - if cols is not None: - result[entry_iter.name] = cols + result.update( + {name: all_columns[key] for name, key in scoped.items() if key in all_columns} + ) return result diff --git a/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py b/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py index 1d0be9b2f1..54d68acc94 100644 --- a/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py +++ b/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py @@ -6,6 +6,10 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo +from sqlbuild.adapter.relations.main.get_columns_for_inspection import get_columns_for_inspection +from sqlbuild.adapter.relations.main.list_relations_for_inspection import ( + list_relations_for_inspection, +) from sqlbuild.compiler.compile.models import CompiledObjectKey, CompiledProject from sqlbuild.compiler.planner.types import ContractPolicy from sqlbuild.compiler.references.types import SqlReferenceKind @@ -44,15 +48,10 @@ def get_semantic_source_columns( by_database.setdefault(entry.database, []).append(entry) result: dict[str, tuple[ColumnInfo, ...]] = dict(columns) for database, entries in by_database.items(): - relations: tuple[RelationInfo, ...] = adapter.list_relations( - connection=connection, - database=database, - schemas=None, - names=tuple(sorted({entry.table or entry.name for entry in entries})), - ) - inspected: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( - adapter.get_columns_for_relations(connection=connection, relations=relations) + relations: tuple[RelationInfo, ...] = _list_source_candidates( + adapter=adapter, connection=connection, database=database, entries=tuple(entries) ) + chosen: dict[str, RelationInfo] = {} for entry in entries: candidates: tuple[RelationInfo, ...] = tuple( relation @@ -63,14 +62,56 @@ def get_semantic_source_columns( or (relation.schema or "").casefold() == entry.schema.casefold() ) ) - if len(candidates) != 1: - continue - relation: RelationInfo = candidates[0] + if len(candidates) == 1: + chosen[entry.name] = candidates[0] + inspected: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( + get_columns_for_inspection( + adapter=adapter, + connection=connection, + relations=tuple(dict.fromkeys(chosen.values())), + ) + if chosen + else {} + ) + name: str + relation: RelationInfo + for name, relation in chosen.items(): identity: tuple[str | None, str | None, str] = ( relation.database.lower() if relation.database else None, relation.schema.lower() if relation.schema else None, relation.name.lower(), ) if identity in inspected: - result[entry.name] = inspected[identity] + result[name] = inspected[identity] return result + + +def _list_source_candidates( + *, adapter: BaseAdapter, connection: Any, database: str | None, entries: tuple[SourceEntry, ...] +) -> tuple[RelationInfo, ...]: + """List schema-qualified sources from their schemas and unqualified ones by name.""" + + scoped: tuple[SourceEntry, ...] = tuple(entry for entry in entries if entry.schema is not None) + unscoped: tuple[SourceEntry, ...] = tuple(entry for entry in entries if entry.schema is None) + relations: list[RelationInfo] = [] + if scoped: + relations.extend( + list_relations_for_inspection( + adapter=adapter, + connection=connection, + database=database, + schemas=tuple(sorted({entry.schema or "" for entry in scoped})), + names=tuple(sorted({entry.table or entry.name for entry in scoped})), + ) + ) + if unscoped: + relations.extend( + list_relations_for_inspection( + adapter=adapter, + connection=connection, + database=database, + schemas=None, + names=tuple(sorted({entry.table or entry.name for entry in unscoped})), + ) + ) + return tuple(dict.fromkeys(relations)) diff --git a/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py b/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py index c1935082a9..958eda3ce6 100644 --- a/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py +++ b/src/sqlbuild/compiler/planner/_helpers/warehouse/snapshot.py @@ -12,6 +12,18 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo from sqlbuild.adapter.contract.types import AdapterExecute +from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog +from sqlbuild.adapter.relations.main.active_inspection_catalog import active_inspection_catalog +from sqlbuild.adapter.relations.main.get_columns_for_inspection import get_columns_for_inspection +from sqlbuild.adapter.relations.main.list_relations_for_inspection import ( + list_relations_for_inspection, +) +from sqlbuild.adapter.relations.main.record_inspection_query import record_inspection_query +from sqlbuild.adapter.relations.main.run_bounded_inspections import run_bounded_inspections +from sqlbuild.adapter.relations.main.run_recorded_inspection_query import ( + run_recorded_inspection_query, +) +from sqlbuild.adapter.relations.models import InspectionQueryRecord from sqlbuild.compiler.compile.main._cursor_roles import resolve_cursor_input_roles from sqlbuild.compiler.compile.models import ( CompiledFunction, @@ -106,6 +118,15 @@ class _PhysicalCursorQuery: cursor_type: str = CursorType.TIMESTAMP +@dataclass(frozen=True) +class _CursorQueryOutcome: + """Result of one physical cursor statement; ``row`` is None when nothing was returned.""" + + row: Any + elapsed_seconds: float + error: Exception | None = None + + @dataclass(frozen=True) class _UpstreamCursorInfo: """Pre-resolved upstream cursor metadata for one ref.""" @@ -240,6 +261,14 @@ def gather_warehouse_snapshot( if not schemas and metadata_names is None: return WarehouseSnapshot() query_schemas: tuple[str, ...] | None = schemas or None + _prefetch_inspection_schemas( + project=project, + adapter=adapter, + connection=connection, + database=database, + schemas=schemas, + relevant_keys=relevant_keys, + ) relations: dict[str, RelationInfo] state_schemas: _StateTableSchemas @@ -254,6 +283,7 @@ def gather_warehouse_snapshot( fingerprint_state_schemas: frozenset[str] = state_schemas.fingerprints freshness_state_schemas: frozenset[str] = state_schemas.source_freshness columns: dict[str, tuple[ColumnInfo, ...]] = _gather_columns( + project=project, adapter=adapter, connection=connection, relations=relations, @@ -363,6 +393,36 @@ def gather_redirected_cursor_snapshots( } +def _prefetch_inspection_schemas( + *, + project: CompiledProject, + adapter: BaseAdapter, + connection: Any, + database: str | None, + schemas: tuple[str, ...], + relevant_keys: frozenset[CompiledObjectKey] | None, +) -> None: + """List target and source schemas together, concurrently; later reads report failures.""" + + catalog: InspectionCatalog | None = active_inspection_catalog( + adapter=adapter, connection=connection + ) + if catalog is None: + return + scopes: list[tuple[str | None, str]] = [(database, schema) for schema in schemas] + source: CompiledSource + for source in project.sources: + entry: SourceEntry = source.source_entry + if ( + (relevant_keys is not None and source.key not in relevant_keys) + or entry.expression is not None + or entry.schema is None + ): + continue + scopes.append((entry.database, entry.schema)) + catalog.prefetch_schema_listings(scopes=tuple(scopes), best_effort=True) + + def _relevant_state_keys( *, project: CompiledProject, @@ -535,8 +595,8 @@ def _gather_relations( ) -> tuple[dict[str, RelationInfo], _StateTableSchemas]: """Fetch relations and the schemas where fingerprint, freshness and old-name state live.""" - relations: tuple[RelationInfo, ...] = adapter.list_relations( - connection=connection, database=database, schemas=schemas, names=names + relations: tuple[RelationInfo, ...] = list_relations_for_inspection( + adapter=adapter, connection=connection, database=database, schemas=schemas, names=names ) result: dict[str, RelationInfo] = {} logical_names_by_identity: dict[tuple[str | None, str | None, str], str] = {} @@ -586,19 +646,35 @@ def _location_identity( def _gather_columns( *, + project: CompiledProject, adapter: BaseAdapter, connection: Any, relations: dict[str, RelationInfo], ) -> dict[str, tuple[ColumnInfo, ...]]: - """Fetch column metadata for all relations across target schemas.""" + """Fetch columns only for listed relations that planning looks up by node name.""" - physical_relations: tuple[RelationInfo, ...] = tuple(relations.values()) + node_names: frozenset[str] = frozenset( + ( + *(model.name for model in project.models), + *(seed.name for seed in project.seeds), + *(function.name for function in project.functions), + ) + ) + needed: dict[str, RelationInfo] = { + logical_name: relation + for logical_name, relation in relations.items() + if logical_name in node_names + } all_columns: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( - adapter.get_columns_for_relations(connection=connection, relations=physical_relations) + get_columns_for_inspection( + adapter=adapter, connection=connection, relations=tuple(needed.values()) + ) + if needed + else {} ) return { logical_name: all_columns[relation.identity] - for logical_name, relation in relations.items() + for logical_name, relation in needed.items() if relation.identity in all_columns } @@ -696,11 +772,13 @@ def _gather_cursor_snapshots( queries: list[_PhysicalCursorQuery] = _build_cursor_queries(cursor_models) cursor_start: float = time.monotonic() + concurrency: int = _inspection_concurrency(adapter=adapter, connection=connection) results: dict[str, CursorScalar] = _execute_cursor_queries( queries=queries, connection=connection, execute=execute, on_progress=on_progress, + concurrency=concurrency, ) results = _gather_eligible_target_maxes( cursor_models=cursor_models, @@ -1020,76 +1098,168 @@ def _execute_cursor_queries( connection: Any, execute: AdapterExecute[Any, Any], on_progress: Callable[[str], None] | None, + concurrency: int = 1, ) -> dict[str, CursorScalar]: - """Execute standalone physical cursor queries and fan values out to logical tags.""" + """Run one unmerged MIN/MAX statement per relation, bounded-parallel, merged in query order.""" - results: dict[str, CursorScalar] = {} total: int = len(queries) - query_index: int - query: _PhysicalCursorQuery - for query_index, query in enumerate(queries, start=1): - query_results: dict[str, CursorScalar] = _execute_cursor_query( - query=query, - query_index=query_index, - total=total, - connection=connection, - execute=execute, - on_progress=on_progress, + if concurrency <= 1 or total <= 1: + results: dict[str, CursorScalar] = {} + query_index: int + query: _PhysicalCursorQuery + for query_index, query in enumerate(queries, start=1): + identity: str = _cursor_query_identity(query=query, index=query_index, total=total) + if on_progress is not None: + on_progress(f"Inspecting cursor bounds {identity}...") + outcome: _CursorQueryOutcome = _read_cursor_query( + query=query, connection=connection, execute=execute + ) + _report_cursor_query( + query=query, outcome=outcome, identity=identity, on_progress=on_progress + ) + results.update(_fan_out_cursor_row(query=query, row=outcome.row)) + return results + if on_progress is not None: + on_progress( + f"Inspecting cursor bounds for {total} relations " + f"({min(concurrency, total)} concurrent)..." + ) + progress: _CursorBoundsProgress = _CursorBoundsProgress( + queries=tuple(queries), on_progress=on_progress + ) + outcomes: list[_CursorQueryOutcome] = run_bounded_inspections( + tasks=tuple( + _bind_cursor_query(query=query, connection=connection, execute=execute) + for query in queries + ), + concurrency=concurrency, + on_complete=progress.report, + ) + merged: dict[str, CursorScalar] = {} + outcome_item: _CursorQueryOutcome + for query, outcome_item in zip(queries, outcomes, strict=True): + merged.update(_fan_out_cursor_row(query=query, row=outcome_item.row)) + return merged + + +class _CursorBoundsProgress: + """Report concurrent cursor-bound completions in finish order from the planning thread.""" + + def __init__( + self, + *, + queries: tuple[_PhysicalCursorQuery, ...], + on_progress: Callable[[str], None] | None, + ) -> None: + self._queries: tuple[_PhysicalCursorQuery, ...] = queries + self._on_progress: Callable[[str], None] | None = on_progress + self._completed: int = 0 + + def report(self, *, index: int, result: _CursorQueryOutcome) -> None: + """Report one finished relation read.""" + + self._completed += 1 + _report_cursor_query( + query=self._queries[index], + outcome=result, + identity=_cursor_query_identity( + query=self._queries[index], index=self._completed, total=len(self._queries) + ), + on_progress=self._on_progress, ) - results.update(query_results) - return results +def _inspection_concurrency(*, adapter: BaseAdapter, connection: Any) -> int: + catalog: InspectionCatalog | None = active_inspection_catalog( + adapter=adapter, connection=connection + ) + return 1 if catalog is None else catalog.concurrency -def _execute_cursor_query( - *, - query: _PhysicalCursorQuery, - query_index: int, - total: int, - connection: Any, - execute: AdapterExecute[Any, Any], - on_progress: Callable[[str], None] | None, -) -> dict[str, CursorScalar]: - """Execute one physical cursor query and fan out its returned values.""" +def _cursor_query_identity(*, query: _PhysicalCursorQuery, index: int, total: int) -> str: + return f"({index}/{total}): {query.relation}.{query.cursor_column} [{_cursor_bounds(query)}]" + + +def _cursor_bounds(query: _PhysicalCursorQuery) -> str: bounds: str = ",".join(("min",) if query.min_tags else ()) if query.max_tags: bounds = f"{bounds},max" if bounds else "max" - identity: str = f"({query_index}/{total}): {query.relation}.{query.cursor_column} [{bounds}]" + return bounds + + +def _cursor_query_sql(query: _PhysicalCursorQuery) -> str: select_parts: list[str] = [] if query.min_tags: select_parts.append(f"CAST(MIN({query.cursor_column}) AS VARCHAR) AS _min") if query.max_tags: select_parts.append(f"CAST(MAX({query.cursor_column}) AS VARCHAR) AS _max") - sql: str = f"SELECT {', '.join(select_parts)} FROM {query.relation}" - if on_progress is not None: - on_progress(f"Inspecting cursor bounds {identity}...") + return f"SELECT {', '.join(select_parts)} FROM {query.relation}" + + +def _bind_cursor_query( + *, query: _PhysicalCursorQuery, connection: Any, execute: AdapterExecute[Any, Any] +) -> Callable[[], _CursorQueryOutcome]: + def read() -> _CursorQueryOutcome: + return _read_cursor_query(query=query, connection=connection, execute=execute) + + return read + +def _read_cursor_query( + *, query: _PhysicalCursorQuery, connection: Any, execute: AdapterExecute[Any, Any] +) -> _CursorQueryOutcome: + """Run one physical cursor statement; a failure leaves only that relation unavailable.""" + + sql: str = _cursor_query_sql(query) query_start: float = time.monotonic() try: result: Any = execute(connection=connection, sql=sql) rows: list[Any] = result.fetchall() except Exception as error: elapsed: float = time.monotonic() - query_start - if on_progress is not None: - on_progress(f"Failed cursor bounds {identity} ({elapsed:.2f}s): {error}") - log_debug_event( - logger=_DEBUG_LOGGER, - message="cursor bounds physical query failed; treating relation as unavailable", - sqlbuild_relation=query.relation, - sqlbuild_cursor_column=query.cursor_column, - sqlbuild_bounds=bounds, - sqlbuild_elapsed_seconds=f"{elapsed:.2f}", - sqlbuild_error=str(error), + record_inspection_query( + record=InspectionQueryRecord( + sql=sql, elapsed_seconds=elapsed, row_count=None, error=str(error) + ) ) - return {} + return _CursorQueryOutcome(row=None, elapsed_seconds=elapsed, error=error) elapsed = time.monotonic() - query_start + record_inspection_query( + record=InspectionQueryRecord(sql=sql, elapsed_seconds=elapsed, row_count=len(rows)) + ) + return _CursorQueryOutcome(row=rows[0] if rows else None, elapsed_seconds=elapsed) + + +def _report_cursor_query( + *, + query: _PhysicalCursorQuery, + outcome: _CursorQueryOutcome, + identity: str, + on_progress: Callable[[str], None] | None, +) -> None: + if outcome.error is None: + if on_progress is not None: + on_progress(f"Inspected cursor bounds {identity} ({outcome.elapsed_seconds:.2f}s)") + return if on_progress is not None: - on_progress(f"Inspected cursor bounds {identity} ({elapsed:.2f}s)") - if not rows: + on_progress( + f"Failed cursor bounds {identity} ({outcome.elapsed_seconds:.2f}s): {outcome.error}" + ) + log_debug_event( + logger=_DEBUG_LOGGER, + message="cursor bounds physical query failed; treating relation as unavailable", + sqlbuild_relation=query.relation, + sqlbuild_cursor_column=query.cursor_column, + sqlbuild_bounds=_cursor_bounds(query), + sqlbuild_elapsed_seconds=f"{outcome.elapsed_seconds:.2f}", + sqlbuild_error=str(outcome.error), + ) + + +def _fan_out_cursor_row(*, query: _PhysicalCursorQuery, row: Any) -> dict[str, CursorScalar]: + if row is None: return {} output: dict[str, CursorScalar] = {} - row: Any = rows[0] value_index: int = 0 if query.min_tags: min_value: Any = row[value_index] @@ -1130,6 +1300,7 @@ def _gather_eligible_target_maxes( if invocation_time is None: return results eligible_results: dict[str, CursorScalar] = dict(results) + requests: list[tuple[_CursorModelInfo, str]] = [] info: _CursorModelInfo for info in cursor_models: config: StartCursorsConfig | None = info.start_cursor_config @@ -1154,22 +1325,29 @@ def _gather_eligible_target_maxes( ) if compare(left=target_max, right=horizon) <= 0: continue - horizon_is_date: bool = isinstance(horizon, DateValue) - eligible_sql: str = adapter.render_max_cursor_at_or_before( - relation=info.target_relation, - cursor_column=info.physical_cursor_column, - maximum_allowed=render(value=horizon), - cursor_type=info.cursor_type, - is_date=horizon_is_date, + requests.append( + ( + info, + adapter.render_max_cursor_at_or_before( + relation=info.target_relation, + cursor_column=info.physical_cursor_column, + maximum_allowed=render(value=horizon), + cursor_type=info.cursor_type, + is_date=isinstance(horizon, DateValue), + ), + ) ) - try: - query_result: Any = execute(connection=connection, sql=eligible_sql) - rows: list[Any] = query_result.fetchall() - except Exception as error: - raise PlannerInputError( - f"model '{info.model_name}': failed to query highest eligible target cursor " - f"for {info.target_relation}.{info.cursor_column}: {error}" - ) from error + rows_by_request: list[list[Any]] = run_bounded_inspections( + tasks=tuple( + _bind_eligible_target_max( + info=request_info, sql=sql, connection=connection, execute=execute + ) + for request_info, sql in requests + ), + concurrency=_inspection_concurrency(adapter=adapter, connection=connection), + ) + rows: list[Any] + for (info, _sql), rows in zip(requests, rows_by_request, strict=True): if rows and rows[0][0] is not None: eligible_results[f"{info.model_name}__target__eligible_max"] = _normalize_cursor_grain( value=parse(raw=rows[0][0], cursor_type=info.cursor_type or CursorType.TIMESTAMP), @@ -1179,6 +1357,28 @@ def _gather_eligible_target_maxes( return eligible_results +def _bind_eligible_target_max( + *, + info: _CursorModelInfo, + sql: str, + connection: Any, + execute: AdapterExecute[Any, Any], +) -> Callable[[], list[Any]]: + def read() -> list[Any]: + try: + return run_recorded_inspection_query( + sql=sql, + run=lambda: execute(connection=connection, sql=sql).fetchall(), + ) + except Exception as error: + raise PlannerInputError( + f"model '{info.model_name}': failed to query highest eligible target cursor " + f"for {info.target_relation}.{info.cursor_column}: {error}" + ) from error + + return read + + def _normalize_cursor_grain( *, value: CursorScalar, cursor_type: str | None, cursor_grain: str | None ) -> CursorScalar: diff --git a/src/sqlbuild/compiler/planner/classes/migration_state_inspection.py b/src/sqlbuild/compiler/planner/classes/migration_state_inspection.py index 3824981075..309a1441a3 100644 --- a/src/sqlbuild/compiler/planner/classes/migration_state_inspection.py +++ b/src/sqlbuild/compiler/planner/classes/migration_state_inspection.py @@ -6,6 +6,10 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo +from sqlbuild.adapter.relations.main.get_columns_for_inspection import get_columns_for_inspection +from sqlbuild.adapter.relations.main.list_relations_for_inspection import ( + list_relations_for_inspection, +) from sqlbuild.compiler.compile.models import CompiledRelationLocation from sqlbuild.compiler.fingerprints.constants import FINGERPRINT_TABLE_NAME, NODE_TYPE_MODEL from sqlbuild.compiler.fingerprints.main.read import read_latest_fingerprints @@ -78,7 +82,8 @@ def inspect_column_events(self, *, schemas: set[str]) -> None: if not pending: return self._column_event_schemas.update(schema.lower() for schema in pending) - listed: tuple[RelationInfo, ...] = self._adapter.list_relations( + listed: tuple[RelationInfo, ...] = list_relations_for_inspection( + adapter=self._adapter, connection=self._connection, database=self._database, schemas=pending, @@ -117,7 +122,8 @@ def inspect_old_name_views(self, *, schemas: set[str]) -> None: if not pending: return self._old_name_schemas.update(schema.lower() for schema in pending) - listed: tuple[RelationInfo, ...] = self._adapter.list_relations( + listed: tuple[RelationInfo, ...] = list_relations_for_inspection( + adapter=self._adapter, connection=self._connection, database=self._database, schemas=pending, @@ -167,7 +173,8 @@ def inspect_schemas(self, *, schemas: set[str]) -> None: if not pending: return self._read_schemas.update(schema.lower() for schema in pending) - listed: tuple[RelationInfo, ...] = self._adapter.list_relations( + listed: tuple[RelationInfo, ...] = list_relations_for_inspection( + adapter=self._adapter, connection=self._connection, database=self._database, schemas=pending, @@ -217,7 +224,8 @@ def inspect_relations(self, *, locations: tuple[CompiledRelationLocation, ...]) ) if not pending: return - listed: tuple[RelationInfo, ...] = self._adapter.list_relations( + listed: tuple[RelationInfo, ...] = list_relations_for_inspection( + adapter=self._adapter, connection=self._connection, database=self._database, schemas=tuple(sorted({location.schema or "" for location in pending})), @@ -228,7 +236,9 @@ def inspect_relations(self, *, locations: tuple[CompiledRelationLocation, ...]) relation for relation in listed if _listed_key(relation) in wanted ) columns: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( - self._adapter.get_columns_for_relations(connection=self._connection, relations=matched) + get_columns_for_inspection( + adapter=self._adapter, connection=self._connection, relations=matched + ) if matched else {} ) diff --git a/src/sqlbuild/compiler/planner/main/execution/execution.py b/src/sqlbuild/compiler/planner/main/execution/execution.py index 3348cf14a0..5d236e68b1 100644 --- a/src/sqlbuild/compiler/planner/main/execution/execution.py +++ b/src/sqlbuild/compiler/planner/main/execution/execution.py @@ -7,6 +7,7 @@ from typing import Any from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.adapter.relations.main.open_inspection_catalog import open_inspection_catalog from sqlbuild.compiler.compile.models import CompiledProject from sqlbuild.compiler.planner._helpers.changes.detect import detect_changes from sqlbuild.compiler.planner._helpers.planning.buildability import ( @@ -84,12 +85,15 @@ def build_execution_plan( selection=selection, policies=policies, ) - with BackgroundSqlTestPlanning( - project=project, - adapter=adapter, - selected_keys=scopes.inspection_scope.selected_keys, - enabled=policies.plan_sql_tests, - ) as test_planning: + with ( + BackgroundSqlTestPlanning( + project=project, + adapter=adapter, + selected_keys=scopes.inspection_scope.selected_keys, + enabled=policies.plan_sql_tests, + ) as test_planning, + open_inspection_catalog(adapter=adapter, connection=connection), + ): warehouse: PlannerWarehouseState = gather_planner_warehouse_state( runtime=runtime, scopes=scopes, diff --git a/tests/integration/src/sqlbuild/compiler/planner/_helpers/test_warehouse_snapshot.py b/tests/integration/src/sqlbuild/compiler/planner/_helpers/test_warehouse_snapshot.py index 94fa2037e0..c701f5e8df 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/_helpers/test_warehouse_snapshot.py +++ b/tests/integration/src/sqlbuild/compiler/planner/_helpers/test_warehouse_snapshot.py @@ -197,7 +197,10 @@ expected_fingerprint_names=frozenset(), ), GatherWarehouseSnapshotTestCase( - description="falls back to schema-wide metadata when selected scope exceeds threshold", + description=( + "lists the whole schema beyond the name threshold but reads columns only for " + "project relations" + ), setup_sql=( *(f"CREATE TABLE staging.model_{index} (id INTEGER)" for index in range(251)), "CREATE TABLE staging.unrelated (value VARCHAR)", @@ -213,9 +216,7 @@ expected_relation_names=frozenset( {*(f"model_{index}" for index in range(251)), "unrelated"} ), - expected_column_table_names=frozenset( - {*(f"model_{index}" for index in range(251)), "unrelated"} - ), + expected_column_table_names=frozenset(f"model_{index}" for index in range(251)), expected_fingerprint_names=frozenset(), ), ], diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py b/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py index fa0f96534c..0137602668 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py @@ -101,3 +101,31 @@ class FutureCursorPlannerErrorTestCase: class TableTypePlanAssemblyTestCase: description: str expected_entry_names: tuple[str, ...] + + +@dataclass(frozen=True) +class SnowflakeInspectionBudgetTestCase: + """A synthetic multi-schema Snowflake project planned against a recording warehouse.""" + + description: str + unmanaged_relations_per_schema: int + expected_schema_reads: dict[str, int] + expected_metadata_budget: int + expected_in_list_limit: int + + +@dataclass(frozen=True) +class SnowflakeCursorBoundsBudgetTestCase: + """Cursor-bound reads issued while planning a synthetic Snowflake project.""" + + description: str + statement_latency_seconds: float + expected_max_concurrency: int + + +@dataclass(frozen=True) +class SnowflakeReplanTestCase: + """Planning the same project twice in separate invocations.""" + + description: str + expected_tables_reads: int diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py b/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py index ca25110464..50cdc792f5 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py @@ -2,6 +2,8 @@ from __future__ import annotations +import re +from collections import Counter from collections.abc import Callable from dataclasses import fields, replace from datetime import UTC, datetime @@ -24,7 +26,9 @@ FunctionArgument, ) from sqlbuild.compiler.compile.types import CompiledResourceType, FunctionLanguage +from sqlbuild.compiler.discovery.main.discover import discover_project_inputs from sqlbuild.compiler.discovery.models import ( + DiscoveredProjectInputs, DiscoveredSchemaFile, DiscoveredSeedFile, DiscoveredSourceFile, @@ -32,6 +36,8 @@ from sqlbuild.compiler.fingerprints.main.compute_query_hash import compute_query_hash from sqlbuild.compiler.fingerprints.main.write import write_fingerprint from sqlbuild.compiler.fingerprints.models import Fingerprint +from sqlbuild.compiler.pipeline.main.compile import run_compile_pipeline +from sqlbuild.compiler.pipeline.models import CompilePipelineOptions, CompilePipelineResult from sqlbuild.compiler.planner._helpers.identity.functions import ( build_compiled_function_fingerprint_sql, ) @@ -44,6 +50,7 @@ PlanOutput, ) from sqlbuild.compiler.references.types import SqlReferenceKind +from sqlbuild.runtime.contracts.models import ConnectionHooks from sqlbuild.spec.contracts.models import ( SchemaSeedEntry, SettingsConfig, @@ -54,6 +61,12 @@ FormatPlanIntegrationTestCase, SourceCursorInputPlanErrorTestCase, ) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + OfflineSnowflakeAdapter, + RecordedQuery, + RecordingSnowflakeWarehouse, + SyntheticSnowflakeProject, +) def build_project_from_test_case( @@ -447,3 +460,53 @@ def grouped(model: type) -> dict[str, Any]: policies=policies, **kwargs, ) + + +_OFFLINE_METADATA_KINDS: frozenset[str] = frozenset( + {"tables", "columns", "show_columns", "other_metadata"} +) +_CURSOR_BOUND_RELATION: re.Pattern[str] = re.compile(r"\sFROM\s+(\S+)\s*$", re.IGNORECASE) + + +def plan_offline_snowflake_project( + *, + project: SyntheticSnowflakeProject, + warehouse: RecordingSnowflakeWarehouse, + no_cache: bool = True, +) -> CompilePipelineResult: + """Compile and plan a synthetic Snowflake project against the recording warehouse.""" + + inputs: DiscoveredProjectInputs = discover_project_inputs(project_dir=project.project_dir) + return run_compile_pipeline( + discovered_inputs=inputs, + adapter=OfflineSnowflakeAdapter(warehouse=warehouse), + options=CompilePipelineOptions( + no_sql_validation=True, + no_cache=no_cache, + connection_config=dict(inputs.project_config.connection), + ), + hooks=ConnectionHooks(), + ) + + +def offline_metadata_queries(warehouse: RecordingSnowflakeWarehouse) -> tuple[RecordedQuery, ...]: + """Return every catalog metadata statement the warehouse received.""" + + return tuple(filter(lambda query: query.kind in _OFFLINE_METADATA_KINDS, warehouse.queries)) + + +def offline_metadata_reads_by_schema( + *, warehouse: RecordingSnowflakeWarehouse, kind: str +) -> dict[str, int]: + """Count schema-scoped metadata reads of one kind by their schema parameter.""" + + return dict(Counter(str(query.params[0]) for query in warehouse.queries_of_kind(kind))) + + +def offline_cursor_bound_relations(warehouse: RecordingSnowflakeWarehouse) -> tuple[str, ...]: + """Return the single relation each cursor-bound statement reads.""" + + return tuple( + _CURSOR_BOUND_RELATION.findall(query.sql)[0] + for query in warehouse.queries_of_kind("cursor_bounds") + ) diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py b/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py new file mode 100644 index 0000000000..064b6d2ed1 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py @@ -0,0 +1,155 @@ +"""Round-trip budget guard for planning a large multi-schema Snowflake project offline. + +The project is compiled and planned through the real pipeline; only the network is replaced by +a recording warehouse that answers metadata SQL from in-memory synthetic catalog rows. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from sqlbuild.adapter.relations.constants import INSPECTION_IN_LIST_LIMIT +from sqlbuild.compiler.pipeline.models import CompilePipelineResult +from tests.integration.src.sqlbuild.compiler.planner.main._test_types import ( + SnowflakeCursorBoundsBudgetTestCase, + SnowflakeInspectionBudgetTestCase, + SnowflakeReplanTestCase, +) +from tests.integration.src.sqlbuild.compiler.planner.main.helpers import ( + offline_cursor_bound_relations, + offline_metadata_queries, + offline_metadata_reads_by_schema, + plan_offline_snowflake_project, +) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + OfflineSnowflakeAdapter, + RecordedQuery, + RecordingSnowflakeWarehouse, + SyntheticSnowflakeProject, + write_synthetic_snowflake_project, +) + +_ONE_READ_PER_SCHEMA: dict[str, int] = {"STAGING": 1, "INTERMEDIATE": 1, "MARTS": 1, "RAW": 1} + + +@pytest.mark.parametrize( + "test_case", + [ + SnowflakeInspectionBudgetTestCase( + description="about 2,500 relations", + unmanaged_relations_per_schema=450, + expected_schema_reads=_ONE_READ_PER_SCHEMA, + expected_metadata_budget=8, + expected_in_list_limit=INSPECTION_IN_LIST_LIMIT, + ), + SnowflakeInspectionBudgetTestCase( + description="about 4,300 relations", + unmanaged_relations_per_schema=900, + expected_schema_reads=_ONE_READ_PER_SCHEMA, + expected_metadata_budget=8, + expected_in_list_limit=INSPECTION_IN_LIST_LIMIT, + ), + ], + ids=lambda case: case.description, +) +def test_given_large_multi_schema_project_when_planning_then_metadata_reads_stay_in_budget( + test_case: SnowflakeInspectionBudgetTestCase, tmp_path: Path +) -> None: + project: SyntheticSnowflakeProject = write_synthetic_snowflake_project( + project_dir=tmp_path / "orders_platform", + models_per_schema=60, + incremental_every=2, + sources=260, + unmanaged_relations_per_schema=test_case.unmanaged_relations_per_schema, + ) + warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( + relations=project.relations + ) + + result: CompilePipelineResult = plan_offline_snowflake_project( + project=project, warehouse=warehouse + ) + + metadata: tuple[RecordedQuery, ...] = offline_metadata_queries(warehouse) + assert len(result.plan_output.model_entries) == len(project.model_schemas) + assert ( + offline_metadata_reads_by_schema(warehouse=warehouse, kind="tables") + == test_case.expected_schema_reads + ) + assert ( + offline_metadata_reads_by_schema(warehouse=warehouse, kind="columns") + == test_case.expected_schema_reads + ) + assert warehouse.queries_of_kind("show_columns") == () + assert warehouse.queries_of_kind("other_metadata") == () + assert max(query.largest_in_list for query in metadata) <= test_case.expected_in_list_limit + assert len({(query.sql, query.params) for query in metadata}) == len(metadata) + assert len(metadata) <= test_case.expected_metadata_budget + + +@pytest.mark.parametrize( + "test_case", + [ + SnowflakeCursorBoundsBudgetTestCase( + description="one statement per relation with bounded parallelism", + statement_latency_seconds=0.002, + expected_max_concurrency=OfflineSnowflakeAdapter.metadata_inspection_concurrency, + ) + ], + ids=lambda case: case.description, +) +def test_given_incremental_models_when_planning_then_cursor_bounds_run_one_per_relation_in_parallel( + test_case: SnowflakeCursorBoundsBudgetTestCase, tmp_path: Path +) -> None: + project: SyntheticSnowflakeProject = write_synthetic_snowflake_project( + project_dir=tmp_path / "orders_platform", + models_per_schema=60, + incremental_every=2, + unmanaged_relations_per_schema=50, + ) + warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( + relations=project.relations, statement_latency_seconds=test_case.statement_latency_seconds + ) + + _ = plan_offline_snowflake_project(project=project, warehouse=warehouse) + + bounds: tuple[RecordedQuery, ...] = warehouse.queries_of_kind("cursor_bounds") + relations: tuple[str, ...] = offline_cursor_bound_relations(warehouse) + assert len(bounds) >= len(project.incremental_model_names) + assert all("UNION" not in query.sql.upper() for query in bounds) + assert len(relations) == len(bounds) == len(set(relations)) + assert 1 < warehouse.max_concurrent_cursor_bounds <= test_case.expected_max_concurrency + + +@pytest.mark.parametrize( + "test_case", + [ + SnowflakeReplanTestCase( + description="a second invocation re-reads bounds and listings", expected_tables_reads=4 + ) + ], + ids=lambda case: case.description, +) +def test_given_planned_project_when_planning_again_then_cursor_bounds_are_read_again( + test_case: SnowflakeReplanTestCase, tmp_path: Path +) -> None: + project: SyntheticSnowflakeProject = write_synthetic_snowflake_project( + project_dir=tmp_path / "orders_platform", unmanaged_relations_per_schema=10 + ) + warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( + relations=project.relations + ) + _ = plan_offline_snowflake_project(project=project, warehouse=warehouse) + first: int = len(warehouse.queries_of_kind("cursor_bounds")) + warehouse.reset() + + _ = plan_offline_snowflake_project(project=project, warehouse=warehouse, no_cache=False) + + assert len(warehouse.queries_of_kind("cursor_bounds")) == first + assert len(warehouse.queries_of_kind("tables")) == test_case.expected_tables_reads + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapter/relations/_test_types.py b/tests/unit/src/sqlbuild/adapter/relations/_test_types.py new file mode 100644 index 0000000000..456b2f5bd4 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/_test_types.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class MetadataSqlOwnershipTestCase: + """Source modules allowed to contain warehouse catalog SQL.""" + + description: str + allowed_modules: frozenset[str] + expected_unlisted_modules: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/adapter/relations/helpers.py b/tests/unit/src/sqlbuild/adapter/relations/helpers.py new file mode 100644 index 0000000000..a23b468735 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/helpers.py @@ -0,0 +1,26 @@ +"""Helpers for adapter relation architecture guards.""" + +from __future__ import annotations + +import re +from pathlib import Path + +SOURCE_ROOT: Path = Path(__file__).resolve().parents[6] / "src" / "sqlbuild" +_CATALOG_SQL_PATTERNS: tuple[re.Pattern[str], ...] = ( + re.compile(r"information_schema", re.IGNORECASE), + re.compile(r"[\"']\s*SHOW\s+[A-Z]"), +) + + +def metadata_sql_modules() -> frozenset[str]: + """Return source modules containing INFORMATION_SCHEMA or SHOW statement text.""" + + return frozenset( + path.relative_to(SOURCE_ROOT).as_posix() + for path in filter(_contains_catalog_sql, SOURCE_ROOT.rglob("*.py")) + ) + + +def _contains_catalog_sql(path: Path) -> bool: + text: str = path.read_text(encoding="utf-8") + return any(pattern.search(text) is not None for pattern in _CATALOG_SQL_PATTERNS) diff --git a/tests/unit/src/sqlbuild/adapter/relations/test_metadata_sql_ownership.py b/tests/unit/src/sqlbuild/adapter/relations/test_metadata_sql_ownership.py new file mode 100644 index 0000000000..ee36d55226 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/test_metadata_sql_ownership.py @@ -0,0 +1,51 @@ +"""Architecture guard: warehouse catalog SQL lives only in adapter inspection modules. + +Planner and executor code must reach metadata through adapter methods and the shared inspection +catalog, so query shapes, IN-list caps, and per-schema read budgets stay enforceable in one place. +""" + +from __future__ import annotations + +import pytest + +from tests.unit.src.sqlbuild.adapter.relations._test_types import MetadataSqlOwnershipTestCase +from tests.unit.src.sqlbuild.adapter.relations.helpers import metadata_sql_modules + +_ALLOWED_METADATA_SQL_MODULES: frozenset[str] = frozenset( + { + "adapter/contract/classes/base_adapter.py", + "adapter/contract/classes/duckdb_backed_adapter.py", + "adapters/bigquery/classes/bigquery_adapter.py", + "adapters/databricks/classes/databricks_adapter.py", + "adapters/postgres/classes/postgres_adapter.py", + "adapters/snowflake/classes/snowflake_adapter.py", + "adapters/sqlserver/classes/sqlserver_adapter.py", + "adapters/sqlserver/constants.py", + "cost/_helpers/snowflake.py", + "cost/models.py", + } +) + + +@pytest.mark.parametrize( + "test_case", + [ + MetadataSqlOwnershipTestCase( + description="catalog SQL appears only in adapter inspection modules", + allowed_modules=_ALLOWED_METADATA_SQL_MODULES, + expected_unlisted_modules=(), + ) + ], + ids=lambda case: case.description, +) +def test_given_source_tree_when_scanning_then_metadata_sql_stays_in_adapter_modules( + test_case: MetadataSqlOwnershipTestCase, +) -> None: + found: frozenset[str] = metadata_sql_modules() + + assert tuple(sorted(found - test_case.allowed_modules)) == test_case.expected_unlisted_modules + assert tuple(sorted(test_case.allowed_modules - found)) == test_case.expected_unlisted_modules + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py index 9ed9362053..57accfef1e 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py @@ -8,3 +8,20 @@ class SemanticSourceBatchCase: description: str expected_batches: int expected_sources: frozenset[str] + + +@dataclass(frozen=True) +class ConcurrentCursorBoundsTestCase: + description: str + relation_count: int + concurrency: int + failing_relations: frozenset[str] + expected_failure_relations: tuple[str, ...] + expected_start_progress: str + + +@dataclass(frozen=True) +class EligibleTargetMaxFailureTestCase: + description: str + failing_relation: str + expected_error_fragment: str diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py new file mode 100644 index 0000000000..40ec4e109f --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py @@ -0,0 +1,145 @@ +"""Helpers for warehouse snapshot planner tests.""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Callable +from typing import Any + +from sqlbuild.compiler.planner._helpers.warehouse.snapshot import ( + _CursorModelInfo, + _PhysicalCursorQuery, +) +from sqlbuild.compiler.planner.types import CursorType +from sqlbuild.spec.contracts.models import StartCursorsConfig + + +class CursorBoundRows: + """Fetchable result of one cursor-bound statement double.""" + + def __init__(self, rows: list[tuple[str | None, ...]]) -> None: + self._rows: list[tuple[str | None, ...]] = rows + + def fetchall(self) -> list[tuple[str | None, ...]]: + return self._rows + + +def _answer_cursor_relation(relation: str) -> CursorBoundRows: + day: int = int(relation.rsplit("_", 1)[1]) % 28 + 1 + return CursorBoundRows([(f"2026-01-{day:02d}", f"2026-02-{day:02d}")]) + + +def _reject_cursor_relation(relation: str) -> CursorBoundRows: + raise RuntimeError(f"Object '{relation}' does not exist or not authorized.") + + +_CURSOR_RELATION_ANSWERS: dict[bool, Callable[[str], CursorBoundRows]] = { + False: _answer_cursor_relation, + True: _reject_cursor_relation, +} + + +class ConcurrencyTrackingCursorExecute: + """Execute double answering one MIN/MAX statement per relation while counting overlap.""" + + def __init__(self, *, failing_relations: frozenset[str], latency_seconds: float) -> None: + self._failing_relations: frozenset[str] = failing_relations + self._latency_seconds: float = latency_seconds + self._lock: threading.Lock = threading.Lock() + self._active: int = 0 + self.max_active: int = 0 + self.sql: list[str] = [] + + def __call__(self, *, connection: Any, sql: str) -> CursorBoundRows: + del connection + with self._lock: + self._active += 1 + self.max_active = max(self.max_active, self._active) + self.sql.append(sql) + try: + time.sleep(self._latency_seconds) + relation: str = sql.rsplit(" FROM ", 1)[1] + return _CURSOR_RELATION_ANSWERS[relation in self._failing_relations](relation) + finally: + with self._lock: + self._active -= 1 + + +def build_cursor_bound_queries(count: int) -> list[_PhysicalCursorQuery]: + """Return one MIN/MAX physical cursor query per synthetic orders relation.""" + + return [ + _PhysicalCursorQuery( + relation=f"analytics.raw.orders_{index}", + cursor_column="ordered_at", + min_tags=(f"model_{index}__orders__min",), + max_tags=(f"model_{index}__orders__max",), + ) + for index in range(count) + ] + + +class EligibleMaxAdapter: + """Adapter double rendering bounded MAX statements and allowing parallel reads.""" + + metadata_inspection_concurrency: int = 4 + + def render_max_cursor_at_or_before( + self, + *, + relation: str, + cursor_column: str, + maximum_allowed: str, + cursor_type: str, + is_date: bool, + ) -> str: + del cursor_type, is_date + return ( + f"SELECT CAST(MAX({cursor_column}) AS VARCHAR) FROM {relation} " + f"WHERE {cursor_column} <= '{maximum_allowed}'" + ) + + +def _answer_eligible_max(relation: str) -> CursorBoundRows: + del relation + return CursorBoundRows([("2026-01-15",)]) + + +_ELIGIBLE_MAX_ANSWERS: dict[bool, Callable[[str], CursorBoundRows]] = { + False: _answer_eligible_max, + True: _reject_cursor_relation, +} + + +class EligibleMaxExecute: + """Execute double failing the bounded MAX read of one relation.""" + + def __init__(self, *, failing_relation: str, latency_seconds: float) -> None: + self._failing_relation: str = failing_relation + self._latency_seconds: float = latency_seconds + + def __call__(self, *, connection: Any, sql: str) -> CursorBoundRows: + del connection + time.sleep(self._latency_seconds) + relation: str = sql.split(" FROM ", 1)[1].split(" ", 1)[0] + return _ELIGIBLE_MAX_ANSWERS[relation == self._failing_relation](relation) + + +def build_eligible_max_cursor_models(count: int) -> list[_CursorModelInfo]: + """Return timestamp cursor models whose target MAX is beyond a one-day horizon.""" + + return [ + _CursorModelInfo( + model_name=f"model_{index}", + target_tag=f"model_{index}__target__max", + target_relation=f"analytics.marts.orders_{index}", + cursor_column="ordered_at", + upstreams=(), + cursor_type=CursorType.TIMESTAMP, + cursor_grain="day", + effective_cursor_grain="day", + start_cursor_config=StartCursorsConfig(max_ahead="1d"), + ) + for index in range(count) + ] diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_warehouse_cursor_concurrency.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_warehouse_cursor_concurrency.py new file mode 100644 index 0000000000..218993e9f6 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_warehouse_cursor_concurrency.py @@ -0,0 +1,139 @@ +"""Bounded-parallel cursor-bound reads must match sequential reads exactly.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any, cast + +import pytest + +from sqlbuild.adapter.relations.main.open_inspection_catalog import open_inspection_catalog +from sqlbuild.compiler.planner._helpers.warehouse.snapshot import ( + _CursorModelInfo, + _execute_cursor_queries, + _gather_eligible_target_maxes, + _PhysicalCursorQuery, +) +from sqlbuild.compiler.planner.exceptions import PlannerInputError +from sqlbuild.compiler.planner.types import CursorType +from sqlbuild.cursor_algebra.main.parse import parse +from sqlbuild.cursor_algebra.types import CursorScalar +from tests.unit.src.sqlbuild.compiler.planner._helpers.warehouse._test_types import ( + ConcurrentCursorBoundsTestCase, + EligibleTargetMaxFailureTestCase, +) +from tests.unit.src.sqlbuild.compiler.planner._helpers.warehouse.helpers import ( + ConcurrencyTrackingCursorExecute, + EligibleMaxAdapter, + EligibleMaxExecute, + build_cursor_bound_queries, + build_eligible_max_cursor_models, +) + +_LATENCY_SECONDS: float = 0.01 + + +@pytest.mark.parametrize( + "test_case", + [ + ConcurrentCursorBoundsTestCase( + description="parallel reads stay within the bound", + relation_count=24, + concurrency=4, + failing_relations=frozenset(), + expected_failure_relations=(), + expected_start_progress="Inspecting cursor bounds for 24 relations (4 concurrent)...", + ), + ConcurrentCursorBoundsTestCase( + description="one failed relation stays unavailable and others succeed", + relation_count=12, + concurrency=8, + failing_relations=frozenset({"analytics.raw.orders_5"}), + expected_failure_relations=("analytics.raw.orders_5",), + expected_start_progress="Inspecting cursor bounds for 12 relations (8 concurrent)...", + ), + ], + ids=lambda case: case.description, +) +def test_given_concurrency_when_reading_cursor_bounds_then_matches_sequential_results( + test_case: ConcurrentCursorBoundsTestCase, +) -> None: + queries: list[_PhysicalCursorQuery] = build_cursor_bound_queries(test_case.relation_count) + sequential_execute: ConcurrencyTrackingCursorExecute = ConcurrencyTrackingCursorExecute( + failing_relations=test_case.failing_relations, latency_seconds=_LATENCY_SECONDS + ) + sequential: dict[str, CursorScalar] = _execute_cursor_queries( + queries=queries, connection=None, execute=sequential_execute, on_progress=None + ) + parallel_execute: ConcurrencyTrackingCursorExecute = ConcurrencyTrackingCursorExecute( + failing_relations=test_case.failing_relations, latency_seconds=_LATENCY_SECONDS + ) + progress: list[str] = [] + + parallel: dict[str, CursorScalar] = _execute_cursor_queries( + queries=queries, + connection=None, + execute=parallel_execute, + on_progress=progress.append, + concurrency=test_case.concurrency, + ) + + failures: tuple[str, ...] = tuple( + filter(lambda message: message.startswith("Failed"), progress) + ) + assert list(parallel.items()) == list(sequential.items()) + assert sequential_execute.max_active == 1 + assert 1 < parallel_execute.max_active <= test_case.concurrency + assert sorted(parallel_execute.sql) == sorted(sequential_execute.sql) + assert len(parallel_execute.sql) == test_case.relation_count + assert progress[0] == test_case.expected_start_progress + assert len(progress) == test_case.relation_count + 1 + assert ( + tuple(message.split(": ")[1].split(".ordered_at")[0] for message in failures) + == test_case.expected_failure_relations + ) + + +@pytest.mark.parametrize( + "test_case", + [ + EligibleTargetMaxFailureTestCase( + description="failure names the model and relation", + failing_relation="analytics.marts.orders_3", + expected_error_fragment=( + "model 'model_3': failed to query highest eligible target cursor for " + "analytics.marts.orders_3.ordered_at" + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_failing_eligible_max_read_when_reading_in_parallel_then_plan_fails_naming_relation( + test_case: EligibleTargetMaxFailureTestCase, +) -> None: + infos: list[_CursorModelInfo] = build_eligible_max_cursor_models(6) + results: dict[str, CursorScalar] = { + f"model_{index}__target__max": parse(raw="2099-01-01", cursor_type=CursorType.TIMESTAMP) + for index in range(6) + } + adapter: EligibleMaxAdapter = EligibleMaxAdapter() + connection: object = object() + + with ( + open_inspection_catalog(adapter=adapter, connection=connection), + pytest.raises(PlannerInputError, match=test_case.expected_error_fragment), + ): + _ = _gather_eligible_target_maxes( + cursor_models=infos, + results=results, + adapter=cast(Any, adapter), + connection=connection, + execute=EligibleMaxExecute( + failing_relation=test_case.failing_relation, latency_seconds=_LATENCY_SECONDS + ), + invocation_time=datetime(2026, 1, 20, tzinfo=UTC), + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) From 764f07e5baf721ec260bed40a41a9289d67308ee Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 11:35:59 +0100 Subject: [PATCH 3/7] feat(cli): log warehouse inspection queries under plan and build -v With -v, sqb plan and sqb build write each planning inspection read (truncated SQL, elapsed time, row count or error) and a final total to stderr, keeping JSON stdout machine readable. Refs CHI-455 --- .../_helpers/build_planning/planning.py | 8 ++ .../cli/commands/main/project/_plan.py | 8 +- .../classes/verbose_inspection_writer.py | 63 ++++++++++ .../main/verbose_inspection_output.py | 22 ++++ .../cli/commands/main/plan/_test_types.py | 8 ++ .../main/plan/test_verbose_inspection.py | 97 +++++++++++++++ .../verbose_inspection_output/_test_types.py | 25 ++++ .../test_verbose_inspection_output.py | 113 ++++++++++++++++++ 8 files changed, 340 insertions(+), 4 deletions(-) create mode 100644 src/sqlbuild/cli/progress/classes/verbose_inspection_writer.py create mode 100644 src/sqlbuild/cli/progress/main/verbose_inspection_output.py create mode 100644 tests/e2e/src/sqlbuild/cli/commands/main/plan/test_verbose_inspection.py create mode 100644 tests/unit/src/sqlbuild/cli/progress/main/verbose_inspection_output/_test_types.py create mode 100644 tests/unit/src/sqlbuild/cli/progress/main/verbose_inspection_output/test_verbose_inspection_output.py diff --git a/src/sqlbuild/cli/commands/_helpers/build_planning/planning.py b/src/sqlbuild/cli/commands/_helpers/build_planning/planning.py index 6d00103e64..49e05cb74d 100644 --- a/src/sqlbuild/cli/commands/_helpers/build_planning/planning.py +++ b/src/sqlbuild/cli/commands/_helpers/build_planning/planning.py @@ -9,6 +9,7 @@ BuildCommandRequest, BuildInvocation, ) +from sqlbuild.cli.progress.main.verbose_inspection_output import verbose_inspection_output from sqlbuild.compiler.pipeline.main.compile import run_compile_pipeline from sqlbuild.compiler.pipeline.models import ( CompilePipelineOptions, @@ -26,6 +27,13 @@ def compile_build_plan( invocation.progress_stream.write("\n") invocation.progress_stream.flush() + with verbose_inspection_output(enabled=request.verbose): + return _run_build_compile_pipeline(request=request, invocation=invocation) + + +def _run_build_compile_pipeline( + *, request: BuildCommandRequest, invocation: BuildInvocation +) -> CompilePipelineResult: return run_compile_pipeline( discovered_inputs=invocation.discovered_inputs, adapter=invocation.adapter, diff --git a/src/sqlbuild/cli/commands/main/project/_plan.py b/src/sqlbuild/cli/commands/main/project/_plan.py index 917724e117..6f47ec724c 100644 --- a/src/sqlbuild/cli/commands/main/project/_plan.py +++ b/src/sqlbuild/cli/commands/main/project/_plan.py @@ -6,6 +6,7 @@ from sqlbuild.cli.commands._helpers.plan.outputs import write_plan_command_output from sqlbuild.cli.commands._helpers.plan.planning import compile_plan_pipeline from sqlbuild.cli.commands.models import PlanCommandRequest, PlanInvocation +from sqlbuild.cli.progress.main.verbose_inspection_output import verbose_inspection_output from sqlbuild.compiler.pipeline.models import CompilePipelineResult @@ -16,10 +17,9 @@ def run_plan(request: PlanCommandRequest) -> int: if not request.json_output: invocation.progress_stream.write("\n") invocation.progress_stream.flush() - pipeline_result: CompilePipelineResult = compile_plan_pipeline( - request=request, - invocation=invocation, - ) + pipeline_result: CompilePipelineResult + with verbose_inspection_output(enabled=request.verbose): + pipeline_result = compile_plan_pipeline(request=request, invocation=invocation) write_plan_command_output( request=request, invocation=invocation, diff --git a/src/sqlbuild/cli/progress/classes/verbose_inspection_writer.py b/src/sqlbuild/cli/progress/classes/verbose_inspection_writer.py new file mode 100644 index 0000000000..67409dd694 --- /dev/null +++ b/src/sqlbuild/cli/progress/classes/verbose_inspection_writer.py @@ -0,0 +1,63 @@ +"""Verbose per-query warehouse inspection lines written to stderr.""" + +from __future__ import annotations + +import sys +import threading + +from sqlbuild.adapter.relations.models import InspectionQueryRecord +from sqlbuild.presentation.main.transient_line_coordinator import shared_transient_line_coordinator + +_SQL_PREVIEW_LENGTH: int = 160 + + +class VerboseInspectionWriter: + """Write each inspection read as one truncated stderr line and keep a running total.""" + + def __init__(self) -> None: + self._lock: threading.Lock = threading.Lock() + self._count: int = 0 + self._seconds: float = 0.0 + + def write(self, record: InspectionQueryRecord) -> None: + """Write one inspection read; safe to call from inspection worker threads.""" + + with self._lock: + self._count += 1 + self._seconds += record.elapsed_seconds + self._write_line(text=self.format_record(record=record)) + + def write_total(self) -> None: + """Write the read count and cumulative query time when anything was read.""" + + if self._count: + self._write_line( + text=( + f"Inspection queries: {self._count} " + f"({self._seconds:.2f}s cumulative query time)" + ) + ) + + @staticmethod + def format_record(*, record: InspectionQueryRecord) -> str: + """Render one inspection read as a single truncated diagnostic line.""" + + text: str = " ".join(record.sql.split()) + if len(text) > _SQL_PREVIEW_LENGTH: + text = text[: _SQL_PREVIEW_LENGTH - 3] + "..." + outcome: str = ( + f"failed: {' '.join(record.error.split())}" + if record.error is not None + else _row_count_text(record.row_count) + ) + return f" inspect {record.elapsed_seconds:6.2f}s {outcome} {text}" + + @staticmethod + def _write_line(*, text: str) -> None: + shared_transient_line_coordinator().write_persistent(stream=sys.stderr, text=f"{text}\n") + + +def _row_count_text(row_count: int | None) -> str: + if row_count is None: + return "? rows" + return f"{row_count} row" if row_count == 1 else f"{row_count} rows" diff --git a/src/sqlbuild/cli/progress/main/verbose_inspection_output.py b/src/sqlbuild/cli/progress/main/verbose_inspection_output.py new file mode 100644 index 0000000000..d3eafe222e --- /dev/null +++ b/src/sqlbuild/cli/progress/main/verbose_inspection_output.py @@ -0,0 +1,22 @@ +"""Verbose per-query warehouse inspection output on stderr.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager + +from sqlbuild.adapter.relations.main.inspection_query_sink import inspection_query_sink +from sqlbuild.cli.progress.classes.verbose_inspection_writer import VerboseInspectionWriter + + +@contextmanager +def verbose_inspection_output(*, enabled: bool) -> Iterator[None]: + """Write each inspection read, then a total, to stderr while ``enabled``.""" + + if not enabled: + yield + return + writer: VerboseInspectionWriter = VerboseInspectionWriter() + with inspection_query_sink(writer.write): + yield + writer.write_total() diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/plan/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/plan/_test_types.py index 76027ea038..f79268a235 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/plan/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/plan/_test_types.py @@ -46,3 +46,11 @@ class DiamondPlanE2ETestCase: class DiamondPlanJsonE2ETestCase: description: str expected_reasons: dict[str, str] + + +@dataclass(frozen=True) +class VerboseInspectionPlanE2ETestCase: + description: str + command: tuple[str, ...] + expected_inspection_output: bool + expected_json_stdout: bool diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/plan/test_verbose_inspection.py b/tests/e2e/src/sqlbuild/cli/commands/main/plan/test_verbose_inspection.py new file mode 100644 index 0000000000..5376a0ac92 --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/plan/test_verbose_inspection.py @@ -0,0 +1,97 @@ +"""E2E coverage for verbose warehouse inspection query output during plan.""" + +from __future__ import annotations + +import json +import re +import subprocess +from pathlib import Path + +import pytest + +from tests.e2e.src.sqlbuild.cli.commands.main.plan._test_types import ( + VerboseInspectionPlanE2ETestCase, +) +from tests.e2e.src.sqlbuild.cli.commands.main.plan.helpers import ( + prepare_changed_incremental_diamond_project, +) +from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import run_sqb + +_INSPECTION_LINE: re.Pattern[str] = re.compile(r"^ inspect +\d+\.\d{2}s \d+ rows? \S", re.M) +_METADATA_LINE: re.Pattern[str] = re.compile(r"^ inspect .*list_relations\(", re.M) +_CURSOR_LINE: re.Pattern[str] = re.compile(r"^ inspect .*SELECT CAST\(MAX\(", re.M) +_TOTAL_LINE: re.Pattern[str] = re.compile(r"^Inspection queries: (\d+) \(", re.M) + + +@pytest.mark.parametrize( + "test_case", + [ + VerboseInspectionPlanE2ETestCase( + description="verbose plan logs each inspection read to stderr", + command=("--no-color", "plan", "-v"), + expected_inspection_output=True, + expected_json_stdout=False, + ), + VerboseInspectionPlanE2ETestCase( + description="verbose json plan keeps stdout machine readable", + command=("--no-color", "plan", "-v", "--json"), + expected_inspection_output=True, + expected_json_stdout=True, + ), + VerboseInspectionPlanE2ETestCase( + description="default plan omits inspection reads", + command=("--no-color", "plan"), + expected_inspection_output=False, + expected_json_stdout=False, + ), + ], + ids=lambda case: case.description, +) +def test_given_built_project_when_planning_then_verbose_reports_inspection_queries( + test_case: VerboseInspectionPlanE2ETestCase, tmp_path: Path +) -> None: + project_dir: Path = prepare_changed_incremental_diamond_project(tmp_path=tmp_path) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=test_case.command, project_dir=project_dir + ) + + assert result.returncode == 0, result.stdout + result.stderr + reads: list[str] = _INSPECTION_LINE.findall(result.stderr) + totals: list[str] = _TOTAL_LINE.findall(result.stderr) + assert _INSPECTION_LINE.search(result.stdout) is None + assert bool(reads) is test_case.expected_inspection_output + assert bool(_METADATA_LINE.search(result.stderr)) is test_case.expected_inspection_output + assert bool(_CURSOR_LINE.search(result.stderr)) is test_case.expected_inspection_output + assert totals == [str(len(reads))] * test_case.expected_inspection_output + assert result.stdout.lstrip().startswith("{") is test_case.expected_json_stdout + + +@pytest.mark.parametrize( + "test_case", + [ + VerboseInspectionPlanE2ETestCase( + description="verbose json plan stdout parses as one json document", + command=("--no-color", "plan", "-v", "--json"), + expected_inspection_output=True, + expected_json_stdout=True, + ) + ], + ids=lambda case: case.description, +) +def test_given_verbose_json_plan_when_planning_then_stdout_parses_as_json( + test_case: VerboseInspectionPlanE2ETestCase, tmp_path: Path +) -> None: + project_dir: Path = prepare_changed_incremental_diamond_project(tmp_path=tmp_path) + + result: subprocess.CompletedProcess[str] = run_sqb( + command=test_case.command, project_dir=project_dir + ) + + assert result.returncode == 0, result.stdout + result.stderr + assert isinstance(json.loads(result.stdout), dict) is test_case.expected_json_stdout + assert bool(_TOTAL_LINE.search(result.stderr)) is test_case.expected_inspection_output + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/cli/progress/main/verbose_inspection_output/_test_types.py b/tests/unit/src/sqlbuild/cli/progress/main/verbose_inspection_output/_test_types.py new file mode 100644 index 0000000000..2986b6d893 --- /dev/null +++ b/tests/unit/src/sqlbuild/cli/progress/main/verbose_inspection_output/_test_types.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class VerboseInspectionFormatTestCase: + """One recorded inspection read rendered as a verbose line.""" + + description: str + sql: str + row_count: int | None + error: str | None + expected_fragments: tuple[str, ...] + expected_maximum_length: int + + +@dataclass(frozen=True) +class VerboseInspectionOutputTestCase: + """Inspection reads recorded inside and after a verbose output context.""" + + description: str + enabled: bool + records: tuple[tuple[str, float, int], ...] + expected_stderr_lines: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/cli/progress/main/verbose_inspection_output/test_verbose_inspection_output.py b/tests/unit/src/sqlbuild/cli/progress/main/verbose_inspection_output/test_verbose_inspection_output.py new file mode 100644 index 0000000000..5ae0078cb1 --- /dev/null +++ b/tests/unit/src/sqlbuild/cli/progress/main/verbose_inspection_output/test_verbose_inspection_output.py @@ -0,0 +1,113 @@ +"""Verbose inspection lines are single-line, truncated, and go to stderr only.""" + +from __future__ import annotations + +import pytest +from _pytest.capture import CaptureResult + +from sqlbuild.adapter.relations.main.record_inspection_query import record_inspection_query +from sqlbuild.adapter.relations.models import InspectionQueryRecord +from sqlbuild.cli.progress.classes.verbose_inspection_writer import VerboseInspectionWriter +from sqlbuild.cli.progress.main.verbose_inspection_output import verbose_inspection_output +from tests.unit.src.sqlbuild.cli.progress.main.verbose_inspection_output._test_types import ( + VerboseInspectionFormatTestCase, + VerboseInspectionOutputTestCase, +) + +_LONG_SQL: str = ( + 'SELECT table_name, column_name, data_type FROM "ANALYTICS".information_schema.columns\n' + "WHERE table_schema = 'STAGING' AND table_catalog = 'ANALYTICS' " + + "ORDER BY table_name, ordinal_position " + * 8 +) + + +@pytest.mark.parametrize( + "test_case", + [ + VerboseInspectionFormatTestCase( + description="long multi-line SQL is collapsed and truncated", + sql=_LONG_SQL, + row_count=652, + error=None, + expected_fragments=("0.42s", "652 rows", "information_schema.columns WHERE", "..."), + expected_maximum_length=200, + ), + VerboseInspectionFormatTestCase( + description="single row is singular", + sql="SELECT CAST(MAX(ordered_at) AS VARCHAR) AS _max FROM analytics.marts.orders", + row_count=1, + error=None, + expected_fragments=(" 1 row ", "FROM analytics.marts.orders"), + expected_maximum_length=200, + ), + VerboseInspectionFormatTestCase( + description="failure shows the error instead of a row count", + sql="SHOW COLUMNS IN TABLE analytics.raw.orders", + row_count=None, + error="Object does not exist\nor not authorized", + expected_fragments=("failed: Object does not exist or not authorized", "SHOW COLUMNS"), + expected_maximum_length=200, + ), + ], + ids=lambda case: case.description, +) +def test_given_inspection_record_when_formatting_then_renders_one_bounded_line( + test_case: VerboseInspectionFormatTestCase, +) -> None: + line: str = VerboseInspectionWriter.format_record( + record=InspectionQueryRecord( + sql=test_case.sql, + elapsed_seconds=0.42, + row_count=test_case.row_count, + error=test_case.error, + ) + ) + + assert "\n" not in line + assert len(line) <= test_case.expected_maximum_length + for fragment in test_case.expected_fragments: + assert fragment in line + + +@pytest.mark.parametrize( + "test_case", + [ + VerboseInspectionOutputTestCase( + description="enabled output writes each read and a total to stderr", + enabled=True, + records=(("SELECT 1", 0.5, 3), ("SELECT 2", 0.25, 0)), + expected_stderr_lines=( + " inspect 0.50s 3 rows SELECT 1", + " inspect 0.25s 0 rows SELECT 2", + "Inspection queries: 2 (0.75s cumulative query time)", + ), + ), + VerboseInspectionOutputTestCase( + description="disabled output writes nothing", + enabled=False, + records=(("SELECT 1", 0.5, 3),), + expected_stderr_lines=(), + ), + ], + ids=lambda case: case.description, +) +def test_given_verbose_setting_when_recording_reads_then_writes_only_inside_context_to_stderr( + test_case: VerboseInspectionOutputTestCase, capsys: pytest.CaptureFixture[str] +) -> None: + with verbose_inspection_output(enabled=test_case.enabled): + for sql, elapsed, rows in test_case.records: + record_inspection_query( + record=InspectionQueryRecord(sql=sql, elapsed_seconds=elapsed, row_count=rows) + ) + record_inspection_query( + record=InspectionQueryRecord(sql="SELECT 3", elapsed_seconds=0.1, row_count=0) + ) + + captured: CaptureResult[str] = capsys.readouterr() + assert captured.out == "" + assert tuple(captured.err.splitlines()) == test_case.expected_stderr_lines + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) From 9417adf99f9f1019e3d240e4be43c94f698aeb01 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 12:28:28 +0100 Subject: [PATCH 4/7] perf(snowflake): inspect planning metadata with SHOW instead of INFORMATION_SCHEMA INFORMATION_SCHEMA queries need a running warehouse and queue behind other work on a shared warehouse; SHOW commands are served by the metadata service. Planning schema listings now use SHOW TABLES and SHOW VIEWS (paged past the 10,000-row cap) and SHOW COLUMNS IN SCHEMA, mapped onto the same RelationInfo and type strings the INFORMATION_SCHEMA path produced. A missing or unauthorized schema lists nothing, as before. SHOW COLUMNS silently truncates at 10,000 rows, so a schema that reaches the cap is reread from INFORMATION_SCHEMA.COLUMNS. SHOW has no LAST_ALTERED, so source freshness reads it only for observed sources, in IN lists of at most 50 names run with bounded concurrency. SHOW COLUMNS type translation lives in one helper shared with the per-relation path and is checked against INFORMATION_SCHEMA output for every Snowflake type family. The budget guard now requires SHOW reads per schema and no INFORMATION_SCHEMA reads beyond freshness. Refs CHI-455 --- .../docs/concepts/adapters/snowflake.md | 12 + .../relations/classes/inspection_catalog.py | 7 - src/sqlbuild/adapter/relations/models.py | 1 - .../snowflake/_helpers/metadata_types.py | 59 ++++ .../snowflake/_helpers/show_metadata.py | 95 +++++++ .../snowflake/classes/snowflake_adapter.py | 265 +++++++++++------- .../compiler/planner/main/_test_types.py | 3 +- .../sqlbuild/compiler/planner/main/helpers.py | 30 +- .../main/test_snowflake_inspection_budget.py | 28 +- .../relations/test_metadata_sql_ownership.py | 2 + .../snowflake/inspection/_test_types.py | 51 ++++ .../adapters/snowflake/inspection/helpers.py | 256 +++++++++++++++-- .../inspection/test_inspection_catalog.py | 83 +++++- .../inspection/test_metadata_types.py | 239 ++++++++++++++++ .../inspection/test_show_metadata.py | 195 +++++++++++++ .../docs/docs/concepts/adapters/snowflake.mdx | 12 + 16 files changed, 1186 insertions(+), 152 deletions(-) create mode 100644 src/sqlbuild/adapters/snowflake/_helpers/metadata_types.py create mode 100644 src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py create mode 100644 tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_metadata_types.py create mode 100644 tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py diff --git a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/adapters/snowflake.md b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/adapters/snowflake.md index 9bd6b73918..25a09a2f4b 100644 --- a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/adapters/snowflake.md +++ b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/adapters/snowflake.md @@ -41,6 +41,18 @@ On connect, SQLBuild applies the connection's role and warehouse, then applies t target's authoritative database and schema. These ensure the session context is set correctly regardless of the user's default settings. +## Planning metadata + +Planning reads relation and column metadata with `SHOW TABLES`, `SHOW VIEWS`, and `SHOW COLUMNS` +once per schema, several schemas at a time. `SHOW` commands are served by Snowflake's metadata +service, so they don't wait for, or resume, a warehouse. Two reads still use `INFORMATION_SCHEMA` +and need a running warehouse: + +- `LAST_ALTERED` for source freshness, read only for the sources the plan observes. +- Columns for a schema whose `SHOW COLUMNS` output reaches Snowflake's 10,000-row limit. + +Run `sqb plan -v` to see each metadata query with its elapsed time and row count. + ## Shared connections across targets Use multiple targets to reuse one Snowflake connection while selecting different authoritative diff --git a/src/sqlbuild/adapter/relations/classes/inspection_catalog.py b/src/sqlbuild/adapter/relations/classes/inspection_catalog.py index 6a80deb62a..63dab91de2 100644 --- a/src/sqlbuild/adapter/relations/classes/inspection_catalog.py +++ b/src/sqlbuild/adapter/relations/classes/inspection_catalog.py @@ -46,13 +46,6 @@ def concurrency(self) -> int: return self._concurrency - def cached_relation_listing( - self, *, database: str | None, schema: str - ) -> SchemaRelationListing | None: - """Return a schema listing already read in this invocation, without reading.""" - - return self._relation_listings.get(self._schema_key(database=database, schema=schema)) - def list_relations( self, *, diff --git a/src/sqlbuild/adapter/relations/models.py b/src/sqlbuild/adapter/relations/models.py index 535f003c3c..01abc1b699 100644 --- a/src/sqlbuild/adapter/relations/models.py +++ b/src/sqlbuild/adapter/relations/models.py @@ -22,7 +22,6 @@ class SchemaRelationListing: database: str | None schema: str entries: tuple[ListedRelation, ...] - rows: tuple[tuple[object, ...], ...] = () @dataclass(frozen=True) diff --git a/src/sqlbuild/adapters/snowflake/_helpers/metadata_types.py b/src/sqlbuild/adapters/snowflake/_helpers/metadata_types.py new file mode 100644 index 0000000000..2d1846bc3c --- /dev/null +++ b/src/sqlbuild/adapters/snowflake/_helpers/metadata_types.py @@ -0,0 +1,59 @@ +"""Snowflake column type strings from INFORMATION_SCHEMA fields and SHOW COLUMNS metadata.""" + +from __future__ import annotations + +import json + +from sqlbuild.adapter.contract.exceptions import AdapterUserError +from sqlbuild.adapters.snowflake.constants import NUMBER_TYPE_NAME, TEXT_TYPE_NAMES + +_SHOW_TYPE_TO_INFORMATION_SCHEMA_TYPE: dict[str, str] = { + "FIXED": NUMBER_TYPE_NAME, + "REAL": "FLOAT", +} +_INVALID_SHOW_TYPE_MESSAGE: str = "Snowflake SHOW COLUMNS returned invalid type metadata" + + +def information_schema_type( + *, + data_type: str, + numeric_precision: object, + numeric_scale: object, + character_maximum_length: object, +) -> str: + """Render the planner type string for one INFORMATION_SCHEMA.COLUMNS row.""" + + normalized_type: str = data_type.upper() + if ( + normalized_type == NUMBER_TYPE_NAME + and isinstance(numeric_precision, int) + and isinstance(numeric_scale, int) + ): + return f"{NUMBER_TYPE_NAME}({numeric_precision},{numeric_scale})" + if normalized_type in TEXT_TYPE_NAMES and isinstance(character_maximum_length, int): + return f"VARCHAR({character_maximum_length})" + return normalized_type + + +def show_columns_type(raw_data_type: object) -> str: + """Render the same type string as INFORMATION_SCHEMA for one SHOW COLUMNS data_type JSON.""" + + try: + decoded: object = json.loads(str(raw_data_type)) + except (json.JSONDecodeError, TypeError) as error: + raise AdapterUserError(message=_INVALID_SHOW_TYPE_MESSAGE) from error + if not isinstance(decoded, dict) or not isinstance(decoded.get("type"), str): + raise AdapterUserError(message=_INVALID_SHOW_TYPE_MESSAGE) + raw_name: str = str(decoded["type"]).upper() + return information_schema_type( + data_type=_SHOW_TYPE_TO_INFORMATION_SCHEMA_TYPE.get(raw_name, raw_name), + numeric_precision=_metadata_integer(decoded.get("precision")), + numeric_scale=_metadata_integer(decoded.get("scale")), + character_maximum_length=_metadata_integer(decoded.get("length")), + ) + + +def _metadata_integer(value: object) -> int | None: + if isinstance(value, bool) or not isinstance(value, int): + return None + return value diff --git a/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py new file mode 100644 index 0000000000..ac351243be --- /dev/null +++ b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py @@ -0,0 +1,95 @@ +"""Map Snowflake SHOW TABLES and SHOW VIEWS rows onto INFORMATION_SCHEMA relation semantics.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from sqlbuild.adapter.contract.models import RelationInfo +from sqlbuild.adapter.relations.main.relation_age_timestamp import relation_age_timestamp_utc +from sqlbuild.adapter.relations.models import ListedRelation + +_MISSING_OBJECT_ERRNO: int = 2003 +_MISSING_OBJECT_MESSAGE: str = "does not exist or not authorized" +_TRUE_FLAGS: frozenset[str] = frozenset({"Y", "YES", "TRUE"}) +_TRANSIENT_KIND: str = "TRANSIENT" +_TEMPORARY_KIND: str = "TEMPORARY" + + +def listed_relation_from_show_table( + *, row: Mapping[str, object], database: str | None +) -> ListedRelation: + """Translate one SHOW TABLES row into the relation INFORMATION_SCHEMA.TABLES would list.""" + + kind: str = str(row.get("kind") or "").upper() + relation_type: str = "base table" + if _flag(row.get("is_external")): + relation_type = "external table" + elif kind == _TEMPORARY_KIND: + relation_type = "temporary table" + elif _flag(row.get("is_event")): + relation_type = "event table" + return _listed_relation( + row=row, + database=database, + relation_type=relation_type, + is_transient=kind == _TRANSIENT_KIND, + retention_days=_retention_days(row.get("retention_time")), + ) + + +def listed_relation_from_show_view( + *, row: Mapping[str, object], database: str | None +) -> ListedRelation: + """Translate one SHOW VIEWS row into the relation INFORMATION_SCHEMA.TABLES would list.""" + + return _listed_relation( + row=row, + database=database, + relation_type="materialized view" if _flag(row.get("is_materialized")) else "view", + is_transient=False, + retention_days=None, + ) + + +def is_missing_object_error(error: Exception) -> bool: + """Return whether SHOW failed because the schema is absent or invisible to the role.""" + + return ( + getattr(error, "errno", None) == _MISSING_OBJECT_ERRNO + or _MISSING_OBJECT_MESSAGE in str(error).lower() + ) + + +def _listed_relation( + *, + row: Mapping[str, object], + database: str | None, + relation_type: str, + is_transient: bool, + retention_days: int | None, +) -> ListedRelation: + stored_name: str = str(row["name"]) + return ListedRelation( + stored_name=stored_name, + relation=RelationInfo( + database=database, + schema=str(row["schema_name"]).lower(), + name=stored_name.lower(), + relation_type=relation_type, + is_transient=is_transient, + created_at=relation_age_timestamp_utc(row.get("created_on")), + last_altered_at=None, + retention_days=retention_days, + ), + ) + + +def _flag(value: object) -> bool: + if isinstance(value, bool): + return value + return str(value or "").strip().upper() in _TRUE_FLAGS + + +def _retention_days(value: object) -> int | None: + text: str = str(value if value is not None else "").strip() + return int(text) if text.lstrip("-").isdigit() else None diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index b0ccbf6f6d..d36735bbe0 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -5,6 +5,7 @@ import csv import json import logging +from collections.abc import Callable from dataclasses import replace from datetime import date, datetime from decimal import Decimal @@ -81,9 +82,11 @@ TablePromotionMode, ) from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog +from sqlbuild.adapter.relations.constants import INSPECTION_IN_LIST_LIMIT from sqlbuild.adapter.relations.main.active_inspection_catalog import active_inspection_catalog from sqlbuild.adapter.relations.main.relation_age_timestamp import relation_age_timestamp_utc from sqlbuild.adapter.relations.main.render_inspection_sql import render_inspection_sql +from sqlbuild.adapter.relations.main.run_bounded_inspections import run_bounded_inspections from sqlbuild.adapter.relations.main.run_recorded_inspection_query import ( run_recorded_inspection_query, ) @@ -107,19 +110,26 @@ show_grants_object_kind, snowflake_relation_grants, ) +from sqlbuild.adapters.snowflake._helpers.metadata_types import ( + information_schema_type, + show_columns_type, +) +from sqlbuild.adapters.snowflake._helpers.show_metadata import ( + is_missing_object_error, + listed_relation_from_show_table, + listed_relation_from_show_view, +) from sqlbuild.adapters.snowflake.classes.snowflake_connection import _SnowflakeConnection from sqlbuild.adapters.snowflake.constants import ( BASE_TABLE_METADATA_TYPE, EXTERNAL_BROWSER_AUTHENTICATOR, MAX_STATEMENT_TIMEOUT_SECONDS, MFA_AUTHENTICATOR, - NUMBER_TYPE_NAME, OAUTH_AUTHORIZATION_CODE_AUTHENTICATOR, SECONDARY_ROLES_ALL, SECONDARY_ROLES_NONE, STATUS_COLUMN_NAME, SUCCESS_STATUS_TOKENS, - TEXT_TYPE_NAMES, TRUE_METADATA_VALUE, VIEW_RELATION_TYPE_TOKEN, ) @@ -136,6 +146,9 @@ _EXACT_COLUMN_INSPECTION_LIMIT: int = 32 _METADATA_INSPECTION_CONCURRENCY: int = 8 +_SHOW_RESULT_LIMIT: int = 10_000 +_SHOW_COLUMN_KIND: str = "COLUMN" +_SHOW_COLUMN_KINDS: frozenset[str] = frozenset({_SHOW_COLUMN_KIND, "VIRTUAL_COLUMN"}) _RELATION_LISTING_COLUMNS: str = ( "table_name, table_schema, table_type, is_transient, created, last_altered, retention_time" ) @@ -382,55 +395,61 @@ def _freshness_rows( schema: str | None, scoped_requests: list[TableFreshnessRequest], ) -> list[tuple[Any, ...]]: - """Reuse the planning catalog's schema listing, or read the requested tables.""" + """Read LAST_ALTERED for only the requested tables, in bounded IN-list chunks.""" - catalog: InspectionCatalog | None = active_inspection_catalog( - adapter=self, connection=connection - ) - if catalog is not None and schema is not None: - catalog.prefetch_relation_listings( - database=scoped_requests[0].database, schemas=(schema,) - ) - listing: SchemaRelationListing | None = catalog.cached_relation_listing( - database=scoped_requests[0].database, schema=schema + names: tuple[str, ...] = tuple( + dict.fromkeys( + self._information_schema_identifier(request.name) for request in scoped_requests ) - if listing is not None: - wanted: frozenset[str] = frozenset( - self._information_schema_identifier(request.name) for request in scoped_requests - ) - return [ - (database, row[1], row[0], row[2], row[5]) - for row in listing.rows - if str(row[0]) in wanted - ] - clauses: list[str] = [] - params: list[str] = [] + ) + relation: str = self._information_schema_relation( + database=scoped_requests[0].database, name="tables" + ) + scope_clauses: list[str] = [] + scope_params: list[str] = [] if database is not None: - clauses.append("table_catalog = %s") - params.append(database) + scope_clauses.append("table_catalog = %s") + scope_params.append(database) if schema is not None: - clauses.append("table_schema = %s") - params.append(schema) - placeholders: str = ", ".join(["%s"] * len(scoped_requests)) - clauses.append(f"table_name IN ({placeholders})") - params.extend( - self._information_schema_identifier(request.name) for request in scoped_requests + scope_clauses.append("table_schema = %s") + scope_params.append(schema) + chunks: tuple[tuple[str, ...], ...] = tuple( + names[start : start + INSPECTION_IN_LIST_LIMIT] + for start in range(0, len(names), INSPECTION_IN_LIST_LIMIT) ) - cursor: Any = connection.cursor() - try: - cursor.execute( - "SELECT table_catalog, table_schema, table_name, table_type, last_altered " - "FROM " - + self._information_schema_relation( - database=scoped_requests[0].database, name="tables" + catalog: InspectionCatalog | None = active_inspection_catalog( + adapter=self, connection=connection + ) + chunk_rows: list[list[tuple[Any, ...]]] = run_bounded_inspections( + tasks=tuple( + self._bind_freshness_chunk( + connection=connection, + query=( + "SELECT table_catalog, table_schema, table_name, table_type, last_altered " + f"FROM {relation} WHERE " + + " AND ".join( + (*scope_clauses, f"table_name IN ({', '.join(['%s'] * len(chunk))})") + ) + ), + params=(*scope_params, *chunk), ) - + " WHERE " - + " AND ".join(clauses), - tuple(params), - ) - return list(cursor.fetchall()) - finally: - cursor.close() + for chunk in chunks + ), + concurrency=1 if catalog is None else catalog.concurrency, + ) + rows: list[tuple[Any, ...]] = [] + chunk_result: list[tuple[Any, ...]] + for chunk_result in chunk_rows: + rows.extend(chunk_result) + return rows + + def _bind_freshness_chunk( + self, *, connection: Any, query: str, params: tuple[str, ...] + ) -> Callable[[], list[tuple[Any, ...]]]: + def read() -> list[tuple[Any, ...]]: + return self._fetch_inspection_rows(connection=connection, query=query, params=params) + + return read @staticmethod def _freshness_request_key( @@ -1793,38 +1812,59 @@ def _relation_info_from_row(*, row: tuple[Any, ...], database: str | None) -> Re def read_schema_relation_listing( self, *, connection: _SnowflakeConnection, database: str | None, schema: str ) -> SchemaRelationListing: - """List every relation in one schema with a single unfiltered metadata read.""" - - query: str = ( - f"SELECT {_RELATION_LISTING_COLUMNS} " - f"FROM {self._information_schema_relation(database=database, name='tables')} " - "WHERE table_schema = %s" + """List one schema's tables and views with SHOW, which needs no running warehouse.""" + + scope: str = self._show_schema_scope(database=database, schema=schema) + tables: list[dict[str, object]] = self._fetch_paged_show_rows( + connection=connection, query=f"SHOW TABLES IN SCHEMA {scope}" + ) + views: list[dict[str, object]] = self._fetch_paged_show_rows( + connection=connection, query=f"SHOW VIEWS IN SCHEMA {scope}" + ) + entries: dict[str, ListedRelation] = {} + row: dict[str, object] + for row in tables: + listed: ListedRelation = listed_relation_from_show_table(row=row, database=database) + entries[listed.stored_name] = listed + for row in views: + listed = listed_relation_from_show_view(row=row, database=database) + entries[listed.stored_name] = listed + return SchemaRelationListing( + database=database, schema=schema, entries=tuple(entries.values()) ) - params: list[str] = [self._information_schema_identifier(schema)] - if database is not None: - query += " AND table_catalog = %s" - params.append(self._information_schema_identifier(database)) - rows: list[tuple[Any, ...]] = self._fetch_inspection_rows( - connection=connection, query=query, params=tuple(params) + + def read_schema_column_listing( + self, *, connection: _SnowflakeConnection, database: str | None, schema: str + ) -> SchemaColumnListing: + """List one schema's columns with SHOW COLUMNS; reread a truncated result exactly.""" + + scope: str = self._show_schema_scope(database=database, schema=schema) + rows: list[dict[str, object]] = self._fetch_show_rows( + connection=connection, query=f"SHOW COLUMNS IN SCHEMA {scope}" ) - return SchemaRelationListing( + if len(rows) >= _SHOW_RESULT_LIMIT: + return self._read_information_schema_column_listing( + connection=connection, database=database, schema=schema + ) + columns: dict[str, list[ColumnInfo]] = {} + for row in rows: + if str(row.get("kind") or _SHOW_COLUMN_KIND).upper() not in _SHOW_COLUMN_KINDS: + continue + columns.setdefault(str(row["table_name"]), []).append( + ColumnInfo( + name=str(row["column_name"]).lower(), + type=show_columns_type(row["data_type"]), + ) + ) + return SchemaColumnListing( database=database, schema=schema, - entries=tuple( - ListedRelation( - stored_name=str(row[0]), - relation=self._relation_info_from_row(row=row, database=database), - ) - for row in rows - ), - rows=tuple(tuple(row) for row in rows), + columns_by_stored_name={name: tuple(value) for name, value in columns.items()}, ) - def read_schema_column_listing( + def _read_information_schema_column_listing( self, *, connection: _SnowflakeConnection, database: str | None, schema: str ) -> SchemaColumnListing: - """List every column in one schema with a single unfiltered metadata read.""" - query: str = ( "SELECT table_name, column_name, data_type, numeric_precision, numeric_scale, " "character_maximum_length FROM " @@ -1845,7 +1885,7 @@ def read_schema_column_listing( columns.setdefault(str(row[0]), []).append( ColumnInfo( name=str(row[1]).lower(), - type=self._build_information_schema_type( + type=information_schema_type( data_type=str(row[2]), numeric_precision=row[3], numeric_scale=row[4], @@ -1859,6 +1899,50 @@ def read_schema_column_listing( columns_by_stored_name={name: tuple(value) for name, value in columns.items()}, ) + def _show_schema_scope(self, *, database: str | None, schema: str) -> str: + if database is None: + return self.render_identifier(schema) + return f"{self.render_identifier(database)}.{self.render_identifier(schema)}" + + def _fetch_paged_show_rows( + self, *, connection: _SnowflakeConnection, query: str + ) -> list[dict[str, object]]: + """Page SHOW output past its row cap; SHOW orders rows by name.""" + + rows: list[dict[str, object]] = [] + page_query: str = f"{query} LIMIT {_SHOW_RESULT_LIMIT}" + while True: + page: list[dict[str, object]] = self._fetch_show_rows( + connection=connection, query=page_query + ) + rows.extend(page) + if len(page) < _SHOW_RESULT_LIMIT: + return rows + last_name: str = str(page[-1]["name"]).replace("'", "''") + page_query = f"{query} LIMIT {_SHOW_RESULT_LIMIT} FROM '{last_name}'" + + @staticmethod + def _fetch_show_rows( + *, connection: _SnowflakeConnection, query: str + ) -> list[dict[str, object]]: + """Run one SHOW command; a missing or unauthorized schema lists nothing.""" + + def fetch() -> list[dict[str, object]]: + cursor: Any = connection.cursor() + try: + try: + cursor.execute(query) + except Exception as error: + if is_missing_object_error(error): + return [] + raise + names: list[str] = [str(column[0]).lower() for column in cursor.description or ()] + return [dict(zip(names, row, strict=False)) for row in cursor.fetchall()] + finally: + cursor.close() + + return run_recorded_inspection_query(sql=query, run=fetch) + def read_relation_columns( self, *, connection: _SnowflakeConnection, relation: RelationInfo ) -> tuple[ColumnInfo, ...]: @@ -2084,7 +2168,7 @@ def fetch() -> list[tuple[Any, ...]]: return tuple( ColumnInfo( name=str(row[2]).lower(), - type=self._build_show_columns_type(raw_data_type=row[3]), + type=show_columns_type(row[3]), ) for row in rows ) @@ -2162,29 +2246,6 @@ def _get_columns_for_relations_bulk_chunk( finally: cursor.close() - def _build_show_columns_type(self, *, raw_data_type: object) -> str: - try: - decoded_data_type: object = json.loads(str(raw_data_type)) - except (json.JSONDecodeError, TypeError) as error: - raise AdapterUserError( - message="Snowflake SHOW COLUMNS returned invalid type metadata" - ) from error - if not isinstance(decoded_data_type, dict): - raise AdapterUserError(message="Snowflake SHOW COLUMNS returned invalid type metadata") - data_type: dict[str, object] = decoded_data_type - raw_name: str = str(data_type.get("type", "")) - normalized_name: str = { - "FIXED": NUMBER_TYPE_NAME, - "TEXT": "VARCHAR", - "REAL": "FLOAT", - }.get(raw_name.upper(), raw_name.upper()) - return self._build_information_schema_type( - data_type=normalized_name, - numeric_precision=data_type.get("precision"), - numeric_scale=data_type.get("scale"), - character_maximum_length=data_type.get("length"), - ) - def close(self, connection: _SnowflakeConnection) -> None: """Close a Snowflake connection.""" @@ -3210,16 +3271,12 @@ def _build_information_schema_type( numeric_scale: object, character_maximum_length: object, ) -> str: - normalized_type: str = data_type.upper() - if ( - normalized_type == NUMBER_TYPE_NAME - and isinstance(numeric_precision, int) - and isinstance(numeric_scale, int) - ): - return f"{NUMBER_TYPE_NAME}({numeric_precision},{numeric_scale})" - if normalized_type in TEXT_TYPE_NAMES and isinstance(character_maximum_length, int): - return f"VARCHAR({character_maximum_length})" - return normalized_type + return information_schema_type( + data_type=data_type, + numeric_precision=numeric_precision, + numeric_scale=numeric_scale, + character_maximum_length=character_maximum_length, + ) def validate_row_diff_keys( self, diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py b/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py index 0137602668..92ebedd06c 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py @@ -112,6 +112,7 @@ class SnowflakeInspectionBudgetTestCase: expected_schema_reads: dict[str, int] expected_metadata_budget: int expected_in_list_limit: int + expected_freshness_reads: int @dataclass(frozen=True) @@ -128,4 +129,4 @@ class SnowflakeReplanTestCase: """Planning the same project twice in separate invocations.""" description: str - expected_tables_reads: int + expected_listing_reads: int diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py b/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py index 50cdc792f5..c26340c0df 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py @@ -463,8 +463,18 @@ def grouped(model: type) -> dict[str, Any]: _OFFLINE_METADATA_KINDS: frozenset[str] = frozenset( - {"tables", "columns", "show_columns", "other_metadata"} + { + "tables", + "columns", + "show_columns", + "show_tables", + "show_views", + "show_schema_columns", + "other_metadata", + } ) +_OFFLINE_INFORMATION_SCHEMA_KINDS: frozenset[str] = frozenset({"tables", "columns"}) +_SHOW_SCHEMA_SCOPE: re.Pattern[str] = re.compile(r' IN SCHEMA (?:"[^"]*"\.)?"([^"]*)"') _CURSOR_BOUND_RELATION: re.Pattern[str] = re.compile(r"\sFROM\s+(\S+)\s*$", re.IGNORECASE) @@ -498,9 +508,23 @@ def offline_metadata_queries(warehouse: RecordingSnowflakeWarehouse) -> tuple[Re def offline_metadata_reads_by_schema( *, warehouse: RecordingSnowflakeWarehouse, kind: str ) -> dict[str, int]: - """Count schema-scoped metadata reads of one kind by their schema parameter.""" + """Count schema-scoped SHOW reads of one kind by the schema they list.""" + + return dict( + Counter( + _SHOW_SCHEMA_SCOPE.findall(query.sql)[0] for query in warehouse.queries_of_kind(kind) + ) + ) + - return dict(Counter(str(query.params[0]) for query in warehouse.queries_of_kind(kind))) +def offline_information_schema_queries( + warehouse: RecordingSnowflakeWarehouse, +) -> tuple[RecordedQuery, ...]: + """Return metadata statements that need a running warehouse.""" + + return tuple( + filter(lambda query: query.kind in _OFFLINE_INFORMATION_SCHEMA_KINDS, warehouse.queries) + ) def offline_cursor_bound_relations(warehouse: RecordingSnowflakeWarehouse) -> tuple[str, ...]: diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py b/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py index 064b6d2ed1..26122e543e 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py @@ -19,6 +19,7 @@ ) from tests.integration.src.sqlbuild.compiler.planner.main.helpers import ( offline_cursor_bound_relations, + offline_information_schema_queries, offline_metadata_queries, offline_metadata_reads_by_schema, plan_offline_snowflake_project, @@ -32,6 +33,9 @@ ) _ONE_READ_PER_SCHEMA: dict[str, int] = {"STAGING": 1, "INTERMEDIATE": 1, "MARTS": 1, "RAW": 1} +_SOURCES: int = 260 +_SHOW_READS_PER_SCHEMA: int = 3 +_FRESHNESS_READS: int = -(-_SOURCES // INSPECTION_IN_LIST_LIMIT) @pytest.mark.parametrize( @@ -41,15 +45,17 @@ description="about 2,500 relations", unmanaged_relations_per_schema=450, expected_schema_reads=_ONE_READ_PER_SCHEMA, - expected_metadata_budget=8, + expected_metadata_budget=_SHOW_READS_PER_SCHEMA * 4 + _FRESHNESS_READS, expected_in_list_limit=INSPECTION_IN_LIST_LIMIT, + expected_freshness_reads=_FRESHNESS_READS, ), SnowflakeInspectionBudgetTestCase( description="about 4,300 relations", unmanaged_relations_per_schema=900, expected_schema_reads=_ONE_READ_PER_SCHEMA, - expected_metadata_budget=8, + expected_metadata_budget=_SHOW_READS_PER_SCHEMA * 4 + _FRESHNESS_READS, expected_in_list_limit=INSPECTION_IN_LIST_LIMIT, + expected_freshness_reads=_FRESHNESS_READS, ), ], ids=lambda case: case.description, @@ -61,7 +67,7 @@ def test_given_large_multi_schema_project_when_planning_then_metadata_reads_stay project_dir=tmp_path / "orders_platform", models_per_schema=60, incremental_every=2, - sources=260, + sources=_SOURCES, unmanaged_relations_per_schema=test_case.unmanaged_relations_per_schema, ) warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( @@ -74,16 +80,23 @@ def test_given_large_multi_schema_project_when_planning_then_metadata_reads_stay metadata: tuple[RecordedQuery, ...] = offline_metadata_queries(warehouse) assert len(result.plan_output.model_entries) == len(project.model_schemas) + warehouse_queries: tuple[RecordedQuery, ...] = offline_information_schema_queries(warehouse) assert ( - offline_metadata_reads_by_schema(warehouse=warehouse, kind="tables") + offline_metadata_reads_by_schema(warehouse=warehouse, kind="show_tables") == test_case.expected_schema_reads ) assert ( - offline_metadata_reads_by_schema(warehouse=warehouse, kind="columns") + offline_metadata_reads_by_schema(warehouse=warehouse, kind="show_views") + == test_case.expected_schema_reads + ) + assert ( + offline_metadata_reads_by_schema(warehouse=warehouse, kind="show_schema_columns") == test_case.expected_schema_reads ) assert warehouse.queries_of_kind("show_columns") == () assert warehouse.queries_of_kind("other_metadata") == () + assert len(warehouse_queries) == test_case.expected_freshness_reads + assert all("last_altered" in query.sql for query in warehouse_queries) assert max(query.largest_in_list for query in metadata) <= test_case.expected_in_list_limit assert len({(query.sql, query.params) for query in metadata}) == len(metadata) assert len(metadata) <= test_case.expected_metadata_budget @@ -127,7 +140,8 @@ def test_given_incremental_models_when_planning_then_cursor_bounds_run_one_per_r "test_case", [ SnowflakeReplanTestCase( - description="a second invocation re-reads bounds and listings", expected_tables_reads=4 + description="a second invocation re-reads bounds and listings", + expected_listing_reads=4, ) ], ids=lambda case: case.description, @@ -148,7 +162,7 @@ def test_given_planned_project_when_planning_again_then_cursor_bounds_are_read_a _ = plan_offline_snowflake_project(project=project, warehouse=warehouse, no_cache=False) assert len(warehouse.queries_of_kind("cursor_bounds")) == first - assert len(warehouse.queries_of_kind("tables")) == test_case.expected_tables_reads + assert len(warehouse.queries_of_kind("show_tables")) == test_case.expected_listing_reads if __name__ == "__main__": diff --git a/tests/unit/src/sqlbuild/adapter/relations/test_metadata_sql_ownership.py b/tests/unit/src/sqlbuild/adapter/relations/test_metadata_sql_ownership.py index ee36d55226..f762fe4391 100644 --- a/tests/unit/src/sqlbuild/adapter/relations/test_metadata_sql_ownership.py +++ b/tests/unit/src/sqlbuild/adapter/relations/test_metadata_sql_ownership.py @@ -18,6 +18,8 @@ "adapters/bigquery/classes/bigquery_adapter.py", "adapters/databricks/classes/databricks_adapter.py", "adapters/postgres/classes/postgres_adapter.py", + "adapters/snowflake/_helpers/metadata_types.py", + "adapters/snowflake/_helpers/show_metadata.py", "adapters/snowflake/classes/snowflake_adapter.py", "adapters/sqlserver/classes/sqlserver_adapter.py", "adapters/sqlserver/constants.py", diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py index d6e365e2ff..de0bfa972c 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py @@ -61,3 +61,54 @@ class SpeculativePrefetchTestCase: failing_schema: str expected_error_fragment: str expected_failed_reads: int + + +@dataclass(frozen=True) +class ShowColumnTypeTestCase: + """One SHOW COLUMNS data_type and the INFORMATION_SCHEMA fields of the same column.""" + + description: str + show_data_type: dict[str, object] + information_schema_fields: tuple[str, object, object, object] + expected_type: str + + +@dataclass(frozen=True) +class InvalidShowColumnTypeTestCase: + """A SHOW COLUMNS data_type value that is not valid type metadata.""" + + description: str + raw_data_type: object + expected_error_fragment: str + + +@dataclass(frozen=True) +class ShowRelationRowTestCase: + """One SHOW TABLES or SHOW VIEWS row and the relation INFORMATION_SCHEMA would list.""" + + description: str + is_view: bool + row: dict[str, object] + expected_relation_type: str + expected_is_transient: bool + expected_retention_days: int | None + + +@dataclass(frozen=True) +class MissingObjectErrorTestCase: + """A SHOW failure classified as an absent schema or a real error.""" + + description: str + errno: int | None + message: str + expected_missing: bool + + +@dataclass(frozen=True) +class ShowResultCapTestCase: + """A schema whose SHOW output reaches Snowflake's 10,000-row cap.""" + + description: str + relation_count: int + columns_per_relation: int + expected_query_kinds: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py index e9f55799a3..a614aa8cc6 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -12,7 +12,7 @@ import threading import time from collections.abc import Callable, Iterator -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from datetime import UTC, datetime from itertools import compress from pathlib import Path @@ -35,6 +35,13 @@ _SHOW_COLUMNS_PATTERN: re.Pattern[str] = re.compile( r"^\s*SHOW\s+COLUMNS\s+IN\s+(?:TABLE|VIEW)\s+(?P\S+)\s*$", re.IGNORECASE ) +_SHOW_SCHEMA_RELATIONS_PATTERN: re.Pattern[str] = re.compile( + r"^SHOW (?PTABLES|VIEWS) IN SCHEMA (?P\S+)" + r"(?: LIMIT (?P\d+)(?: FROM '(?P(?:[^']|'')*)')?)?$" +) +_SHOW_SCHEMA_COLUMNS_PATTERN: re.Pattern[str] = re.compile( + r"^SHOW COLUMNS IN SCHEMA (?P\S+)$" +) _SHOW_PATTERN: re.Pattern[str] = re.compile(r"^\s*SHOW\s", re.IGNORECASE) _CURSOR_BOUND_PATTERN: re.Pattern[str] = re.compile( r"^SELECT\s+CAST\((?:MIN|MAX)\(.*\)\s+AS\s+VARCHAR\).*\s+FROM\s+(?P\S+)\s*$", @@ -48,6 +55,9 @@ (_TABLES_PATTERN, "tables"), (_COLUMNS_PATTERN, "columns"), (_SHOW_COLUMNS_PATTERN, "show_columns"), + (re.compile(r"^SHOW TABLES IN SCHEMA "), "show_tables"), + (re.compile(r"^SHOW VIEWS IN SCHEMA "), "show_views"), + (_SHOW_SCHEMA_COLUMNS_PATTERN, "show_schema_columns"), (_CURSOR_BOUND_PATTERN, "cursor_bounds"), (_OTHER_INFORMATION_SCHEMA_PATTERN, "other_metadata"), (_SHOW_PATTERN, "other_metadata"), @@ -63,6 +73,52 @@ ("data_type",), ) _BOUND_MARKERS: tuple[str, str] = ("MIN(", "MAX(") +SHOW_RESULT_CAP: int = 10_000 +_SHOW_TABLE_FIELDS: tuple[str, ...] = ( + "created_on", + "name", + "database_name", + "schema_name", + "kind", + "comment", + "cluster_by", + "rows", + "bytes", + "owner", + "retention_time", + "is_external", + "is_event", + "is_dynamic", +) +_SHOW_VIEW_FIELDS: tuple[str, ...] = ( + "created_on", + "name", + "reserved", + "database_name", + "schema_name", + "owner", + "comment", + "text", + "is_secure", + "is_materialized", +) +_SHOW_SCHEMA_COLUMN_FIELDS: tuple[str, ...] = ( + "table_name", + "schema_name", + "column_name", + "data_type", + "null?", + "default", + "kind", + "expression", + "comment", + "database_name", + "autoincrement", +) +_VIEW_TABLE_TYPES: frozenset[str] = frozenset({"VIEW", "MATERIALIZED VIEW"}) +_TIMESTAMP_TYPES: frozenset[str] = frozenset( + {"TIMESTAMP_NTZ", "TIMESTAMP_LTZ", "TIMESTAMP_TZ", "TIME"} +) @dataclass(frozen=True) @@ -124,6 +180,7 @@ class RecordingSnowflakeWarehouse: statement_latency_seconds: float = 0.0 failing_relations: frozenset[str] = frozenset() failing_metadata_schemas: frozenset[str] = frozenset() + current_database: str = "ANALYTICS" cursor_values: tuple[str, str] = ("2026-01-01 00:00:00.000", "2026-01-31 00:00:00.000") queries: list[RecordedQuery] = field(default_factory=list) attempted_sql: list[str] = field(default_factory=list) @@ -213,6 +270,9 @@ def run(self, *, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[Any, return rows, description def _load(self, *, table: str, rows: list[dict[str, object]]) -> None: + _ROW_LOADERS[bool(rows)](self, table, rows) + + def _insert_rows(self, table: str, rows: list[dict[str, object]]) -> None: staged: pyarrow.Table = pyarrow.Table.from_pylist(rows) self._database.register("staged_rows", staged) self._database.execute(f"INSERT INTO {table} BY NAME SELECT * FROM staged_rows") @@ -223,6 +283,8 @@ def _answer(self, *, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[A (_METADATA_RELATION_PATTERN, self._answer_information_schema), (_OTHER_INFORMATION_SCHEMA_PATTERN, _answer_empty_metadata), (_SHOW_COLUMNS_PATTERN, self._answer_show_columns), + (_SHOW_SCHEMA_RELATIONS_PATTERN, self._answer_show_schema_relations), + (_SHOW_SCHEMA_COLUMNS_PATTERN, self._answer_show_schema_columns), (_CURSOR_BOUND_PATTERN, self._answer_cursor_bounds), (_ANY_PATTERN, _answer_status), ) @@ -237,9 +299,9 @@ def _answer_information_schema( _MISSING_RELATION_ACTIONS[ not self.failing_metadata_schemas.isdisjoint(str(param) for param in params) ](sql) - duckdb_sql: str = _METADATA_RELATION_PATTERN.sub(_scope_metadata_relation, sql).replace( - "%s", "?" - ) + duckdb_sql: str = _METADATA_RELATION_PATTERN.sub( + self._scope_metadata_relation, sql + ).replace("%s", "?") with self._lock: cursor: duckdb.DuckDBPyConnection = self._database.cursor() result: duckdb.DuckDBPyConnection = cursor.execute(duckdb_sql, list(params)) @@ -268,6 +330,75 @@ def _answer_show_columns( _MISSING_RELATION_ACTIONS[not rows](relation) return rows, _SHOW_DESCRIPTION + def _scope_metadata_relation(self, match: re.Match[str]) -> str: + table: str = f"fake_{match.group('view').lower()}" + database: str = (match.group("database") or self.current_database).replace("'", "''") + return f"(SELECT * FROM {table} WHERE table_catalog = '{database}') AS {table}" + + def _schema_relations(self, *, scope: str) -> tuple[FakeRelation, ...]: + parts: tuple[str, ...] = ( + self.current_database, + *(part.strip('"') for part in scope.split(".")), + )[-2:] + _FAILED_SCHEMA_ACTIONS[parts[-1] in self.failing_metadata_schemas](scope) + matches: tuple[FakeRelation, ...] = tuple( + filter( + lambda fake: (fake.database, fake.schema)[-len(parts) :] == parts, self.relations + ) + ) + _MISSING_SCHEMA_ACTIONS[not matches](scope) + return matches + + def _answer_show_schema_relations( + self, *, sql: str, params: tuple[object, ...] + ) -> tuple[list[tuple[Any, ...]], Any]: + del params + match: re.Match[str] = next(_SHOW_SCHEMA_RELATIONS_PATTERN.finditer(sql)) + wants_views: bool = match.group("kind") == "VIEWS" + after: str = (match.group("after") or "").replace("''", "'") + selected: list[FakeRelation] = sorted( + filter( + lambda fake: ( + (fake.table_type in _VIEW_TABLE_TYPES) is wants_views and fake.name > after + ), + self._schema_relations(scope=match.group("scope")), + ), + key=lambda fake: fake.name, + ) + limit: int = int(match.group("limit") or SHOW_RESULT_CAP) + row_builder: Callable[[FakeRelation], tuple[Any, ...]] = ( + _show_table_row, + _show_view_row, + )[wants_views] + fields: tuple[str, ...] = (_SHOW_TABLE_FIELDS, _SHOW_VIEW_FIELDS)[wants_views] + return [row_builder(fake) for fake in selected[:limit]], tuple((f,) for f in fields) + + def _answer_show_schema_columns( + self, *, sql: str, params: tuple[object, ...] + ) -> tuple[list[tuple[Any, ...]], Any]: + del params + scope: str = next(_SHOW_SCHEMA_COLUMNS_PATTERN.finditer(sql)).group("scope") + rows: list[tuple[Any, ...]] = [] + fake: FakeRelation + for fake in sorted(self._schema_relations(scope=scope), key=lambda item: item.name): + rows.extend( + ( + fake.name, + fake.schema, + column.name, + json.dumps(_show_columns_type(column)), + "true", + "", + "COLUMN", + "", + "", + fake.database, + "", + ) + for column in fake.columns + ) + return rows[:SHOW_RESULT_CAP], tuple((f,) for f in _SHOW_SCHEMA_COLUMN_FIELDS) + def _answer_cursor_bounds( self, *, sql: str, params: tuple[object, ...] ) -> tuple[list[tuple[Any, ...]], Any]: @@ -295,6 +426,79 @@ def _raise_unavailable_relation(relation: str) -> None: } +def _skip_rows(warehouse: RecordingSnowflakeWarehouse, table: str, rows: list[Any]) -> None: + del warehouse, table, rows + + +_ROW_LOADERS: dict[bool, Callable[[RecordingSnowflakeWarehouse, str, list[Any]], None]] = { + False: _skip_rows, + True: RecordingSnowflakeWarehouse._insert_rows, +} + + +class FakeSnowflakeProgrammingError(RuntimeError): + """Driver-style error carrying a Snowflake error number.""" + + def __init__(self, message: str, *, errno: int) -> None: + super().__init__(message) + self.errno: int = errno + + +def _raise_missing_schema(scope: str) -> None: + raise FakeSnowflakeProgrammingError( + f"SQL compilation error: Schema '{scope}' does not exist or not authorized.", errno=2003 + ) + + +def _raise_failed_schema(scope: str) -> None: + raise FakeSnowflakeProgrammingError(f"Metadata service unavailable for {scope}.", errno=390) + + +_MISSING_SCHEMA_ACTIONS: dict[bool, Callable[[str], None]] = { + False: _accept_relation, + True: _raise_missing_schema, +} +_FAILED_SCHEMA_ACTIONS: dict[bool, Callable[[str], None]] = { + False: _accept_relation, + True: _raise_failed_schema, +} + + +def _show_table_row(relation: FakeRelation) -> tuple[Any, ...]: + external: bool = relation.table_type == "EXTERNAL TABLE" + return ( + _FIXED_CREATED_AT, + relation.name, + relation.database, + relation.schema, + ("TABLE", "TRANSIENT")[relation.is_transient], + "", + "", + 0, + 0, + "SYSADMIN", + str(relation.retention_time), + ("N", "Y")[external], + "N", + "N", + ) + + +def _show_view_row(relation: FakeRelation) -> tuple[Any, ...]: + return ( + _FIXED_CREATED_AT, + relation.name, + "", + relation.database, + relation.schema, + "SYSADMIN", + "", + "", + "false", + ("false", "true")[relation.table_type == "MATERIALIZED VIEW"], + ) + + def _answer_empty_metadata( *, sql: str, params: tuple[object, ...] ) -> tuple[list[tuple[Any, ...]], Any]: @@ -307,13 +511,6 @@ def _answer_status(*, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[ return [], _STATUS_DESCRIPTION -def _scope_metadata_relation(match: re.Match[str]) -> str: - table: str = f"fake_{match.group('view').lower()}" - database: str = (match.group("database") or "").replace("'", "''") - scoped: str = f"(SELECT * FROM {table} WHERE table_catalog = '{database}') AS {table}" - return (scoped, table)[match.group("database") is None] - - def _table_row(relation: FakeRelation) -> dict[str, object]: return { "table_catalog": relation.database, @@ -323,7 +520,7 @@ def _table_row(relation: FakeRelation) -> dict[str, object]: "is_transient": ("NO", "YES")[relation.is_transient], "created": _FIXED_CREATED_AT, "last_altered": _FIXED_CREATED_AT, - "retention_time": relation.retention_time, + "retention_time": (relation.retention_time, None)[relation.table_type in _VIEW_TABLE_TYPES], } @@ -352,8 +549,8 @@ def _show_columns_type(column: FakeColumn) -> dict[str, object]: candidates: tuple[tuple[str, object], ...] = ( ("type", _SHOW_TYPE_NAMES.get(column.data_type, column.data_type)), ("nullable", True), - ("precision", column.numeric_precision), - ("scale", column.numeric_scale), + ("precision", (column.numeric_precision, 0)[column.data_type in _TIMESTAMP_TYPES]), + ("scale", (column.numeric_scale, 9)[column.data_type in _TIMESTAMP_TYPES]), ("length", (None, column.character_maximum_length)[column.data_type == "TEXT"]), ) return dict(filter(lambda item: item[1] is not None, candidates)) @@ -581,7 +778,7 @@ def build_inspection_catalog_relations() -> tuple[FakeRelation, ...]: ), *( FakeRelation(database="ANALYTICS", schema="STAGING", name=f"INVENTORY_{index}") - for index in range(40) + for index in range(60) ), FakeRelation(database="ANALYTICS", schema="MARTS", name="REVENUE", columns=_ORDER_COLUMNS), FakeRelation(database="ANALYTICS", schema="MARTS", name="revenue_v", table_type="VIEW"), @@ -603,9 +800,9 @@ def build_offline_snowflake( def sorted_relation_reprs(relations: tuple[RelationInfo, ...]) -> list[str]: - """Return an order-independent comparable form of listed relations.""" + """Return an order-independent form of listed relations without SHOW-absent LAST_ALTERED.""" - return sorted(repr(relation) for relation in relations) + return sorted(repr(replace(relation, last_altered_at=None)) for relation in relations) def _project_toml(*, database: str, model_schemas: tuple[str, ...]) -> str: @@ -646,3 +843,28 @@ def _sources_yml(*, database: str, source_schema: str, sources: int) -> str: ) ) return "\n".join(lines) + "\n" + + +def build_wide_schema_relations( + *, relation_count: int, columns_per_relation: int +) -> tuple[FakeRelation, ...]: + """Return one schema of identical relations sized to exercise SHOW result caps.""" + + columns: tuple[FakeColumn, ...] = tuple( + FakeColumn(name=f"ATTRIBUTE_{index}") for index in range(columns_per_relation) + ) + return tuple( + FakeRelation( + database="ANALYTICS", + schema="INVENTORY", + name=f"STOCK_{index:05d}", + columns=columns, + ) + for index in range(relation_count) + ) + + +def show_relation_row(*, created_on: datetime, **fields: object) -> dict[str, object]: + """Return one SHOW TABLES or SHOW VIEWS row for a STAGING.ORDERS relation.""" + + return {"created_on": created_on, "name": "ORDERS", "schema_name": "STAGING", **fields} diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py index 3f69723535..6a08d99792 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py @@ -22,6 +22,7 @@ FreshnessReuseTestCase, RelationRequestEquivalenceTestCase, RepeatedLookupTestCase, + ShowResultCapTestCase, SpeculativePrefetchTestCase, ) from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( @@ -29,6 +30,7 @@ RecordingSnowflakeWarehouse, build_inspection_catalog_relations, build_offline_snowflake, + build_wide_schema_relations, sorted_relation_reprs, ) @@ -47,7 +49,7 @@ database="analytics", schemas=("staging",), names=None, - expected_relation_count=47, + expected_relation_count=67, ), RelationRequestEquivalenceTestCase( description="name filter across schemas ignores quoted lowercase twins", @@ -71,11 +73,11 @@ expected_relation_count=0, ), RelationRequestEquivalenceTestCase( - description="session database lists every catalog", + description="session database lists only the current catalog", database=None, schemas=("staging",), names=("orders",), - expected_relation_count=2, + expected_relation_count=1, ), RelationRequestEquivalenceTestCase( description="database-wide name lookup", @@ -89,7 +91,7 @@ database="analytics", schemas=None, names=_MANY_NAMES, - expected_relation_count=42, + expected_relation_count=62, ), ], ids=lambda case: case.description, @@ -130,7 +132,7 @@ def test_given_metadata_request_when_served_by_catalog_then_matches_direct_adapt description="broad request reads each schema's columns once without IN lists", schemas=("staging",), names=None, - expected_query_kinds=("columns",), + expected_query_kinds=("show_schema_columns",), ), ], ids=lambda case: case.description, @@ -172,19 +174,20 @@ def test_given_relations_when_reading_columns_through_catalog_then_matches_direc "test_case", [ FreshnessReuseTestCase( - description="freshness reuses listings without another table read", + description="freshness reads LAST_ALTERED only for requested tables in capped chunks", requests=( ("analytics", "staging", "orders"), ("analytics", "staging", "transient_events"), + *(("analytics", "staging", f"inventory_{index}") for index in range(60)), ("analytics", "marts", "revenue"), ), listed_schemas=("staging", "marts"), - expected_metadata_reads=2, + expected_metadata_reads=3, ) ], ids=lambda case: case.description, ) -def test_given_open_catalog_when_reading_source_freshness_then_reuses_schema_listing( +def test_given_open_catalog_when_reading_source_freshness_then_reads_requested_tables_only( test_case: FreshnessReuseTestCase, ) -> None: adapter, connection, warehouse = build_offline_snowflake( @@ -201,12 +204,15 @@ def test_given_open_catalog_when_reading_source_freshness_then_reuses_schema_lis with open_inspection_catalog(adapter=adapter, connection=connection) as catalog: _ = catalog.list_relations(database="analytics", schemas=test_case.listed_schemas) + warehouse.reset() served: dict[TableFreshnessRequest, TableFreshnessMetadata] = ( adapter.get_tables_freshness_metadata(connection=connection, requests=requests) ) assert served == direct + assert len(warehouse.queries_of_kind("tables")) == test_case.expected_metadata_reads assert len(warehouse.queries) == test_case.expected_metadata_reads + assert max(query.largest_in_list for query in warehouse.queries) <= INSPECTION_IN_LIST_LIMIT @pytest.mark.parametrize( @@ -244,8 +250,8 @@ def test_given_open_catalog_when_freshness_targets_a_view_then_raises_like_direc RepeatedLookupTestCase( description="state table, model, and whole-schema lookups share one listing", lookups=(("_sqlbuild_fingerprints",), ("orders",), None), - expected_relation_counts=(0, 1, 47), - expected_query_kinds=("tables",), + expected_relation_counts=(0, 1, 67), + expected_query_kinds=("show_tables", "show_views"), ) ], ids=lambda case: case.description, @@ -273,7 +279,7 @@ def test_given_repeated_lookups_when_catalog_is_open_then_lists_schema_once( SpeculativePrefetchTestCase( description="unreadable speculative schema fails only the request that needs it", failing_schema="RESTRICTED", - expected_error_fragment="not authorized", + expected_error_fragment="unavailable", expected_failed_reads=2, ) ], @@ -302,9 +308,62 @@ def test_given_unreadable_schema_when_prefetching_best_effort_then_only_real_rea ) assert len(staging) == 1 - assert len(warehouse.queries_of_kind("tables")) == 1 + assert len(warehouse.queries_of_kind("show_tables")) == 1 assert len(warehouse.attempted_sql) - len(warehouse.queries) == test_case.expected_failed_reads +@pytest.mark.parametrize( + "test_case", + [ + ShowResultCapTestCase( + description="more than 10,000 tables are paged", + relation_count=10_050, + columns_per_relation=0, + expected_query_kinds=( + "show_tables", + "show_tables", + "show_views", + "show_schema_columns", + ), + ), + ShowResultCapTestCase( + description="a truncated column listing is reread from INFORMATION_SCHEMA", + relation_count=40, + columns_per_relation=260, + expected_query_kinds=("show_tables", "show_views", "show_schema_columns", "columns"), + ), + ], + ids=lambda case: case.description, +) +def test_given_schema_beyond_show_cap_when_listing_then_matches_direct_adapter( + test_case: ShowResultCapTestCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=build_wide_schema_relations( + relation_count=test_case.relation_count, + columns_per_relation=test_case.columns_per_relation, + ) + ) + direct: tuple[RelationInfo, ...] = adapter.list_relations( + connection=connection, database="analytics", schemas=("inventory",) + ) + direct_columns: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( + adapter.get_columns_for_relations(connection=connection, relations=direct) + ) + warehouse.reset() + catalog: InspectionCatalog = InspectionCatalog(adapter=adapter, connection=connection) + + served: tuple[RelationInfo, ...] = catalog.list_relations( + database="analytics", schemas=("inventory",) + ) + served_columns: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( + catalog.get_columns_for_relations(relations=served) + ) + + assert sorted_relation_reprs(served) == sorted_relation_reprs(direct) + assert served_columns == direct_columns + assert tuple(query.kind for query in warehouse.queries) == test_case.expected_query_kinds + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_metadata_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_metadata_types.py new file mode 100644 index 0000000000..46937d1994 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_metadata_types.py @@ -0,0 +1,239 @@ +"""SHOW COLUMNS type metadata must render exactly the INFORMATION_SCHEMA type strings.""" + +from __future__ import annotations + +import json + +import pytest + +from sqlbuild.adapter.contract.exceptions import AdapterUserError +from sqlbuild.adapters.snowflake._helpers.metadata_types import ( + information_schema_type, + show_columns_type, +) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection._test_types import ( + InvalidShowColumnTypeTestCase, + ShowColumnTypeTestCase, +) + + +@pytest.mark.parametrize( + "test_case", + [ + ShowColumnTypeTestCase( + description="integer number", + show_data_type={"type": "FIXED", "precision": 38, "scale": 0, "nullable": True}, + information_schema_fields=("NUMBER", 38, 0, None), + expected_type="NUMBER(38,0)", + ), + ShowColumnTypeTestCase( + description="decimal number", + show_data_type={"type": "FIXED", "precision": 10, "scale": 2, "nullable": False}, + information_schema_fields=("NUMBER", 10, 2, None), + expected_type="NUMBER(10,2)", + ), + ShowColumnTypeTestCase( + description="float", + show_data_type={"type": "REAL", "nullable": True}, + information_schema_fields=("FLOAT", None, None, None), + expected_type="FLOAT", + ), + ShowColumnTypeTestCase( + description="default-length text", + show_data_type={ + "type": "TEXT", + "length": 16777216, + "byteLength": 16777216, + "nullable": True, + "fixed": False, + }, + information_schema_fields=("TEXT", None, None, 16777216), + expected_type="VARCHAR(16777216)", + ), + ShowColumnTypeTestCase( + description="bounded text", + show_data_type={"type": "TEXT", "length": 32, "byteLength": 128, "fixed": False}, + information_schema_fields=("TEXT", None, None, 32), + expected_type="VARCHAR(32)", + ), + ShowColumnTypeTestCase( + description="fixed-width char", + show_data_type={"type": "TEXT", "length": 1, "byteLength": 4, "fixed": True}, + information_schema_fields=("TEXT", None, None, 1), + expected_type="VARCHAR(1)", + ), + ShowColumnTypeTestCase( + description="binary keeps the bare name", + show_data_type={"type": "BINARY", "length": 8388608, "byteLength": 8388608}, + information_schema_fields=("BINARY", None, None, 8388608), + expected_type="BINARY", + ), + ShowColumnTypeTestCase( + description="boolean", + show_data_type={"type": "BOOLEAN", "nullable": True}, + information_schema_fields=("BOOLEAN", None, None, None), + expected_type="BOOLEAN", + ), + ShowColumnTypeTestCase( + description="date", + show_data_type={"type": "DATE", "nullable": True}, + information_schema_fields=("DATE", None, None, None), + expected_type="DATE", + ), + ShowColumnTypeTestCase( + description="time ignores its precision", + show_data_type={"type": "TIME", "precision": 0, "scale": 9, "nullable": True}, + information_schema_fields=("TIME", None, None, None), + expected_type="TIME", + ), + ShowColumnTypeTestCase( + description="timestamp without time zone ignores its precision", + show_data_type={"type": "TIMESTAMP_NTZ", "precision": 0, "scale": 9}, + information_schema_fields=("TIMESTAMP_NTZ", None, None, None), + expected_type="TIMESTAMP_NTZ", + ), + ShowColumnTypeTestCase( + description="millisecond timestamp without time zone", + show_data_type={"type": "TIMESTAMP_NTZ", "precision": 0, "scale": 3}, + information_schema_fields=("TIMESTAMP_NTZ", None, None, None), + expected_type="TIMESTAMP_NTZ", + ), + ShowColumnTypeTestCase( + description="local time zone timestamp", + show_data_type={"type": "TIMESTAMP_LTZ", "precision": 0, "scale": 9}, + information_schema_fields=("TIMESTAMP_LTZ", None, None, None), + expected_type="TIMESTAMP_LTZ", + ), + ShowColumnTypeTestCase( + description="offset timestamp", + show_data_type={"type": "TIMESTAMP_TZ", "precision": 0, "scale": 9}, + information_schema_fields=("TIMESTAMP_TZ", None, None, None), + expected_type="TIMESTAMP_TZ", + ), + ShowColumnTypeTestCase( + description="variant", + show_data_type={"type": "VARIANT", "nullable": True}, + information_schema_fields=("VARIANT", None, None, None), + expected_type="VARIANT", + ), + ShowColumnTypeTestCase( + description="semi-structured object", + show_data_type={"type": "OBJECT", "nullable": True}, + information_schema_fields=("OBJECT", None, None, None), + expected_type="OBJECT", + ), + ShowColumnTypeTestCase( + description="semi-structured array", + show_data_type={"type": "ARRAY", "nullable": True}, + information_schema_fields=("ARRAY", None, None, None), + expected_type="ARRAY", + ), + ShowColumnTypeTestCase( + description="structured map", + show_data_type={"type": "MAP", "keyType": {"type": "TEXT"}, "valueType": {}}, + information_schema_fields=("MAP", None, None, None), + expected_type="MAP", + ), + ShowColumnTypeTestCase( + description="geography", + show_data_type={"type": "GEOGRAPHY", "nullable": True}, + information_schema_fields=("GEOGRAPHY", None, None, None), + expected_type="GEOGRAPHY", + ), + ShowColumnTypeTestCase( + description="geometry", + show_data_type={"type": "GEOMETRY", "nullable": True}, + information_schema_fields=("GEOMETRY", None, None, None), + expected_type="GEOMETRY", + ), + ShowColumnTypeTestCase( + description="vector", + show_data_type={"type": "VECTOR", "vectorElementType": "FLOAT", "dimension": 3}, + information_schema_fields=("VECTOR", None, None, None), + expected_type="VECTOR", + ), + ShowColumnTypeTestCase( + description="decfloat", + show_data_type={"type": "DECFLOAT", "nullable": True}, + information_schema_fields=("DECFLOAT", None, None, None), + expected_type="DECFLOAT", + ), + ShowColumnTypeTestCase( + description="file", + show_data_type={"type": "FILE", "nullable": True}, + information_schema_fields=("FILE", None, None, None), + expected_type="FILE", + ), + ShowColumnTypeTestCase( + description="lowercase type name", + show_data_type={"type": "fixed", "precision": 38, "scale": 0}, + information_schema_fields=("number", 38, 0, None), + expected_type="NUMBER(38,0)", + ), + ShowColumnTypeTestCase( + description="number without precision keeps the bare name", + show_data_type={"type": "FIXED"}, + information_schema_fields=("NUMBER", None, None, None), + expected_type="NUMBER", + ), + ShowColumnTypeTestCase( + description="boolean-valued precision is not a precision", + show_data_type={"type": "FIXED", "precision": True, "scale": 0}, + information_schema_fields=("NUMBER", None, 0, None), + expected_type="NUMBER", + ), + ], + ids=lambda case: case.description, +) +def test_given_column_metadata_when_rendering_type_then_show_matches_information_schema( + test_case: ShowColumnTypeTestCase, +) -> None: + data_type, precision, scale, length = test_case.information_schema_fields + + from_show: str = show_columns_type(json.dumps(test_case.show_data_type)) + from_information_schema: str = information_schema_type( + data_type=data_type, + numeric_precision=precision, + numeric_scale=scale, + character_maximum_length=length, + ) + + assert from_show == test_case.expected_type + assert from_information_schema == test_case.expected_type + + +@pytest.mark.parametrize( + "test_case", + [ + InvalidShowColumnTypeTestCase( + description="not json", + raw_data_type="NUMBER(38,0)", + expected_error_fragment="invalid type metadata", + ), + InvalidShowColumnTypeTestCase( + description="json array", + raw_data_type='["FIXED"]', + expected_error_fragment="invalid type metadata", + ), + InvalidShowColumnTypeTestCase( + description="missing type", + raw_data_type='{"precision": 38}', + expected_error_fragment="invalid type metadata", + ), + InvalidShowColumnTypeTestCase( + description="null", + raw_data_type=None, + expected_error_fragment="invalid type metadata", + ), + ], + ids=lambda case: case.description, +) +def test_given_invalid_show_type_when_rendering_then_raises_adapter_error( + test_case: InvalidShowColumnTypeTestCase, +) -> None: + with pytest.raises(AdapterUserError, match=test_case.expected_error_fragment): + _ = show_columns_type(test_case.raw_data_type) + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py new file mode 100644 index 0000000000..d0b88af2fa --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py @@ -0,0 +1,195 @@ +"""SHOW TABLES and SHOW VIEWS rows map onto INFORMATION_SCHEMA.TABLES relation semantics.""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import UTC, datetime + +import pytest + +from sqlbuild.adapter.relations.models import ListedRelation +from sqlbuild.adapters.snowflake._helpers.show_metadata import ( + is_missing_object_error, + listed_relation_from_show_table, + listed_relation_from_show_view, +) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection._test_types import ( + MissingObjectErrorTestCase, + ShowRelationRowTestCase, +) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + FakeSnowflakeProgrammingError, + show_relation_row, +) + +_CREATED_ON: datetime = datetime(2026, 3, 1, 8, 30, tzinfo=UTC) + + +@pytest.mark.parametrize( + "test_case", + [ + ShowRelationRowTestCase( + description="permanent table", + is_view=False, + row=show_relation_row( + created_on=_CREATED_ON, kind="TABLE", retention_time="1", is_external="N" + ), + expected_relation_type="base table", + expected_is_transient=False, + expected_retention_days=1, + ), + ShowRelationRowTestCase( + description="transient table", + is_view=False, + row=show_relation_row( + created_on=_CREATED_ON, kind="TRANSIENT", retention_time="0", is_external="N" + ), + expected_relation_type="base table", + expected_is_transient=True, + expected_retention_days=0, + ), + ShowRelationRowTestCase( + description="dynamic table is a base table", + is_view=False, + row=show_relation_row( + created_on=_CREATED_ON, + kind="TABLE", + retention_time="1", + is_external="N", + is_dynamic="Y", + ), + expected_relation_type="base table", + expected_is_transient=False, + expected_retention_days=1, + ), + ShowRelationRowTestCase( + description="temporary table", + is_view=False, + row=show_relation_row( + created_on=_CREATED_ON, kind="TEMPORARY", retention_time="1", is_external="N" + ), + expected_relation_type="temporary table", + expected_is_transient=False, + expected_retention_days=1, + ), + ShowRelationRowTestCase( + description="external table without retention", + is_view=False, + row=show_relation_row( + created_on=_CREATED_ON, kind="TABLE", retention_time="", is_external="Y" + ), + expected_relation_type="external table", + expected_is_transient=False, + expected_retention_days=None, + ), + ShowRelationRowTestCase( + description="event table", + is_view=False, + row=show_relation_row( + created_on=_CREATED_ON, + kind="TABLE", + retention_time="1", + is_external="N", + is_event="Y", + ), + expected_relation_type="event table", + expected_is_transient=False, + expected_retention_days=1, + ), + ShowRelationRowTestCase( + description="older output without flag columns", + is_view=False, + row=show_relation_row(created_on=_CREATED_ON, kind="TABLE", retention_time=None), + expected_relation_type="base table", + expected_is_transient=False, + expected_retention_days=None, + ), + ShowRelationRowTestCase( + description="view", + is_view=True, + row=show_relation_row( + created_on=_CREATED_ON, is_materialized="false", is_secure="false" + ), + expected_relation_type="view", + expected_is_transient=False, + expected_retention_days=None, + ), + ShowRelationRowTestCase( + description="materialized view as text flag", + is_view=True, + row=show_relation_row(created_on=_CREATED_ON, is_materialized="true"), + expected_relation_type="materialized view", + expected_is_transient=False, + expected_retention_days=None, + ), + ShowRelationRowTestCase( + description="materialized view as boolean flag", + is_view=True, + row=show_relation_row(created_on=_CREATED_ON, is_materialized=True), + expected_relation_type="materialized view", + expected_is_transient=False, + expected_retention_days=None, + ), + ], + ids=lambda case: case.description, +) +def test_given_show_row_when_mapping_then_matches_information_schema_relation( + test_case: ShowRelationRowTestCase, +) -> None: + mappers: tuple[Callable[..., ListedRelation], ...] = ( + listed_relation_from_show_table, + listed_relation_from_show_view, + ) + + listed: ListedRelation = mappers[test_case.is_view](row=test_case.row, database="analytics") + + assert listed.stored_name == "ORDERS" + assert (listed.relation.database, listed.relation.schema, listed.relation.name) == ( + "analytics", + "staging", + "orders", + ) + assert listed.relation.relation_type == test_case.expected_relation_type + assert listed.relation.is_transient is test_case.expected_is_transient + assert listed.relation.retention_days == test_case.expected_retention_days + assert listed.relation.created_at == _CREATED_ON + assert listed.relation.last_altered_at is None + + +@pytest.mark.parametrize( + "test_case", + [ + MissingObjectErrorTestCase( + description="driver error number", + errno=2003, + message="SQL compilation error", + expected_missing=True, + ), + MissingObjectErrorTestCase( + description="message without error number", + errno=None, + message="Schema 'ANALYTICS.STAGING' does not exist or not authorized.", + expected_missing=True, + ), + MissingObjectErrorTestCase( + description="other failure", + errno=390, + message="Metadata service unavailable", + expected_missing=False, + ), + ], + ids=lambda case: case.description, +) +def test_given_show_failure_when_classifying_then_only_missing_schemas_list_nothing( + test_case: MissingObjectErrorTestCase, +) -> None: + error: Exception = ( + RuntimeError(test_case.message), + FakeSnowflakeProgrammingError(test_case.message, errno=test_case.errno or 0), + )[test_case.errno is not None] + + assert is_missing_object_error(error) is test_case.expected_missing + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/website/src/content/docs/docs/concepts/adapters/snowflake.mdx b/website/src/content/docs/docs/concepts/adapters/snowflake.mdx index f3a12b355d..ba0608c913 100644 --- a/website/src/content/docs/docs/concepts/adapters/snowflake.mdx +++ b/website/src/content/docs/docs/concepts/adapters/snowflake.mdx @@ -38,6 +38,18 @@ On connect, SQLBuild applies the connection's role and warehouse, then applies t target's authoritative database and schema. These ensure the session context is set correctly regardless of the user's default settings. +## Planning metadata + +Planning reads relation and column metadata with `SHOW TABLES`, `SHOW VIEWS`, and `SHOW COLUMNS` +once per schema, several schemas at a time. `SHOW` commands are served by Snowflake's metadata +service, so they don't wait for, or resume, a warehouse. Two reads still use `INFORMATION_SCHEMA` +and need a running warehouse: + +- `LAST_ALTERED` for source freshness, read only for the sources the plan observes. +- Columns for a schema whose `SHOW COLUMNS` output reaches Snowflake's 10,000-row limit. + +Run `sqb plan -v` to see each metadata query with its elapsed time and row count. + ## Shared connections across targets Use multiple targets to reuse one Snowflake connection while selecting different authoritative From e75e2b2b7c749baaf603e538cc02118781ac75b6 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 12:35:26 +0100 Subject: [PATCH 5/7] fix(snowflake): treat only a cap-sized SHOW COLUMNS result as truncated SHOW COLUMNS IN SCHEMA returned every row on a real account well past 10,000, so a result is now reread from INFORMATION_SCHEMA only when it is exactly 10,000 rows. SHOW TABLES/VIEWS paging already continues only on a full 10,000-row page. Views and materialized views now report is_transient as None, matching INFORMATION_SCHEMA.TABLES. Refs CHI-455 --- .../docs/concepts/adapters/snowflake.md | 3 ++- .../snowflake/_helpers/show_metadata.py | 4 ++-- .../snowflake/classes/snowflake_adapter.py | 4 ++-- .../snowflake/inspection/_test_types.py | 2 +- .../adapters/snowflake/inspection/helpers.py | 6 +++-- .../inspection/test_inspection_catalog.py | 22 +++++++++++++++++-- .../inspection/test_show_metadata.py | 6 ++--- .../docs/docs/concepts/adapters/snowflake.mdx | 3 ++- 8 files changed, 36 insertions(+), 14 deletions(-) diff --git a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/adapters/snowflake.md b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/adapters/snowflake.md index 25a09a2f4b..6672418ab2 100644 --- a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/adapters/snowflake.md +++ b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/adapters/snowflake.md @@ -49,7 +49,8 @@ service, so they don't wait for, or resume, a warehouse. Two reads still use `IN and need a running warehouse: - `LAST_ALTERED` for source freshness, read only for the sources the plan observes. -- Columns for a schema whose `SHOW COLUMNS` output reaches Snowflake's 10,000-row limit. +- Columns for a schema whose `SHOW COLUMNS` output is exactly 10,000 rows, the documented + limit at which the output may be cut off. Run `sqb plan -v` to see each metadata query with its elapsed time and row count. diff --git a/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py index ac351243be..9f6af4e1b7 100644 --- a/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py +++ b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py @@ -46,7 +46,7 @@ def listed_relation_from_show_view( row=row, database=database, relation_type="materialized view" if _flag(row.get("is_materialized")) else "view", - is_transient=False, + is_transient=None, retention_days=None, ) @@ -65,7 +65,7 @@ def _listed_relation( row: Mapping[str, object], database: str | None, relation_type: str, - is_transient: bool, + is_transient: bool | None, retention_days: int | None, ) -> ListedRelation: stored_name: str = str(row["name"]) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index d36735bbe0..9eb28676db 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -1836,13 +1836,13 @@ def read_schema_relation_listing( def read_schema_column_listing( self, *, connection: _SnowflakeConnection, database: str | None, schema: str ) -> SchemaColumnListing: - """List one schema's columns with SHOW COLUMNS; reread a truncated result exactly.""" + """List one schema's columns with SHOW COLUMNS; reread a result of exactly the row cap.""" scope: str = self._show_schema_scope(database=database, schema=schema) rows: list[dict[str, object]] = self._fetch_show_rows( connection=connection, query=f"SHOW COLUMNS IN SCHEMA {scope}" ) - if len(rows) >= _SHOW_RESULT_LIMIT: + if len(rows) == _SHOW_RESULT_LIMIT: return self._read_information_schema_column_listing( connection=connection, database=database, schema=schema ) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py index de0bfa972c..b34f447571 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py @@ -90,7 +90,7 @@ class ShowRelationRowTestCase: is_view: bool row: dict[str, object] expected_relation_type: str - expected_is_transient: bool + expected_is_transient: bool | None expected_retention_days: int | None diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py index a614aa8cc6..d7a9201e5a 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -397,7 +397,7 @@ def _answer_show_schema_columns( ) for column in fake.columns ) - return rows[:SHOW_RESULT_CAP], tuple((f,) for f in _SHOW_SCHEMA_COLUMN_FIELDS) + return rows, tuple((f,) for f in _SHOW_SCHEMA_COLUMN_FIELDS) def _answer_cursor_bounds( self, *, sql: str, params: tuple[object, ...] @@ -517,7 +517,9 @@ def _table_row(relation: FakeRelation) -> dict[str, object]: "table_schema": relation.schema, "table_name": relation.name, "table_type": relation.table_type, - "is_transient": ("NO", "YES")[relation.is_transient], + "is_transient": (("NO", "YES")[relation.is_transient], None)[ + relation.table_type in _VIEW_TABLE_TYPES + ], "created": _FIXED_CREATED_AT, "last_altered": _FIXED_CREATED_AT, "retention_time": (relation.retention_time, None)[relation.table_type in _VIEW_TABLE_TYPES], diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py index 6a08d99792..f45571fedb 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py @@ -327,11 +327,29 @@ def test_given_unreadable_schema_when_prefetching_best_effort_then_only_real_rea ), ), ShowResultCapTestCase( - description="a truncated column listing is reread from INFORMATION_SCHEMA", + description="exactly 10,000 columns is treated as truncated and reread", relation_count=40, - columns_per_relation=260, + columns_per_relation=250, expected_query_kinds=("show_tables", "show_views", "show_schema_columns", "columns"), ), + ShowResultCapTestCase( + description="10,001 columns is a complete listing", + relation_count=73, + columns_per_relation=137, + expected_query_kinds=("show_tables", "show_views", "show_schema_columns"), + ), + ShowResultCapTestCase( + description="fewer than 10,000 columns is a complete listing", + relation_count=39, + columns_per_relation=256, + expected_query_kinds=("show_tables", "show_views", "show_schema_columns"), + ), + ShowResultCapTestCase( + description="more than 10,000 columns is a complete listing", + relation_count=120, + columns_per_relation=260, + expected_query_kinds=("show_tables", "show_views", "show_schema_columns"), + ), ], ids=lambda case: case.description, ) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py index d0b88af2fa..880bb77dd4 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py @@ -111,7 +111,7 @@ created_on=_CREATED_ON, is_materialized="false", is_secure="false" ), expected_relation_type="view", - expected_is_transient=False, + expected_is_transient=None, expected_retention_days=None, ), ShowRelationRowTestCase( @@ -119,7 +119,7 @@ is_view=True, row=show_relation_row(created_on=_CREATED_ON, is_materialized="true"), expected_relation_type="materialized view", - expected_is_transient=False, + expected_is_transient=None, expected_retention_days=None, ), ShowRelationRowTestCase( @@ -127,7 +127,7 @@ is_view=True, row=show_relation_row(created_on=_CREATED_ON, is_materialized=True), expected_relation_type="materialized view", - expected_is_transient=False, + expected_is_transient=None, expected_retention_days=None, ), ], diff --git a/website/src/content/docs/docs/concepts/adapters/snowflake.mdx b/website/src/content/docs/docs/concepts/adapters/snowflake.mdx index ba0608c913..ff8a6787ae 100644 --- a/website/src/content/docs/docs/concepts/adapters/snowflake.mdx +++ b/website/src/content/docs/docs/concepts/adapters/snowflake.mdx @@ -46,7 +46,8 @@ service, so they don't wait for, or resume, a warehouse. Two reads still use `IN and need a running warehouse: - `LAST_ALTERED` for source freshness, read only for the sources the plan observes. -- Columns for a schema whose `SHOW COLUMNS` output reaches Snowflake's 10,000-row limit. +- Columns for a schema whose `SHOW COLUMNS` output is exactly 10,000 rows, the documented + limit at which the output may be cut off. Run `sqb plan -v` to see each metadata query with its elapsed time and row count. From 91bd1cc5907111254349097de23641a580dfd2ab Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 13:28:39 +0100 Subject: [PATCH 6/7] fix(planner): harden shared inspection reads from review findings - Semantic source fallback lists candidates by name in every schema and matches schemas case-insensitively, as before the catalog; schema-scoped catalogs answer declared schemas from listings first. - Bounded inspections cancel queued reads and re-raise promptly on any exception in the coordinating thread (including Ctrl-C); after a read fails, running reads finish and the lowest-index failure is raised. - Snowflake SHOW is always qualified with a database (the session's when none is given); only a missing schema lists empty, a missing database fails as before. - Column reads choose per-relation SHOW COLUMNS by the request's total size, as before, not per schema; a relation dropped after listing has no columns instead of failing. - A cap-sized SHOW COLUMNS result rereads only the needed relations from INFORMATION_SCHEMA in concurrent IN chunks of at most 50 names. Refs CHI-455 --- .../relations/classes/inspection_catalog.py | 117 +++++++---- .../relations/main/run_bounded_inspections.py | 54 ++++-- src/sqlbuild/adapter/relations/models.py | 1 + src/sqlbuild/adapter/relations/types.py | 14 +- .../snowflake/_helpers/show_metadata.py | 13 ++ .../snowflake/classes/snowflake_adapter.py | 182 +++++++++++++----- .../_helpers/warehouse/semantic_sources.py | 29 ++- .../compiler/planner/main/_test_types.py | 11 ++ .../main/test_snowflake_inspection_budget.py | 44 +++++ .../run_bounded_inspections/_test_types.py | 24 +++ .../main/run_bounded_inspections/helpers.py | 64 ++++++ .../test_run_bounded_inspections.py | 83 ++++++++ .../snowflake/inspection/_test_types.py | 34 ++++ .../adapters/snowflake/inspection/helpers.py | 23 +++ .../inspection/test_inspection_catalog.py | 164 ++++++++++++++++ .../inspection/test_show_metadata.py | 24 ++- .../planner/_helpers/warehouse/_test_types.py | 16 ++ .../planner/_helpers/warehouse/helpers.py | 7 + .../warehouse/test_semantic_sources.py | 122 +++++++++++- 19 files changed, 904 insertions(+), 122 deletions(-) diff --git a/src/sqlbuild/adapter/relations/classes/inspection_catalog.py b/src/sqlbuild/adapter/relations/classes/inspection_catalog.py index 63dab91de2..0be090518a 100644 --- a/src/sqlbuild/adapter/relations/classes/inspection_catalog.py +++ b/src/sqlbuild/adapter/relations/classes/inspection_catalog.py @@ -37,9 +37,16 @@ def __init__(self, *, adapter: Any, connection: Any) -> None: self._concurrency: int = max(1, int(getattr(adapter, "metadata_inspection_concurrency", 1))) self._relation_listings: dict[_SchemaKey, SchemaRelationListing] = {} self._column_listings: dict[_SchemaKey, SchemaColumnListing] = {} + self._capped_column_schemas: set[_SchemaKey] = set() self._relation_requests: dict[_RelationRequestKey, tuple[RelationInfo, ...]] = {} self._relation_columns: dict[_RelationIdentity, tuple[ColumnInfo, ...] | None] = {} + @property + def schema_scoped(self) -> bool: + """Return whether schema listings come from one case-aware read per schema.""" + + return self._reader is not None + @property def concurrency(self) -> int: """Return the bounded number of concurrent reads this adapter supports.""" @@ -237,54 +244,66 @@ def _read_columns_directly(self, *, relations: tuple[RelationInfo, ...]) -> None def _read_columns_by_schema( self, *, reader: SchemaScopedMetadataReader, relations: tuple[RelationInfo, ...] ) -> None: + """Use exact per-relation reads for small requests and one read per schema otherwise.""" + unscoped: list[RelationInfo] = [] - by_schema: dict[_SchemaKey, list[RelationInfo]] = {} + remaining: list[RelationInfo] = [] relation: RelationInfo for relation in relations: if relation.schema is None: unscoped.append(relation) continue - by_schema.setdefault( - self._schema_key(database=relation.database, schema=relation.schema), [] - ).append(relation) + listing: SchemaColumnListing | None = self._column_listings.get( + self._schema_key(database=relation.database, schema=relation.schema) + ) + if listing is None: + remaining.append(relation) + continue + self._relation_columns[relation.identity] = listing.columns_by_stored_name.get( + reader.metadata_name_key(relation.name) + ) if unscoped: self._read_columns_directly(relations=tuple(unscoped)) - listing_keys: list[_SchemaKey] = [] - listing_tasks: list[Callable[[], SchemaColumnListing]] = [] - exact_relations: list[RelationInfo] = [] - schema_key: _SchemaKey - scoped: list[RelationInfo] - for schema_key, scoped in by_schema.items(): - if schema_key in self._column_listings: - continue - if len(scoped) <= reader.exact_column_inspection_limit: - exact_relations.extend(scoped) - continue - listing_keys.append(schema_key) - listing_tasks.append( + if not remaining: + return + if len(remaining) <= reader.exact_column_inspection_limit: + exact_results: list[tuple[ColumnInfo, ...] | None] = run_bounded_inspections( + tasks=tuple( + self._bind_relation_columns(reader=reader, relation=exact) + for exact in remaining + ), + concurrency=self._concurrency, + ) + exact_columns: tuple[ColumnInfo, ...] | None + for relation, exact_columns in zip(remaining, exact_results, strict=True): + self._relation_columns[relation.identity] = exact_columns + return + by_schema: dict[_SchemaKey, list[RelationInfo]] = {} + for relation in remaining: + by_schema.setdefault( + self._schema_key(database=relation.database, schema=relation.schema or ""), [] + ).append(relation) + listings: list[SchemaColumnListing] = run_bounded_inspections( + tasks=tuple( self._bind_column_listing( - reader=reader, database=scoped[0].database, schema=scoped[0].schema or "" + reader=reader, + database=scoped[0].database, + schema=scoped[0].schema or "", + stored_names=self._stored_names(reader=reader, relations=scoped), + known_capped=schema_key in self._capped_column_schemas, ) - ) - exact_tasks: list[Callable[[], tuple[ColumnInfo, ...]]] = [ - self._bind_relation_columns(reader=reader, relation=exact) for exact in exact_relations - ] - results: list[SchemaColumnListing | tuple[ColumnInfo, ...]] = run_bounded_inspections( - tasks=(*listing_tasks, *exact_tasks), concurrency=self._concurrency + for schema_key, scoped in by_schema.items() + ), + concurrency=self._concurrency, ) - listing: SchemaColumnListing | tuple[ColumnInfo, ...] - for schema_key, listing in zip(listing_keys, results[: len(listing_keys)], strict=True): - if isinstance(listing, SchemaColumnListing): - self._column_listings[schema_key] = listing - exact: RelationInfo - exact_columns: SchemaColumnListing | tuple[ColumnInfo, ...] - for exact, exact_columns in zip(exact_relations, results[len(listing_keys) :], strict=True): - if not isinstance(exact_columns, SchemaColumnListing): - self._relation_columns[exact.identity] = exact_columns - for schema_key, scoped in by_schema.items(): - column_listing: SchemaColumnListing | None = self._column_listings.get(schema_key) - if column_listing is None: - continue + schema_key: _SchemaKey + scoped: list[RelationInfo] + column_listing: SchemaColumnListing + for (schema_key, scoped), column_listing in zip(by_schema.items(), listings, strict=True): + if column_listing.complete: + self._column_listings[schema_key] = column_listing + else: + self._capped_column_schemas.add(schema_key) for relation in scoped: self._relation_columns[relation.identity] = ( column_listing.columns_by_stored_name.get( @@ -292,20 +311,36 @@ def _read_columns_by_schema( ) ) + @staticmethod + def _stored_names( + *, reader: SchemaScopedMetadataReader, relations: list[RelationInfo] + ) -> frozenset[str]: + return frozenset(reader.metadata_name_key(relation.name) for relation in relations) + def _bind_column_listing( - self, *, reader: SchemaScopedMetadataReader, database: str | None, schema: str + self, + *, + reader: SchemaScopedMetadataReader, + database: str | None, + schema: str, + stored_names: frozenset[str], + known_capped: bool, ) -> Callable[[], SchemaColumnListing]: def read_listing() -> SchemaColumnListing: return reader.read_schema_column_listing( - connection=self.connection, database=database, schema=schema + connection=self.connection, + database=database, + schema=schema, + stored_names=stored_names, + known_capped=known_capped, ) return read_listing def _bind_relation_columns( self, *, reader: SchemaScopedMetadataReader, relation: RelationInfo - ) -> Callable[[], tuple[ColumnInfo, ...]]: - def read_columns() -> tuple[ColumnInfo, ...]: + ) -> Callable[[], tuple[ColumnInfo, ...] | None]: + def read_columns() -> tuple[ColumnInfo, ...] | None: return reader.read_relation_columns(connection=self.connection, relation=relation) return read_columns diff --git a/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py b/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py index 7389ce64ad..bca0052fd4 100644 --- a/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py +++ b/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py @@ -28,28 +28,46 @@ def run_bounded_inspections[ResultT]( on_complete(index=index, result=result) results.append(result) return results + pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=min(concurrency, len(tasks))) + try: + completed: dict[int, ResultT] = _collect(pool=pool, tasks=tasks, on_complete=on_complete) + except BaseException: + pool.shutdown(wait=False, cancel_futures=True) + raise + pool.shutdown(wait=True) + return [completed[index] for index in range(len(tasks))] + + +def _collect[ResultT]( + *, + pool: ThreadPoolExecutor, + tasks: Sequence[Callable[[], ResultT]], + on_complete: InspectionCompletion[ResultT] | None, +) -> dict[int, ResultT]: + """Stop starting tasks after a failure, let running ones finish, raise the lowest failure.""" + + futures: dict[Future[ResultT], int] = { + cast(Future[ResultT], pool.submit(copy_context().run, task)): index + for index, task in enumerate(tasks) + } completed: dict[int, ResultT] = {} failures: dict[int, BaseException] = {} - with ThreadPoolExecutor(max_workers=min(concurrency, len(tasks))) as pool: - futures: dict[Future[ResultT], int] = { - cast(Future[ResultT], pool.submit(copy_context().run, task)): index - for index, task in enumerate(tasks) - } - future: Future[ResultT] - for future in as_completed(futures): - if future.cancelled(): - continue - future_index: int = futures[future] - error: BaseException | None = future.exception() - if error is not None: - failures[future_index] = error + future: Future[ResultT] + for future in as_completed(futures): + if future.cancelled(): + continue + future_index: int = futures[future] + error: BaseException | None = future.exception() + if error is not None: + if not failures: pending: Future[ResultT] for pending in futures: _ = pending.cancel() - continue - completed[future_index] = future.result() - if on_complete is not None and not failures: - on_complete(index=future_index, result=completed[future_index]) + failures[future_index] = error + continue + completed[future_index] = future.result() + if on_complete is not None and not failures: + on_complete(index=future_index, result=completed[future_index]) if failures: raise failures[min(failures)] - return [completed[index] for index in range(len(tasks))] + return completed diff --git a/src/sqlbuild/adapter/relations/models.py b/src/sqlbuild/adapter/relations/models.py index 01abc1b699..6c7fba5721 100644 --- a/src/sqlbuild/adapter/relations/models.py +++ b/src/sqlbuild/adapter/relations/models.py @@ -31,6 +31,7 @@ class SchemaColumnListing: database: str | None schema: str columns_by_stored_name: dict[str, tuple[ColumnInfo, ...]] + complete: bool = True @dataclass(frozen=True) diff --git a/src/sqlbuild/adapter/relations/types.py b/src/sqlbuild/adapter/relations/types.py index bb9135356c..abce8d457f 100644 --- a/src/sqlbuild/adapter/relations/types.py +++ b/src/sqlbuild/adapter/relations/types.py @@ -31,15 +31,21 @@ def read_schema_relation_listing( ... def read_schema_column_listing( - self, *, connection: Any, database: str | None, schema: str + self, + *, + connection: Any, + database: str | None, + schema: str, + stored_names: frozenset[str], + known_capped: bool = False, ) -> SchemaColumnListing: - """Read every column in one schema.""" + """Read one schema's columns; an incomplete listing covers only ``stored_names``.""" ... def read_relation_columns( self, *, connection: Any, relation: RelationInfo - ) -> tuple[ColumnInfo, ...]: - """Read one relation's columns.""" + ) -> tuple[ColumnInfo, ...] | None: + """Read one relation's columns, or None when it no longer exists.""" ... def metadata_name_key(self, name: str) -> str: diff --git a/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py index 9f6af4e1b7..4c2efefa35 100644 --- a/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py +++ b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py @@ -11,6 +11,8 @@ _MISSING_OBJECT_ERRNO: int = 2003 _MISSING_OBJECT_MESSAGE: str = "does not exist or not authorized" _TRUE_FLAGS: frozenset[str] = frozenset({"Y", "YES", "TRUE"}) +_MISSING_SCHEMA_MARKER: str = "schema '" +_MISSING_DATABASE_MARKER: str = "database '" _TRANSIENT_KIND: str = "TRANSIENT" _TEMPORARY_KIND: str = "TEMPORARY" @@ -60,6 +62,17 @@ def is_missing_object_error(error: Exception) -> bool: ) +def is_missing_schema_error(error: Exception) -> bool: + """Return whether SHOW failed for a missing schema; a missing database is a real failure.""" + + message: str = str(error).lower() + return ( + is_missing_object_error(error) + and _MISSING_SCHEMA_MARKER in message + and _MISSING_DATABASE_MARKER not in message + ) + + def _listed_relation( *, row: Mapping[str, object], diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index 9eb28676db..ddc631f580 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -10,7 +10,7 @@ from datetime import date, datetime from decimal import Decimal from pathlib import Path -from typing import Any, ClassVar +from typing import Any, ClassVar, cast from sqlbuild.adapter.contract.classes.base_adapter import ( BaseAdapter, @@ -116,6 +116,7 @@ ) from sqlbuild.adapters.snowflake._helpers.show_metadata import ( is_missing_object_error, + is_missing_schema_error, listed_relation_from_show_table, listed_relation_from_show_view, ) @@ -149,6 +150,8 @@ _SHOW_RESULT_LIMIT: int = 10_000 _SHOW_COLUMN_KIND: str = "COLUMN" _SHOW_COLUMN_KINDS: frozenset[str] = frozenset({_SHOW_COLUMN_KIND, "VIRTUAL_COLUMN"}) +_CURRENT_DATABASE_ATTRIBUTE: str = "_sqlbuild_current_database" +_UNRESOLVED: object = object() _RELATION_LISTING_COLUMNS: str = ( "table_name, table_schema, table_type, is_transient, created, last_altered, retention_time" ) @@ -1814,7 +1817,9 @@ def read_schema_relation_listing( ) -> SchemaRelationListing: """List one schema's tables and views with SHOW, which needs no running warehouse.""" - scope: str = self._show_schema_scope(database=database, schema=schema) + scope: str = self._show_schema_scope( + connection=connection, database=database, schema=schema + ) tables: list[dict[str, object]] = self._fetch_paged_show_rows( connection=connection, query=f"SHOW TABLES IN SCHEMA {scope}" ) @@ -1834,37 +1839,60 @@ def read_schema_relation_listing( ) def read_schema_column_listing( - self, *, connection: _SnowflakeConnection, database: str | None, schema: str + self, + *, + connection: _SnowflakeConnection, + database: str | None, + schema: str, + stored_names: frozenset[str], + known_capped: bool = False, ) -> SchemaColumnListing: - """List one schema's columns with SHOW COLUMNS; reread a result of exactly the row cap.""" + """List one schema's columns with SHOW COLUMNS; a cap-sized result rereads only names.""" - scope: str = self._show_schema_scope(database=database, schema=schema) - rows: list[dict[str, object]] = self._fetch_show_rows( - connection=connection, query=f"SHOW COLUMNS IN SCHEMA {scope}" - ) - if len(rows) == _SHOW_RESULT_LIMIT: - return self._read_information_schema_column_listing( + if not known_capped: + scope: str = self._show_schema_scope( connection=connection, database=database, schema=schema ) - columns: dict[str, list[ColumnInfo]] = {} - for row in rows: - if str(row.get("kind") or _SHOW_COLUMN_KIND).upper() not in _SHOW_COLUMN_KINDS: - continue - columns.setdefault(str(row["table_name"]), []).append( - ColumnInfo( - name=str(row["column_name"]).lower(), - type=show_columns_type(row["data_type"]), - ) + rows: list[dict[str, object]] = self._fetch_show_rows( + connection=connection, query=f"SHOW COLUMNS IN SCHEMA {scope}" ) - return SchemaColumnListing( - database=database, - schema=schema, - columns_by_stored_name={name: tuple(value) for name, value in columns.items()}, + if len(rows) != _SHOW_RESULT_LIMIT: + columns: dict[str, list[ColumnInfo]] = {} + row: dict[str, object] + for row in rows: + if str(row.get("kind") or _SHOW_COLUMN_KIND).upper() not in _SHOW_COLUMN_KINDS: + continue + columns.setdefault(str(row["table_name"]), []).append( + ColumnInfo( + name=str(row["column_name"]).lower(), + type=show_columns_type(row["data_type"]), + ) + ) + return SchemaColumnListing( + database=database, + schema=schema, + columns_by_stored_name={name: tuple(value) for name, value in columns.items()}, + ) + return replace( + self._read_information_schema_column_listing( + connection=connection, + database=database, + schema=schema, + stored_names=stored_names, + ), + complete=False, ) def _read_information_schema_column_listing( - self, *, connection: _SnowflakeConnection, database: str | None, schema: str + self, + *, + connection: _SnowflakeConnection, + database: str | None, + schema: str, + stored_names: frozenset[str] | None = None, ) -> SchemaColumnListing: + """Read columns from INFORMATION_SCHEMA, unfiltered or for names in capped IN chunks.""" + query: str = ( "SELECT table_name, column_name, data_type, numeric_precision, numeric_scale, " "character_maximum_length FROM " @@ -1875,34 +1903,81 @@ def _read_information_schema_column_listing( if database is not None: query += " AND table_catalog = %s" params.append(self._information_schema_identifier(database)) - query += " ORDER BY table_name, ordinal_position" - rows: list[tuple[Any, ...]] = self._fetch_inspection_rows( - connection=connection, query=query, params=tuple(params) + names: tuple[str, ...] = tuple(sorted(stored_names or ())) + chunks: tuple[tuple[str, ...], ...] = ( + ((),) + if stored_names is None + else tuple( + names[start : start + INSPECTION_IN_LIST_LIMIT] + for start in range(0, len(names), INSPECTION_IN_LIST_LIMIT) + ) + ) + catalog: InspectionCatalog | None = active_inspection_catalog( + adapter=self, connection=connection + ) + chunk_rows: list[list[tuple[Any, ...]]] = run_bounded_inspections( + tasks=tuple( + self._bind_freshness_chunk( + connection=connection, + query=query + + (f" AND table_name IN ({', '.join(['%s'] * len(chunk))})" if chunk else "") + + " ORDER BY table_name, ordinal_position", + params=(*params, *chunk), + ) + for chunk in chunks + ), + concurrency=1 if catalog is None else catalog.concurrency, ) columns: dict[str, list[ColumnInfo]] = {} + rows: list[tuple[Any, ...]] row: tuple[Any, ...] - for row in rows: - columns.setdefault(str(row[0]), []).append( - ColumnInfo( - name=str(row[1]).lower(), - type=information_schema_type( - data_type=str(row[2]), - numeric_precision=row[3], - numeric_scale=row[4], - character_maximum_length=row[5], - ), + for rows in chunk_rows: + for row in rows: + columns.setdefault(str(row[0]), []).append( + ColumnInfo( + name=str(row[1]).lower(), + type=information_schema_type( + data_type=str(row[2]), + numeric_precision=row[3], + numeric_scale=row[4], + character_maximum_length=row[5], + ), + ) ) - ) return SchemaColumnListing( database=database, schema=schema, columns_by_stored_name={name: tuple(value) for name, value in columns.items()}, ) - def _show_schema_scope(self, *, database: str | None, schema: str) -> str: - if database is None: - return self.render_identifier(schema) - return f"{self.render_identifier(database)}.{self.render_identifier(schema)}" + def _show_schema_scope( + self, *, connection: _SnowflakeConnection, database: str | None, schema: str + ) -> str: + """Qualify SHOW with the database so a missing one is reported, never listed empty.""" + + resolved: str | None = ( + database if database is not None else self._current_database(connection=connection) + ) + if resolved is None: + raise AdapterUserError( + message=f"Snowflake schema {schema} needs a database; the session has none" + ) + return f"{self.render_identifier(resolved)}.{self.render_identifier(schema)}" + + @staticmethod + def _current_database(*, connection: _SnowflakeConnection) -> str | None: + cached: object = getattr(connection, _CURRENT_DATABASE_ATTRIBUTE, _UNRESOLVED) + if cached is not _UNRESOLVED: + return cast(str | None, cached) + cursor: Any = connection.cursor() + try: + cursor.execute("SELECT CURRENT_DATABASE()") + row: tuple[Any, ...] | None = cursor.fetchone() + finally: + cursor.close() + current: str | None = None if row is None or row[0] is None else str(row[0]) + setattr(connection, _CURRENT_DATABASE_ATTRIBUTE, current) + return current def _fetch_paged_show_rows( self, *, connection: _SnowflakeConnection, query: str @@ -1933,7 +2008,7 @@ def fetch() -> list[dict[str, object]]: try: cursor.execute(query) except Exception as error: - if is_missing_object_error(error): + if is_missing_schema_error(error): return [] raise names: list[str] = [str(column[0]).lower() for column in cursor.description or ()] @@ -1945,10 +2020,15 @@ def fetch() -> list[dict[str, object]]: def read_relation_columns( self, *, connection: _SnowflakeConnection, relation: RelationInfo - ) -> tuple[ColumnInfo, ...]: - """Read one relation's columns with an exact SHOW COLUMNS call.""" + ) -> tuple[ColumnInfo, ...] | None: + """Read one relation's columns with SHOW COLUMNS; None once it no longer exists.""" - return self._show_columns_for_relation(connection=connection, relation=relation) + try: + return self._show_columns_for_relation(connection=connection, relation=relation) + except Exception as error: + if is_missing_object_error(error): + return None + raise def metadata_name_key(self, name: str) -> str: """Unquoted logical identifiers match the uppercase name Snowflake stores.""" @@ -2128,12 +2208,14 @@ def get_columns_for_relations( ) ) if len(ordered_relations) <= _EXACT_COLUMN_INSPECTION_LIMIT: - return { - relation.identity: self._show_columns_for_relation( + exact: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = {} + for relation in ordered_relations: + columns: tuple[ColumnInfo, ...] | None = self.read_relation_columns( connection=connection, relation=relation ) - for relation in ordered_relations - } + if columns is not None: + exact[relation.identity] = columns + return exact return self._get_columns_for_relations_bulk( connection=connection, relations=ordered_relations ) diff --git a/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py b/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py index 54d68acc94..351b91bded 100644 --- a/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py +++ b/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py @@ -6,6 +6,8 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo +from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog +from sqlbuild.adapter.relations.main.active_inspection_catalog import active_inspection_catalog from sqlbuild.adapter.relations.main.get_columns_for_inspection import get_columns_for_inspection from sqlbuild.adapter.relations.main.list_relations_for_inspection import ( list_relations_for_inspection, @@ -89,12 +91,15 @@ def get_semantic_source_columns( def _list_source_candidates( *, adapter: BaseAdapter, connection: Any, database: str | None, entries: tuple[SourceEntry, ...] ) -> tuple[RelationInfo, ...]: - """List schema-qualified sources from their schemas and unqualified ones by name.""" + """List candidates by name in any schema; schema listings first answer scoped catalogs.""" - scoped: tuple[SourceEntry, ...] = tuple(entry for entry in entries if entry.schema is not None) - unscoped: tuple[SourceEntry, ...] = tuple(entry for entry in entries if entry.schema is None) + pending: tuple[SourceEntry, ...] = entries relations: list[RelationInfo] = [] - if scoped: + catalog: InspectionCatalog | None = active_inspection_catalog( + adapter=adapter, connection=connection + ) + scoped: tuple[SourceEntry, ...] = tuple(entry for entry in entries if entry.schema is not None) + if catalog is not None and catalog.schema_scoped and scoped: relations.extend( list_relations_for_inspection( adapter=adapter, @@ -104,14 +109,26 @@ def _list_source_candidates( names=tuple(sorted({entry.table or entry.name for entry in scoped})), ) ) - if unscoped: + pending = tuple( + entry for entry in entries if not _has_candidate(entry=entry, relations=relations) + ) + if pending: relations.extend( list_relations_for_inspection( adapter=adapter, connection=connection, database=database, schemas=None, - names=tuple(sorted({entry.table or entry.name for entry in unscoped})), + names=tuple(sorted({entry.table or entry.name for entry in pending})), ) ) return tuple(dict.fromkeys(relations)) + + +def _has_candidate(*, entry: SourceEntry, relations: list[RelationInfo]) -> bool: + name: str = (entry.table or entry.name).casefold() + return any( + relation.name.casefold() == name + and (entry.schema is None or (relation.schema or "").casefold() == entry.schema.casefold()) + for relation in relations + ) diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py b/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py index 92ebedd06c..0d0dec6efc 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py @@ -130,3 +130,14 @@ class SnowflakeReplanTestCase: description: str expected_listing_reads: int + + +@dataclass(frozen=True) +class SnowflakeManySchemasTestCase: + """Many small model schemas planned against a recording warehouse.""" + + description: str + schema_count: int + models_per_schema: int + expected_per_relation_column_reads: int + expected_schema_column_reads: int diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py b/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py index 26122e543e..0aaa5471bc 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py @@ -15,6 +15,7 @@ from tests.integration.src.sqlbuild.compiler.planner.main._test_types import ( SnowflakeCursorBoundsBudgetTestCase, SnowflakeInspectionBudgetTestCase, + SnowflakeManySchemasTestCase, SnowflakeReplanTestCase, ) from tests.integration.src.sqlbuild.compiler.planner.main.helpers import ( @@ -165,5 +166,48 @@ def test_given_planned_project_when_planning_again_then_cursor_bounds_are_read_a assert len(warehouse.queries_of_kind("show_tables")) == test_case.expected_listing_reads +@pytest.mark.parametrize( + "test_case", + [ + SnowflakeManySchemasTestCase( + description="many small schemas read model columns once per schema", + schema_count=24, + models_per_schema=2, + expected_per_relation_column_reads=30, + expected_schema_column_reads=24, + ) + ], + ids=lambda case: case.description, +) +def test_given_many_small_schemas_when_planning_then_columns_are_read_per_schema( + test_case: SnowflakeManySchemasTestCase, tmp_path: Path +) -> None: + project: SyntheticSnowflakeProject = write_synthetic_snowflake_project( + project_dir=tmp_path / "orders_platform", + model_schemas=tuple(f"region_{index}" for index in range(test_case.schema_count)), + models_per_schema=test_case.models_per_schema, + sources=30, + unmanaged_relations_per_schema=5, + ) + warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( + relations=project.relations + ) + + result: CompilePipelineResult = plan_offline_snowflake_project( + project=project, warehouse=warehouse + ) + + assert len(result.plan_output.model_entries) == len(project.model_schemas) + assert ( + len(warehouse.queries_of_kind("show_columns")) + == test_case.expected_per_relation_column_reads + ) + assert ( + len(warehouse.queries_of_kind("show_schema_columns")) + == test_case.expected_schema_column_reads + ) + assert warehouse.queries_of_kind("columns") == () + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/_test_types.py b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/_test_types.py index 5b128cf7b8..1315c6f7ac 100644 --- a/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/_test_types.py +++ b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/_test_types.py @@ -23,3 +23,27 @@ class BoundedInspectionFailureTestCase: concurrency: int failing_indexes: frozenset[int] expected_error: str + + +@dataclass(frozen=True) +class InterruptedInspectionTestCase: + """A KeyboardInterrupt while bounded reads are queued.""" + + description: str + task_count: int + concurrency: int + task_seconds: float + interrupt_in_task: bool + expected_max_seconds: float + expected_max_started: int + + +@dataclass(frozen=True) +class InFlightFailureTestCase: + """Failures from tasks that were already running when the first failure arrived.""" + + description: str + task_seconds: tuple[float, ...] + failing_indexes: frozenset[int] + concurrency: int + expected_error: str diff --git a/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/helpers.py b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/helpers.py index 4c28e099fb..197739c60c 100644 --- a/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/helpers.py +++ b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/helpers.py @@ -45,3 +45,67 @@ def run() -> int: self._active -= 1 return run + + +def _return_index(index: int) -> int: + return index + + +def _interrupt(index: int) -> int: + raise KeyboardInterrupt(f"interrupted at read {index}") + + +class StartRecordingTasks: + """Tasks with per-index latency that record which indexes started.""" + + def __init__( + self, + *, + latencies: tuple[float, ...], + outcomes: dict[int, Callable[[int], int]], + ) -> None: + self._latencies: tuple[float, ...] = latencies + self._outcomes: dict[int, Callable[[int], int]] = outcomes + self._lock: threading.Lock = threading.Lock() + self.started: list[int] = [] + + def tasks(self) -> tuple[Callable[[], int], ...]: + """Return one task per configured latency.""" + + return tuple(self._task(index=index) for index in range(len(self._latencies))) + + def _task(self, *, index: int) -> Callable[[], int]: + def run() -> int: + with self._lock: + self.started.append(index) + time.sleep(self._latencies[index]) + return self._outcomes.get(index, _return_index)(index) + + return run + + +def interrupting_outcomes(*, interrupt_in_task: bool) -> dict[int, Callable[[int], int]]: + """Interrupt read 0 inside the task, or nowhere when the callback interrupts instead.""" + + outcomes: tuple[dict[int, Callable[[int], int]], ...] = ({}, {0: _interrupt}) + return outcomes[interrupt_in_task] + + +def failing_outcomes(indexes: frozenset[int]) -> dict[int, Callable[[int], int]]: + """Fail the given reads with their index in the message.""" + + return dict.fromkeys(indexes, _fail) + + +def interrupt_on_first_completion(*, interrupt: bool) -> Callable[..., None]: + """Return a completion callback that interrupts the coordinating thread when asked.""" + + def complete(*, index: int, result: int) -> None: + del result + (_ignore_completion, _interrupt)[interrupt](index) + + return complete + + +def _ignore_completion(index: int) -> int: + return index diff --git a/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/test_run_bounded_inspections.py b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/test_run_bounded_inspections.py index 3577b1ddc0..6c01d68136 100644 --- a/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/test_run_bounded_inspections.py +++ b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/test_run_bounded_inspections.py @@ -3,6 +3,7 @@ from __future__ import annotations import threading +import time import pytest @@ -10,9 +11,15 @@ from tests.unit.src.sqlbuild.adapter.relations.main.run_bounded_inspections._test_types import ( BoundedInspectionFailureTestCase, BoundedInspectionTestCase, + InFlightFailureTestCase, + InterruptedInspectionTestCase, ) from tests.unit.src.sqlbuild.adapter.relations.main.run_bounded_inspections.helpers import ( ConcurrencyTracker, + StartRecordingTasks, + failing_outcomes, + interrupt_on_first_completion, + interrupting_outcomes, ) @@ -89,5 +96,81 @@ def test_given_failing_read_when_running_then_raises_lowest_index_failure( ) +@pytest.mark.parametrize( + "test_case", + [ + InterruptedInspectionTestCase( + description="interrupt raised by read 0 stops queued reads", + task_count=40, + concurrency=8, + task_seconds=0.2, + interrupt_in_task=True, + expected_max_seconds=0.6, + expected_max_started=9, + ), + InterruptedInspectionTestCase( + description="interrupt raised while reporting read 0 returns without waiting", + task_count=40, + concurrency=8, + task_seconds=0.2, + interrupt_in_task=False, + expected_max_seconds=0.6, + expected_max_started=16, + ), + ], + ids=lambda case: case.description, +) +def test_given_keyboard_interrupt_when_reads_are_queued_then_raises_promptly_without_new_reads( + test_case: InterruptedInspectionTestCase, +) -> None: + latencies: tuple[float, ...] = (0.0,) * test_case.interrupt_in_task + ( + test_case.task_seconds, + ) * (test_case.task_count - test_case.interrupt_in_task) + recorder: StartRecordingTasks = StartRecordingTasks( + latencies=latencies, + outcomes=interrupting_outcomes(interrupt_in_task=test_case.interrupt_in_task), + ) + started_at: float = time.monotonic() + + with pytest.raises(KeyboardInterrupt): + _ = run_bounded_inspections( + tasks=recorder.tasks(), + concurrency=test_case.concurrency, + on_complete=interrupt_on_first_completion(interrupt=not test_case.interrupt_in_task), + ) + + elapsed: float = time.monotonic() - started_at + started_when_raised: int = len(recorder.started) + time.sleep(test_case.task_seconds * 2) + assert elapsed < test_case.expected_max_seconds + assert len(recorder.started) == started_when_raised + assert len(recorder.started) <= test_case.expected_max_started + + +@pytest.mark.parametrize( + "test_case", + [ + InFlightFailureTestCase( + description="slow lower-index failure wins over an earlier higher-index failure", + task_seconds=(0.0, 0.15, 0.0, 0.0, 0.0, 0.3, 0.3, 0.3), + failing_indexes=frozenset({1, 4}), + concurrency=8, + expected_error="read 1 failed", + ), + ], + ids=lambda case: case.description, +) +def test_given_failures_from_running_reads_when_collecting_then_reports_lowest_index( + test_case: InFlightFailureTestCase, +) -> None: + recorder: StartRecordingTasks = StartRecordingTasks( + latencies=test_case.task_seconds, + outcomes=failing_outcomes(test_case.failing_indexes), + ) + + with pytest.raises(RuntimeError, match=test_case.expected_error): + _ = run_bounded_inspections(tasks=recorder.tasks(), concurrency=test_case.concurrency) + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py index b34f447571..e103fe167a 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py @@ -102,6 +102,7 @@ class MissingObjectErrorTestCase: errno: int | None message: str expected_missing: bool + expected_missing_schema: bool @dataclass(frozen=True) @@ -112,3 +113,36 @@ class ShowResultCapTestCase: relation_count: int columns_per_relation: int expected_query_kinds: tuple[str, ...] + + +@dataclass(frozen=True) +class ShowScopeTestCase: + """A schema listing whose database is missing or comes from the session.""" + + description: str + database: str | None + expected_relation_count: int + expected_error_fragment: str | None + expected_attempted_sql: tuple[str, ...] + + +@dataclass(frozen=True) +class DroppedRelationTestCase: + """Columns requested for a relation dropped after it was listed.""" + + description: str + relation_names: tuple[str, ...] + dropped_name: str + expected_column_relations: tuple[str, ...] + + +@dataclass(frozen=True) +class CappedColumnFallbackTestCase: + """Columns needed from a schema whose SHOW COLUMNS output is exactly the cap.""" + + description: str + relation_count: int + columns_per_relation: int + request_batches: tuple[int, ...] + expected_query_kinds: tuple[str, ...] + expected_in_list_sizes: tuple[int, ...] diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py index d7a9201e5a..ab0fc0f70d 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -48,6 +48,7 @@ re.IGNORECASE | re.DOTALL, ) _ANY_PATTERN: re.Pattern[str] = re.compile(r"") +_CURRENT_DATABASE_PATTERN: re.Pattern[str] = re.compile(r"^SELECT CURRENT_DATABASE\(\)$") _IN_LIST_PATTERN: re.Pattern[str] = re.compile(r"\bIN\s*\(([^()]*)\)", re.IGNORECASE) _TABLES_PATTERN: re.Pattern[str] = re.compile(r"information_schema\.tables\b", re.IGNORECASE) _COLUMNS_PATTERN: re.Pattern[str] = re.compile(r"information_schema\.columns\b", re.IGNORECASE) @@ -59,6 +60,7 @@ (re.compile(r"^SHOW VIEWS IN SCHEMA "), "show_views"), (_SHOW_SCHEMA_COLUMNS_PATTERN, "show_schema_columns"), (_CURSOR_BOUND_PATTERN, "cursor_bounds"), + (_CURRENT_DATABASE_PATTERN, "session"), (_OTHER_INFORMATION_SCHEMA_PATTERN, "other_metadata"), (_SHOW_PATTERN, "other_metadata"), (_ANY_PATTERN, "data"), @@ -286,6 +288,7 @@ def _answer(self, *, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[A (_SHOW_SCHEMA_RELATIONS_PATTERN, self._answer_show_schema_relations), (_SHOW_SCHEMA_COLUMNS_PATTERN, self._answer_show_schema_columns), (_CURSOR_BOUND_PATTERN, self._answer_cursor_bounds), + (_CURRENT_DATABASE_PATTERN, self._answer_current_database), (_ANY_PATTERN, _answer_status), ) route: tuple[re.Pattern[str], Callable[..., Any]] = next( @@ -335,12 +338,21 @@ def _scope_metadata_relation(self, match: re.Match[str]) -> str: database: str = (match.group("database") or self.current_database).replace("'", "''") return f"(SELECT * FROM {table} WHERE table_catalog = '{database}') AS {table}" + def _answer_current_database( + self, *, sql: str, params: tuple[object, ...] + ) -> tuple[list[tuple[Any, ...]], Any]: + del sql, params + return [(self.current_database,)], (("CURRENT_DATABASE()",),) + def _schema_relations(self, *, scope: str) -> tuple[FakeRelation, ...]: parts: tuple[str, ...] = ( self.current_database, *(part.strip('"') for part in scope.split(".")), )[-2:] _FAILED_SCHEMA_ACTIONS[parts[-1] in self.failing_metadata_schemas](scope) + _MISSING_DATABASE_ACTIONS[parts[0] not in {fake.database for fake in self.relations}]( + parts[0] + ) matches: tuple[FakeRelation, ...] = tuple( filter( lambda fake: (fake.database, fake.schema)[-len(parts) :] == parts, self.relations @@ -454,6 +466,17 @@ def _raise_failed_schema(scope: str) -> None: raise FakeSnowflakeProgrammingError(f"Metadata service unavailable for {scope}.", errno=390) +def _raise_missing_database(database: str) -> None: + raise FakeSnowflakeProgrammingError( + f"SQL compilation error: Database '{database}' does not exist or not authorized.", + errno=2003, + ) + + +_MISSING_DATABASE_ACTIONS: dict[bool, Callable[[str], None]] = { + False: _accept_relation, + True: _raise_missing_database, +} _MISSING_SCHEMA_ACTIONS: dict[bool, Callable[[str], None]] = { False: _accept_relation, True: _raise_missing_schema, diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py index f45571fedb..1525e9e605 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections import Counter +from dataclasses import replace import pytest @@ -17,12 +18,15 @@ from sqlbuild.adapter.relations.constants import INSPECTION_IN_LIST_LIMIT from sqlbuild.adapter.relations.main.open_inspection_catalog import open_inspection_catalog from tests.unit.src.sqlbuild.adapters.snowflake.inspection._test_types import ( + CappedColumnFallbackTestCase, ColumnRequestEquivalenceTestCase, + DroppedRelationTestCase, FreshnessErrorTestCase, FreshnessReuseTestCase, RelationRequestEquivalenceTestCase, RepeatedLookupTestCase, ShowResultCapTestCase, + ShowScopeTestCase, SpeculativePrefetchTestCase, ) from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( @@ -383,5 +387,165 @@ def test_given_schema_beyond_show_cap_when_listing_then_matches_direct_adapter( assert tuple(query.kind for query in warehouse.queries) == test_case.expected_query_kinds +@pytest.mark.parametrize( + "test_case", + [ + ShowScopeTestCase( + description="session database is resolved once and qualifies every SHOW", + database=None, + expected_relation_count=69, + expected_error_fragment=None, + expected_attempted_sql=( + "SELECT CURRENT_DATABASE()", + 'SHOW TABLES IN SCHEMA "ANALYTICS"."STAGING" LIMIT 10000', + 'SHOW VIEWS IN SCHEMA "ANALYTICS"."STAGING" LIMIT 10000', + 'SHOW TABLES IN SCHEMA "ANALYTICS"."MARTS" LIMIT 10000', + 'SHOW VIEWS IN SCHEMA "ANALYTICS"."MARTS" LIMIT 10000', + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_no_database_when_listing_schemas_then_show_uses_the_session_database( + test_case: ShowScopeTestCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=build_inspection_catalog_relations() + ) + catalog: InspectionCatalog = InspectionCatalog(adapter=adapter, connection=connection) + + staging: tuple[RelationInfo, ...] = catalog.list_relations( + database=test_case.database, schemas=("staging",) + ) + marts: tuple[RelationInfo, ...] = catalog.list_relations( + database=test_case.database, schemas=("marts",) + ) + + assert len(staging) + len(marts) == test_case.expected_relation_count + assert {relation.database for relation in (*staging, *marts)} == {None} + assert tuple(warehouse.attempted_sql) == test_case.expected_attempted_sql + + +@pytest.mark.parametrize( + "test_case", + [ + ShowScopeTestCase( + description="a missing database fails instead of listing nothing", + database="restricted_catalog", + expected_relation_count=0, + expected_error_fragment="Database 'RESTRICTED_CATALOG' does not exist", + expected_attempted_sql=( + 'SHOW TABLES IN SCHEMA "RESTRICTED_CATALOG"."STAGING" LIMIT 10000', + ), + ) + ], + ids=lambda case: case.description, +) +def test_given_missing_database_when_listing_schema_then_inspection_fails( + test_case: ShowScopeTestCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=build_inspection_catalog_relations() + ) + catalog: InspectionCatalog = InspectionCatalog(adapter=adapter, connection=connection) + + with pytest.raises(RuntimeError, match=test_case.expected_error_fragment or ""): + _ = catalog.list_relations(database=test_case.database, schemas=("staging",)) + + assert tuple(warehouse.attempted_sql) == test_case.expected_attempted_sql + + +@pytest.mark.parametrize( + "test_case", + [ + DroppedRelationTestCase( + description="a relation dropped after listing has no columns instead of failing", + relation_names=("orders", "customers_v", "transient_events"), + dropped_name="orders_archive", + expected_column_relations=("customers_v", "orders", "transient_events"), + ) + ], + ids=lambda case: case.description, +) +def test_given_dropped_relation_when_reading_exact_columns_then_it_is_missing( + test_case: DroppedRelationTestCase, +) -> None: + adapter, connection, _ = build_offline_snowflake(relations=build_inspection_catalog_relations()) + listed: tuple[RelationInfo, ...] = adapter.list_relations( + connection=connection, + database="analytics", + schemas=("staging",), + names=test_case.relation_names, + ) + relations: tuple[RelationInfo, ...] = ( + *listed, + replace(listed[0], name=test_case.dropped_name), + ) + + direct: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( + adapter.get_columns_for_relations(connection=connection, relations=relations) + ) + served: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = InspectionCatalog( + adapter=adapter, connection=connection + ).get_columns_for_relations(relations=relations) + + assert served == direct + assert tuple(sorted(identity[2] for identity in served)) == test_case.expected_column_relations + + +@pytest.mark.parametrize( + "test_case", + [ + CappedColumnFallbackTestCase( + description="a cap-sized schema rereads only needed relations in capped chunks", + relation_count=100, + columns_per_relation=100, + request_batches=(60, 40), + expected_query_kinds=("columns", "columns", "columns", "show_schema_columns"), + expected_in_list_sizes=(10, 40, 50), + ) + ], + ids=lambda case: case.description, +) +def test_given_cap_sized_column_listing_when_reading_then_only_needed_relations_are_reread( + test_case: CappedColumnFallbackTestCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=build_wide_schema_relations( + relation_count=test_case.relation_count, + columns_per_relation=test_case.columns_per_relation, + ) + ) + listed: tuple[RelationInfo, ...] = tuple( + sorted( + adapter.list_relations( + connection=connection, database="analytics", schemas=("inventory",) + ), + key=lambda relation: relation.name, + ) + ) + direct: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = ( + adapter.get_columns_for_relations(connection=connection, relations=listed) + ) + warehouse.reset() + catalog: InspectionCatalog = InspectionCatalog(adapter=adapter, connection=connection) + offsets: tuple[int, ...] = tuple( + sum(test_case.request_batches[:index]) for index in range(len(test_case.request_batches)) + ) + + served: dict[tuple[str | None, str | None, str], tuple[ColumnInfo, ...]] = {} + for offset, size in zip(offsets, test_case.request_batches, strict=True): + served.update(catalog.get_columns_for_relations(relations=listed[offset : offset + size])) + + assert served == direct + assert tuple(sorted(query.kind for query in warehouse.queries)) == ( + test_case.expected_query_kinds + ) + assert ( + tuple(sorted(query.largest_in_list for query in warehouse.queries_of_kind("columns"))) + == test_case.expected_in_list_sizes + ) + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py index 880bb77dd4..cfc2cc5dbb 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py @@ -10,6 +10,7 @@ from sqlbuild.adapter.relations.models import ListedRelation from sqlbuild.adapters.snowflake._helpers.show_metadata import ( is_missing_object_error, + is_missing_schema_error, listed_relation_from_show_table, listed_relation_from_show_view, ) @@ -160,22 +161,40 @@ def test_given_show_row_when_mapping_then_matches_information_schema_relation( "test_case", [ MissingObjectErrorTestCase( - description="driver error number", + description="driver error number alone", errno=2003, message="SQL compilation error", expected_missing=True, + expected_missing_schema=False, ), MissingObjectErrorTestCase( - description="message without error number", + description="missing schema message", errno=None, message="Schema 'ANALYTICS.STAGING' does not exist or not authorized.", expected_missing=True, + expected_missing_schema=True, + ), + MissingObjectErrorTestCase( + description="missing database fails the plan", + errno=2003, + message="SQL compilation error: Database 'ANALYTICS' does not exist or not authorized.", + expected_missing=True, + expected_missing_schema=False, + ), + MissingObjectErrorTestCase( + description="missing table", + errno=2003, + message="SQL compilation error: Table 'ANALYTICS.RAW.ORDERS' does not exist or not " + "authorized.", + expected_missing=True, + expected_missing_schema=False, ), MissingObjectErrorTestCase( description="other failure", errno=390, message="Metadata service unavailable", expected_missing=False, + expected_missing_schema=False, ), ], ids=lambda case: case.description, @@ -189,6 +208,7 @@ def test_given_show_failure_when_classifying_then_only_missing_schemas_list_noth )[test_case.errno is not None] assert is_missing_object_error(error) is test_case.expected_missing + assert is_missing_schema_error(error) is test_case.expected_missing_schema if __name__ == "__main__": diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py index 57accfef1e..fbb16d62de 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py @@ -25,3 +25,19 @@ class EligibleTargetMaxFailureTestCase: description: str failing_relation: str expected_error_fragment: str + + +@dataclass(frozen=True) +class SemanticSourceSchemaCaseCase: + description: str + open_catalog: bool + declared_schema: str + expected_columns: dict[str, tuple[str, ...]] + + +@dataclass(frozen=True) +class SnowflakeSourceCandidateCase: + description: str + declared_schema: str + expected_candidates: tuple[tuple[str | None, str | None, str], ...] + expected_query_kinds: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py index 40ec4e109f..a5c2f7b300 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py @@ -7,6 +7,7 @@ from collections.abc import Callable from typing import Any +from sqlbuild.adapter.contract.models import ColumnInfo from sqlbuild.compiler.planner._helpers.warehouse.snapshot import ( _CursorModelInfo, _PhysicalCursorQuery, @@ -143,3 +144,9 @@ def build_eligible_max_cursor_models(count: int) -> list[_CursorModelInfo]: ) for index in range(count) ] + + +def column_names(columns: tuple[ColumnInfo, ...]) -> tuple[str, ...]: + """Return inspected column names in order.""" + + return tuple(column.name for column in columns) diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_semantic_sources.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_semantic_sources.py index b202e513fc..6329dc1072 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_semantic_sources.py +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_semantic_sources.py @@ -3,23 +3,34 @@ from __future__ import annotations from collections.abc import Callable +from contextlib import AbstractContextManager, nullcontext from pathlib import Path from unittest.mock import Mock import duckdb import pytest -from sqlbuild.adapter.contract.models import ColumnInfo +from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo +from sqlbuild.adapter.relations.main.open_inspection_catalog import open_inspection_catalog from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.compiler.compile.models import CompiledProject from sqlbuild.compiler.discovery.main.discover import discover_project_inputs from sqlbuild.compiler.pipeline.main.project import compile_project from sqlbuild.compiler.planner._helpers.warehouse.semantic_sources import ( + _list_source_candidates, get_semantic_source_columns, ) +from sqlbuild.spec.contracts.models import SourceEntry +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + FakeRelation, + build_offline_snowflake, +) from tests.unit.src.sqlbuild.compiler.planner._helpers.warehouse._test_types import ( SemanticSourceBatchCase, + SemanticSourceSchemaCaseCase, + SnowflakeSourceCandidateCase, ) +from tests.unit.src.sqlbuild.compiler.planner._helpers.warehouse.helpers import column_names @pytest.mark.parametrize( @@ -70,5 +81,114 @@ def test_given_multiple_missing_source_schemas_when_inspecting_then_metadata_is_ assert columns.call_count == test_case.expected_batches +@pytest.mark.parametrize( + "test_case", + [ + SemanticSourceSchemaCaseCase( + description="uppercase declared schema finds a lowercase duckdb schema", + open_catalog=False, + declared_schema="RAW", + expected_columns={"orders": ("id", "status")}, + ), + SemanticSourceSchemaCaseCase( + description="uppercase declared schema finds it through the planning catalog", + open_catalog=True, + declared_schema="RAW", + expected_columns={"orders": ("id", "status")}, + ), + ], + ids=lambda case: case.description, +) +def test_given_source_schema_case_differs_when_inspecting_then_columns_are_found( + test_case: SemanticSourceSchemaCaseCase, + tmp_path: Path, + write_repo_files: Callable[[Path, dict[str, str]], None], +) -> None: + write_repo_files( + tmp_path, + { + "sqlbuild_project.toml": 'name = "source_case"\nadapter = "duckdb"\n', + "sources/orders.yml": ( + "sources:\n - name: orders\n" + f" schema: {test_case.declared_schema}\n table: orders\n" + ), + "models/open_orders.sql": 'MODEL ();\nSELECT id FROM __source("orders")', + }, + ) + adapter: DuckDbAdapter = DuckDbAdapter() + project: CompiledProject = compile_project( + discovered_inputs=discover_project_inputs(project_dir=tmp_path), adapter=adapter + ) + with duckdb.connect() as connection: + connection.execute("CREATE SCHEMA raw; CREATE TABLE raw.orders(id INTEGER, status VARCHAR)") + scope: AbstractContextManager[object] = ( + nullcontext(), + open_inspection_catalog(adapter=adapter, connection=connection), + )[test_case.open_catalog] + with scope: + result: dict[str, tuple[ColumnInfo, ...]] = get_semantic_source_columns( + project=project, + adapter=adapter, + connection=connection, + source_read_map={source.name: source.source_entry for source in project.sources}, + columns={}, + selected_keys=frozenset(model.key for model in project.models), + ) + + assert {name: column_names(columns) for name, columns in result.items()} == ( + test_case.expected_columns + ) + + +@pytest.mark.parametrize( + "test_case", + [ + SnowflakeSourceCandidateCase( + description="unquoted schema is answered from its SHOW listing", + declared_schema="raw", + expected_candidates=(("analytics", "raw", "orders"),), + expected_query_kinds=("show_tables", "show_views"), + ), + SnowflakeSourceCandidateCase( + description="quoted lowercase schema falls back to a database-wide name lookup", + declared_schema="landing", + expected_candidates=( + ("analytics", "raw", "orders"), + ("analytics", "landing", "orders"), + ), + expected_query_kinds=("tables",), + ), + ], + ids=lambda case: case.description, +) +def test_given_snowflake_source_schema_when_listing_candidates_then_finds_relation( + test_case: SnowflakeSourceCandidateCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=( + FakeRelation(database="ANALYTICS", schema="RAW", name="ORDERS"), + FakeRelation(database="ANALYTICS", schema="landing", name="ORDERS"), + ) + ) + + with open_inspection_catalog(adapter=adapter, connection=connection): + candidates: tuple[RelationInfo, ...] = _list_source_candidates( + adapter=adapter, + connection=connection, + database="analytics", + entries=( + SourceEntry( + name="orders", + database="analytics", + schema=test_case.declared_schema, + table="orders", + ), + ), + ) + + assert tuple(relation.identity for relation in candidates) == test_case.expected_candidates + assert tuple(query.kind for query in warehouse.queries) == test_case.expected_query_kinds + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) From ed7c6d3abb8464ba28eee55135d3f9ba4b3cc267 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 13:52:41 +0100 Subject: [PATCH 7/7] fix(snowflake): tell missing schemas from missing databases with SHOW SCHEMAS - SHOW TABLES/VIEWS report the same generic 2043 error for a missing schema and a missing database, so any missing-object error from a SHOW ... IN SCHEMA read is disambiguated once per plan with SHOW SCHEMAS LIKE ... IN DATABASE: a missing schema lists empty (a fresh target plans as a first run), a missing or unusable database fails naming it, and an existing schema re-raises the original error. - Bounded inspection reads started inside a bounded worker run serially, so capped-schema fallback chunks share the outer concurrency bound. - Semantic source candidates are deduplicated by relation identity, so a relation seen by both SHOW and the name lookup is matched once. Refs CHI-455 --- .../relations/_helpers/inspection_context.py | 13 ++ .../relations/classes/inspection_catalog.py | 16 ++- .../relations/main/run_bounded_inspections.py | 15 ++- .../snowflake/_helpers/show_metadata.py | 26 ++-- .../snowflake/classes/snowflake_adapter.py | 107 +++++++++++++--- .../_helpers/warehouse/semantic_sources.py | 6 +- .../compiler/planner/main/_test_types.py | 19 +++ .../sqlbuild/compiler/planner/main/helpers.py | 17 +++ .../main/test_snowflake_inspection_budget.py | 65 ++++++++++ .../snowflake/inspection/_test_types.py | 25 +++- .../adapters/snowflake/inspection/helpers.py | 96 +++++++++++--- .../inspection/test_inspection_catalog.py | 119 +++++++++++++++++- .../inspection/test_show_metadata.py | 16 +-- .../planner/_helpers/warehouse/_test_types.py | 2 +- .../warehouse/test_semantic_sources.py | 34 +++-- 15 files changed, 497 insertions(+), 79 deletions(-) diff --git a/src/sqlbuild/adapter/relations/_helpers/inspection_context.py b/src/sqlbuild/adapter/relations/_helpers/inspection_context.py index 007b395852..b05668c885 100644 --- a/src/sqlbuild/adapter/relations/_helpers/inspection_context.py +++ b/src/sqlbuild/adapter/relations/_helpers/inspection_context.py @@ -18,6 +18,7 @@ _ACTIVE_CATALOG: ContextVar[InspectionCatalog | None] = ContextVar( "sqlbuild_inspection_catalog", default=None ) +_INSIDE_WORKER: ContextVar[bool] = ContextVar("sqlbuild_inside_inspection_worker", default=False) _QUERY_SINK: ContextVar[Callable[[InspectionQueryRecord], None] | None] = ContextVar( "sqlbuild_inspection_query_sink", default=None ) @@ -72,3 +73,15 @@ def deactivate_query_sink(token: Token[Callable[[InspectionQueryRecord], None] | """Restore the sink that was active before ``activate_query_sink``.""" _QUERY_SINK.reset(token) + + +def inside_inspection_worker() -> bool: + """Return whether this thread is running one bounded inspection read.""" + + return _INSIDE_WORKER.get() + + +def mark_inspection_worker() -> None: + """Mark this worker context so nested bounded reads run serially.""" + + _ = _INSIDE_WORKER.set(True) diff --git a/src/sqlbuild/adapter/relations/classes/inspection_catalog.py b/src/sqlbuild/adapter/relations/classes/inspection_catalog.py index 0be090518a..8ef8562216 100644 --- a/src/sqlbuild/adapter/relations/classes/inspection_catalog.py +++ b/src/sqlbuild/adapter/relations/classes/inspection_catalog.py @@ -2,9 +2,10 @@ from __future__ import annotations +import threading import time from collections.abc import Callable -from typing import Any +from typing import Any, cast from sqlbuild.adapter.contract.models import ColumnInfo, RelationInfo from sqlbuild.adapter.relations._helpers.inspection_context import emit_inspection_record @@ -38,9 +39,22 @@ def __init__(self, *, adapter: Any, connection: Any) -> None: self._relation_listings: dict[_SchemaKey, SchemaRelationListing] = {} self._column_listings: dict[_SchemaKey, SchemaColumnListing] = {} self._capped_column_schemas: set[_SchemaKey] = set() + self._remembered: dict[tuple[str, ...], object] = {} + self._remember_lock: threading.Lock = threading.Lock() self._relation_requests: dict[_RelationRequestKey, tuple[RelationInfo, ...]] = {} self._relation_columns: dict[_RelationIdentity, tuple[ColumnInfo, ...] | None] = {} + def remember[ValueT](self, *, key: tuple[str, ...], compute: Callable[[], ValueT]) -> ValueT: + """Compute one adapter-owned metadata fact at most once per planning invocation.""" + + with self._remember_lock: + if key in self._remembered: + return cast(ValueT, self._remembered[key]) + value: ValueT = compute() + with self._remember_lock: + self._remembered.setdefault(key, value) + return value + @property def schema_scoped(self) -> bool: """Return whether schema listings come from one case-aware read per schema.""" diff --git a/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py b/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py index bca0052fd4..2617df8230 100644 --- a/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py +++ b/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py @@ -7,6 +7,10 @@ from contextvars import copy_context from typing import cast +from sqlbuild.adapter.relations._helpers.inspection_context import ( + inside_inspection_worker, + mark_inspection_worker, +) from sqlbuild.adapter.relations.types import InspectionCompletion @@ -16,9 +20,9 @@ def run_bounded_inspections[ResultT]( concurrency: int, on_complete: InspectionCompletion[ResultT] | None = None, ) -> list[ResultT]: - """Return results in input order, report completions on this thread, raise lowest failure.""" + """Return ordered results, raising the lowest failure; nested calls run serially.""" - if concurrency <= 1 or len(tasks) <= 1: + if concurrency <= 1 or len(tasks) <= 1 or inside_inspection_worker(): results: list[ResultT] = [] index: int task: Callable[[], ResultT] @@ -47,7 +51,7 @@ def _collect[ResultT]( """Stop starting tasks after a failure, let running ones finish, raise the lowest failure.""" futures: dict[Future[ResultT], int] = { - cast(Future[ResultT], pool.submit(copy_context().run, task)): index + cast(Future[ResultT], pool.submit(copy_context().run, _run_as_worker, task)): index for index, task in enumerate(tasks) } completed: dict[int, ResultT] = {} @@ -71,3 +75,8 @@ def _collect[ResultT]( if failures: raise failures[min(failures)] return completed + + +def _run_as_worker[ResultT](task: Callable[[], ResultT]) -> ResultT: + _ = mark_inspection_worker() + return task() diff --git a/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py index 4c2efefa35..8ed213a0d5 100644 --- a/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py +++ b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py @@ -8,11 +8,12 @@ from sqlbuild.adapter.relations.main.relation_age_timestamp import relation_age_timestamp_utc from sqlbuild.adapter.relations.models import ListedRelation -_MISSING_OBJECT_ERRNO: int = 2003 -_MISSING_OBJECT_MESSAGE: str = "does not exist or not authorized" +_MISSING_OBJECT_ERRNOS: frozenset[int] = frozenset({2003, 2043}) +_MISSING_OBJECT_MESSAGES: tuple[str, ...] = ( + "does not exist or not authorized", + "object does not exist, or operation cannot be performed", +) _TRUE_FLAGS: frozenset[str] = frozenset({"Y", "YES", "TRUE"}) -_MISSING_SCHEMA_MARKER: str = "schema '" -_MISSING_DATABASE_MARKER: str = "database '" _TRANSIENT_KIND: str = "TRANSIENT" _TEMPORARY_KIND: str = "TEMPORARY" @@ -54,22 +55,11 @@ def listed_relation_from_show_view( def is_missing_object_error(error: Exception) -> bool: - """Return whether SHOW failed because the schema is absent or invisible to the role.""" - - return ( - getattr(error, "errno", None) == _MISSING_OBJECT_ERRNO - or _MISSING_OBJECT_MESSAGE in str(error).lower() - ) - - -def is_missing_schema_error(error: Exception) -> bool: - """Return whether SHOW failed for a missing schema; a missing database is a real failure.""" + """Return whether SHOW failed because an object is absent or invisible to the role.""" message: str = str(error).lower() - return ( - is_missing_object_error(error) - and _MISSING_SCHEMA_MARKER in message - and _MISSING_DATABASE_MARKER not in message + return getattr(error, "errno", None) in _MISSING_OBJECT_ERRNOS or any( + marker in message for marker in _MISSING_OBJECT_MESSAGES ) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index ddc631f580..f24ff5da00 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -116,7 +116,6 @@ ) from sqlbuild.adapters.snowflake._helpers.show_metadata import ( is_missing_object_error, - is_missing_schema_error, listed_relation_from_show_table, listed_relation_from_show_view, ) @@ -1817,14 +1816,23 @@ def read_schema_relation_listing( ) -> SchemaRelationListing: """List one schema's tables and views with SHOW, which needs no running warehouse.""" + show_database: str = self._show_database( + connection=connection, database=database, schema=schema + ) scope: str = self._show_schema_scope( connection=connection, database=database, schema=schema ) tables: list[dict[str, object]] = self._fetch_paged_show_rows( - connection=connection, query=f"SHOW TABLES IN SCHEMA {scope}" + connection=connection, + query=f"SHOW TABLES IN SCHEMA {scope}", + show_database=show_database, + schema=schema, ) views: list[dict[str, object]] = self._fetch_paged_show_rows( - connection=connection, query=f"SHOW VIEWS IN SCHEMA {scope}" + connection=connection, + query=f"SHOW VIEWS IN SCHEMA {scope}", + show_database=show_database, + schema=schema, ) entries: dict[str, ListedRelation] = {} row: dict[str, object] @@ -1853,8 +1861,13 @@ def read_schema_column_listing( scope: str = self._show_schema_scope( connection=connection, database=database, schema=schema ) - rows: list[dict[str, object]] = self._fetch_show_rows( - connection=connection, query=f"SHOW COLUMNS IN SCHEMA {scope}" + rows: list[dict[str, object]] = self._fetch_schema_show_rows( + connection=connection, + query=f"SHOW COLUMNS IN SCHEMA {scope}", + show_database=self._show_database( + connection=connection, database=database, schema=schema + ), + schema=schema, ) if len(rows) != _SHOW_RESULT_LIMIT: columns: dict[str, list[ColumnInfo]] = {} @@ -1955,6 +1968,12 @@ def _show_schema_scope( ) -> str: """Qualify SHOW with the database so a missing one is reported, never listed empty.""" + resolved: str = self._show_database(connection=connection, database=database, schema=schema) + return f"{self.render_identifier(resolved)}.{self.render_identifier(schema)}" + + def _show_database( + self, *, connection: _SnowflakeConnection, database: str | None, schema: str + ) -> str: resolved: str | None = ( database if database is not None else self._current_database(connection=connection) ) @@ -1962,7 +1981,7 @@ def _show_schema_scope( raise AdapterUserError( message=f"Snowflake schema {schema} needs a database; the session has none" ) - return f"{self.render_identifier(resolved)}.{self.render_identifier(schema)}" + return resolved @staticmethod def _current_database(*, connection: _SnowflakeConnection) -> str | None: @@ -1980,15 +1999,15 @@ def _current_database(*, connection: _SnowflakeConnection) -> str | None: return current def _fetch_paged_show_rows( - self, *, connection: _SnowflakeConnection, query: str + self, *, connection: _SnowflakeConnection, query: str, show_database: str, schema: str ) -> list[dict[str, object]]: """Page SHOW output past its row cap; SHOW orders rows by name.""" rows: list[dict[str, object]] = [] page_query: str = f"{query} LIMIT {_SHOW_RESULT_LIMIT}" while True: - page: list[dict[str, object]] = self._fetch_show_rows( - connection=connection, query=page_query + page: list[dict[str, object]] = self._fetch_schema_show_rows( + connection=connection, query=page_query, show_database=show_database, schema=schema ) rows.extend(page) if len(page) < _SHOW_RESULT_LIMIT: @@ -2000,17 +2019,12 @@ def _fetch_paged_show_rows( def _fetch_show_rows( *, connection: _SnowflakeConnection, query: str ) -> list[dict[str, object]]: - """Run one SHOW command; a missing or unauthorized schema lists nothing.""" + """Run one SHOW command and return its rows keyed by lowercase output column.""" def fetch() -> list[dict[str, object]]: cursor: Any = connection.cursor() try: - try: - cursor.execute(query) - except Exception as error: - if is_missing_schema_error(error): - return [] - raise + cursor.execute(query) names: list[str] = [str(column[0]).lower() for column in cursor.description or ()] return [dict(zip(names, row, strict=False)) for row in cursor.fetchall()] finally: @@ -2018,6 +2032,67 @@ def fetch() -> list[dict[str, object]]: return run_recorded_inspection_query(sql=query, run=fetch) + def _fetch_schema_show_rows( + self, *, connection: _SnowflakeConnection, query: str, show_database: str, schema: str + ) -> list[dict[str, object]]: + """Run SHOW ... IN SCHEMA; a schema that does not exist lists nothing, as before.""" + + try: + return self._fetch_show_rows(connection=connection, query=query) + except Exception as error: + if not is_missing_object_error(error) or self._schema_exists( + connection=connection, show_database=show_database, schema=schema + ): + raise + return [] + + def _schema_exists( + self, *, connection: _SnowflakeConnection, show_database: str, schema: str + ) -> bool: + """Tell a missing schema from a missing database or another failure, once per plan.""" + + def read() -> bool: + return self._read_schema_exists( + connection=connection, show_database=show_database, schema=schema + ) + + catalog: InspectionCatalog | None = active_inspection_catalog( + adapter=self, connection=connection + ) + if catalog is None: + return read() + return catalog.remember( + key=( + "snowflake_schema_exists", + self._information_schema_identifier(show_database), + self._information_schema_identifier(schema), + ), + compute=read, + ) + + def _read_schema_exists( + self, *, connection: _SnowflakeConnection, show_database: str, schema: str + ) -> bool: + stored_schema: str = self._information_schema_identifier(schema) + query: str = ( + f"SHOW SCHEMAS LIKE '{stored_schema.replace(chr(39), chr(39) * 2)}' " + f"IN DATABASE {self.render_identifier(show_database)}" + ) + try: + rows: list[dict[str, object]] = self._fetch_show_rows( + connection=connection, query=query + ) + except Exception as error: + if not is_missing_object_error(error): + raise + raise AdapterUserError( + message=( + f"Snowflake database {show_database} does not exist or the current role " + f"cannot use it: {error}" + ) + ) from error + return any(str(row.get("name")) == stored_schema for row in rows) + def read_relation_columns( self, *, connection: _SnowflakeConnection, relation: RelationInfo ) -> tuple[ColumnInfo, ...] | None: diff --git a/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py b/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py index 351b91bded..5ab3346d96 100644 --- a/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py +++ b/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py @@ -122,7 +122,11 @@ def _list_source_candidates( names=tuple(sorted({entry.table or entry.name for entry in pending})), ) ) - return tuple(dict.fromkeys(relations)) + unique: dict[tuple[str | None, str | None, str], RelationInfo] = {} + relation: RelationInfo + for relation in relations: + _ = unique.setdefault(relation.identity, relation) + return tuple(unique.values()) def _has_candidate(*, entry: SourceEntry, relations: list[RelationInfo]) -> bool: diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py b/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py index 0d0dec6efc..14c1073dfd 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/_test_types.py @@ -141,3 +141,22 @@ class SnowflakeManySchemasTestCase: models_per_schema: int expected_per_relation_column_reads: int expected_schema_column_reads: int + + +@dataclass(frozen=True) +class SnowflakeFreshTargetTestCase: + """Planning before the first build, when target schemas or the database do not exist.""" + + description: str + warehouse_schemas: frozenset[str] + expected_reason: str + expected_schema_checks: int + + +@dataclass(frozen=True) +class SnowflakeMissingDatabaseTestCase: + """Planning against a database that does not exist or the role cannot use.""" + + description: str + warehouse_database: str + expected_error_fragment: str diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py b/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py index c26340c0df..87c094dd49 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/helpers.py @@ -62,6 +62,7 @@ SourceCursorInputPlanErrorTestCase, ) from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + FakeRelation, OfflineSnowflakeAdapter, RecordedQuery, RecordingSnowflakeWarehouse, @@ -534,3 +535,19 @@ def offline_cursor_bound_relations(warehouse: RecordingSnowflakeWarehouse) -> tu _CURSOR_BOUND_RELATION.findall(query.sql)[0] for query in warehouse.queries_of_kind("cursor_bounds") ) + + +def relations_in_schemas( + *, project: SyntheticSnowflakeProject, schemas: frozenset[str] +) -> tuple[FakeRelation, ...]: + """Keep only the synthetic relations that already exist in the given schemas.""" + + return tuple(filter(lambda relation: relation.schema in schemas, project.relations)) + + +def relations_in_database( + *, project: SyntheticSnowflakeProject, database: str +) -> tuple[FakeRelation, ...]: + """Move every synthetic relation into another database.""" + + return tuple(replace(relation, database=database) for relation in project.relations) diff --git a/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py b/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py index 0aaa5471bc..30a4a9bcd0 100644 --- a/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py +++ b/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py @@ -10,12 +10,16 @@ import pytest +from sqlbuild.adapter.contract.exceptions import AdapterUserError from sqlbuild.adapter.relations.constants import INSPECTION_IN_LIST_LIMIT from sqlbuild.compiler.pipeline.models import CompilePipelineResult +from sqlbuild.compiler.planner.types import PlanReason from tests.integration.src.sqlbuild.compiler.planner.main._test_types import ( SnowflakeCursorBoundsBudgetTestCase, + SnowflakeFreshTargetTestCase, SnowflakeInspectionBudgetTestCase, SnowflakeManySchemasTestCase, + SnowflakeMissingDatabaseTestCase, SnowflakeReplanTestCase, ) from tests.integration.src.sqlbuild.compiler.planner.main.helpers import ( @@ -24,6 +28,8 @@ offline_metadata_queries, offline_metadata_reads_by_schema, plan_offline_snowflake_project, + relations_in_database, + relations_in_schemas, ) from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( OfflineSnowflakeAdapter, @@ -209,5 +215,64 @@ def test_given_many_small_schemas_when_planning_then_columns_are_read_per_schema assert warehouse.queries_of_kind("columns") == () +@pytest.mark.parametrize( + "test_case", + [ + SnowflakeFreshTargetTestCase( + description="target schemas created by the first build do not exist yet", + warehouse_schemas=frozenset({"RAW"}), + expected_reason=PlanReason.FIRST_RUN, + expected_schema_checks=3, + ) + ], + ids=lambda case: case.description, +) +def test_given_fresh_target_when_planning_then_every_model_is_a_first_run( + test_case: SnowflakeFreshTargetTestCase, tmp_path: Path +) -> None: + project: SyntheticSnowflakeProject = write_synthetic_snowflake_project( + project_dir=tmp_path / "orders_platform", unmanaged_relations_per_schema=5 + ) + warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( + relations=relations_in_schemas(project=project, schemas=test_case.warehouse_schemas) + ) + + result: CompilePipelineResult = plan_offline_snowflake_project( + project=project, warehouse=warehouse + ) + + assert {entry.reason for entry in result.plan_output.model_entries} == { + test_case.expected_reason + } + assert len(result.plan_output.model_entries) == len(project.model_schemas) + assert len(warehouse.queries_of_kind("show_schemas")) == test_case.expected_schema_checks + assert warehouse.queries_of_kind("columns") == () + + +@pytest.mark.parametrize( + "test_case", + [ + SnowflakeMissingDatabaseTestCase( + description="target database absent from the account", + warehouse_database="ARCHIVE", + expected_error_fragment="database analytics does not exist", + ) + ], + ids=lambda case: case.description, +) +def test_given_missing_database_when_planning_then_plan_fails_naming_it( + test_case: SnowflakeMissingDatabaseTestCase, tmp_path: Path +) -> None: + project: SyntheticSnowflakeProject = write_synthetic_snowflake_project( + project_dir=tmp_path / "orders_platform", unmanaged_relations_per_schema=5 + ) + warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( + relations=relations_in_database(project=project, database=test_case.warehouse_database) + ) + + with pytest.raises(AdapterUserError, match=test_case.expected_error_fragment): + _ = plan_offline_snowflake_project(project=project, warehouse=warehouse) + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py index e103fe167a..60318f9ea7 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py @@ -102,7 +102,6 @@ class MissingObjectErrorTestCase: errno: int | None message: str expected_missing: bool - expected_missing_schema: bool @dataclass(frozen=True) @@ -146,3 +145,27 @@ class CappedColumnFallbackTestCase: request_batches: tuple[int, ...] expected_query_kinds: tuple[str, ...] expected_in_list_sizes: tuple[int, ...] + + +@dataclass(frozen=True) +class MissingSchemaTestCase: + """SHOW reads of a schema the warehouse does not have, or hides from the role.""" + + description: str + schema: str + forbidden_schemas: frozenset[str] + expected_error_fragment: str + expected_schema_checks: int + + +@dataclass(frozen=True) +class NestedConcurrencyTestCase: + """Several cap-sized schemas whose fallback reads would nest inside schema workers.""" + + description: str + schema_count: int + relations_per_schema: int + columns_per_relation: int + statement_latency_seconds: float + expected_max_concurrent: int + expected_column_reads: int diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py index ab0fc0f70d..b4fc0d1b2d 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Iterator from dataclasses import dataclass, field, replace from datetime import UTC, datetime -from itertools import compress +from itertools import chain, compress from pathlib import Path from typing import Any @@ -49,6 +49,9 @@ ) _ANY_PATTERN: re.Pattern[str] = re.compile(r"") _CURRENT_DATABASE_PATTERN: re.Pattern[str] = re.compile(r"^SELECT CURRENT_DATABASE\(\)$") +_SHOW_SCHEMAS_PATTERN: re.Pattern[str] = re.compile( + r"^SHOW SCHEMAS LIKE '(?P(?:[^']|'')*)' IN DATABASE \"(?P[^\"]*)\"$" +) _IN_LIST_PATTERN: re.Pattern[str] = re.compile(r"\bIN\s*\(([^()]*)\)", re.IGNORECASE) _TABLES_PATTERN: re.Pattern[str] = re.compile(r"information_schema\.tables\b", re.IGNORECASE) _COLUMNS_PATTERN: re.Pattern[str] = re.compile(r"information_schema\.columns\b", re.IGNORECASE) @@ -61,6 +64,7 @@ (_SHOW_SCHEMA_COLUMNS_PATTERN, "show_schema_columns"), (_CURSOR_BOUND_PATTERN, "cursor_bounds"), (_CURRENT_DATABASE_PATTERN, "session"), + (_SHOW_SCHEMAS_PATTERN, "show_schemas"), (_OTHER_INFORMATION_SCHEMA_PATTERN, "other_metadata"), (_SHOW_PATTERN, "other_metadata"), (_ANY_PATTERN, "data"), @@ -183,6 +187,8 @@ class RecordingSnowflakeWarehouse: failing_relations: frozenset[str] = frozenset() failing_metadata_schemas: frozenset[str] = frozenset() current_database: str = "ANALYTICS" + extra_schemas: frozenset[tuple[str, str]] = frozenset() + forbidden_schemas: frozenset[str] = frozenset() cursor_values: tuple[str, str] = ("2026-01-01 00:00:00.000", "2026-01-31 00:00:00.000") queries: list[RecordedQuery] = field(default_factory=list) attempted_sql: list[str] = field(default_factory=list) @@ -289,6 +295,7 @@ def _answer(self, *, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[A (_SHOW_SCHEMA_COLUMNS_PATTERN, self._answer_show_schema_columns), (_CURSOR_BOUND_PATTERN, self._answer_cursor_bounds), (_CURRENT_DATABASE_PATTERN, self._answer_current_database), + (_SHOW_SCHEMAS_PATTERN, self._answer_show_schemas), (_ANY_PATTERN, _answer_status), ) route: tuple[re.Pattern[str], Callable[..., Any]] = next( @@ -344,22 +351,51 @@ def _answer_current_database( del sql, params return [(self.current_database,)], (("CURRENT_DATABASE()",),) - def _schema_relations(self, *, scope: str) -> tuple[FakeRelation, ...]: + def _existing_schemas(self) -> frozenset[tuple[str, str]]: + return self.extra_schemas | {(fake.database, fake.schema) for fake in self.relations} + + def _schema_relations( + self, *, scope: str, reports_missing_database: bool + ) -> tuple[FakeRelation, ...]: + """Mimic Snowflake: SHOW TABLES/VIEWS report 2043 for a missing database or schema.""" + parts: tuple[str, ...] = ( self.current_database, *(part.strip('"') for part in scope.split(".")), )[-2:] _FAILED_SCHEMA_ACTIONS[parts[-1] in self.failing_metadata_schemas](scope) - _MISSING_DATABASE_ACTIONS[parts[0] not in {fake.database for fake in self.relations}]( - parts[0] + existing: frozenset[tuple[str, str]] = self._existing_schemas() + database_missing: bool = parts[0] not in {database for database, _ in existing} + _MISSING_DATABASE_ACTIONS[database_missing and reports_missing_database](parts[0]) + _OBJECT_MISSING_ACTIONS[parts not in existing or parts[-1] in self.forbidden_schemas](scope) + return tuple(filter(lambda fake: (fake.database, fake.schema) == parts, self.relations)) + + def _answer_show_schemas( + self, *, sql: str, params: tuple[object, ...] + ) -> tuple[list[tuple[Any, ...]], Any]: + del params + match: re.Match[str] = next(_SHOW_SCHEMAS_PATTERN.finditer(sql)) + database: str = match.group("database") + existing: frozenset[tuple[str, str]] = self._existing_schemas() + _MISSING_DATABASE_ACTIONS[database not in {name for name, _ in existing}](database) + pattern: re.Pattern[str] = re.compile( + re.escape(match.group("pattern").replace("''", "'")) + .replace("%", ".*") + .replace("_", "."), + re.IGNORECASE, ) - matches: tuple[FakeRelation, ...] = tuple( - filter( - lambda fake: (fake.database, fake.schema)[-len(parts) :] == parts, self.relations + names: list[str] = sorted( + schema + for name, schema in filter( + lambda entry: entry[0] == database and pattern.fullmatch(entry[1]) is not None, + existing, ) ) - _MISSING_SCHEMA_ACTIONS[not matches](scope) - return matches + return [(_FIXED_CREATED_AT, name, database) for name in names], ( + ("created_on",), + ("name",), + ("database_name",), + ) def _answer_show_schema_relations( self, *, sql: str, params: tuple[object, ...] @@ -373,7 +409,7 @@ def _answer_show_schema_relations( lambda fake: ( (fake.table_type in _VIEW_TABLE_TYPES) is wants_views and fake.name > after ), - self._schema_relations(scope=match.group("scope")), + self._schema_relations(scope=match.group("scope"), reports_missing_database=False), ), key=lambda fake: fake.name, ) @@ -392,7 +428,10 @@ def _answer_show_schema_columns( scope: str = next(_SHOW_SCHEMA_COLUMNS_PATTERN.finditer(sql)).group("scope") rows: list[tuple[Any, ...]] = [] fake: FakeRelation - for fake in sorted(self._schema_relations(scope=scope), key=lambda item: item.name): + for fake in sorted( + self._schema_relations(scope=scope, reports_missing_database=True), + key=lambda item: item.name, + ): rows.extend( ( fake.name, @@ -456,9 +495,12 @@ def __init__(self, message: str, *, errno: int) -> None: self.errno: int = errno -def _raise_missing_schema(scope: str) -> None: +def _raise_missing_object(scope: str) -> None: + del scope raise FakeSnowflakeProgrammingError( - f"SQL compilation error: Schema '{scope}' does not exist or not authorized.", errno=2003 + "002043 (02000): SQL compilation error: Object does not exist, or operation cannot be " + "performed.", + errno=2043, ) @@ -468,7 +510,8 @@ def _raise_failed_schema(scope: str) -> None: def _raise_missing_database(database: str) -> None: raise FakeSnowflakeProgrammingError( - f"SQL compilation error: Database '{database}' does not exist or not authorized.", + f"002003 (02000): SQL compilation error: Database '{database}' does not exist or not " + "authorized. Your primary role must have USAGE on the database.", errno=2003, ) @@ -477,9 +520,9 @@ def _raise_missing_database(database: str) -> None: False: _accept_relation, True: _raise_missing_database, } -_MISSING_SCHEMA_ACTIONS: dict[bool, Callable[[str], None]] = { +_OBJECT_MISSING_ACTIONS: dict[bool, Callable[[str], None]] = { False: _accept_relation, - True: _raise_missing_schema, + True: _raise_missing_object, } _FAILED_SCHEMA_ACTIONS: dict[bool, Callable[[str], None]] = { False: _accept_relation, @@ -871,7 +914,7 @@ def _sources_yml(*, database: str, source_schema: str, sources: int) -> str: def build_wide_schema_relations( - *, relation_count: int, columns_per_relation: int + *, relation_count: int, columns_per_relation: int, schema: str = "INVENTORY" ) -> tuple[FakeRelation, ...]: """Return one schema of identical relations sized to exercise SHOW result caps.""" @@ -881,7 +924,7 @@ def build_wide_schema_relations( return tuple( FakeRelation( database="ANALYTICS", - schema="INVENTORY", + schema=schema, name=f"STOCK_{index:05d}", columns=columns, ) @@ -889,6 +932,23 @@ def build_wide_schema_relations( ) +def build_wide_schemas_relations( + *, schema_count: int, relation_count: int, columns_per_relation: int +) -> tuple[FakeRelation, ...]: + """Return several wide schemas named WAREHOUSE_.""" + + return tuple( + chain.from_iterable( + build_wide_schema_relations( + relation_count=relation_count, + columns_per_relation=columns_per_relation, + schema=f"WAREHOUSE_{index}", + ) + for index in range(schema_count) + ) + ) + + def show_relation_row(*, created_on: datetime, **fields: object) -> dict[str, object]: """Return one SHOW TABLES or SHOW VIEWS row for a STAGING.ORDERS relation.""" diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py index 1525e9e605..4ecaa7a34f 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py @@ -17,12 +17,16 @@ from sqlbuild.adapter.relations.classes.inspection_catalog import InspectionCatalog from sqlbuild.adapter.relations.constants import INSPECTION_IN_LIST_LIMIT from sqlbuild.adapter.relations.main.open_inspection_catalog import open_inspection_catalog +from sqlbuild.adapter.relations.models import SchemaColumnListing +from sqlbuild.adapters.snowflake.classes.snowflake_connection import _SnowflakeConnection from tests.unit.src.sqlbuild.adapters.snowflake.inspection._test_types import ( CappedColumnFallbackTestCase, ColumnRequestEquivalenceTestCase, DroppedRelationTestCase, FreshnessErrorTestCase, FreshnessReuseTestCase, + MissingSchemaTestCase, + NestedConcurrencyTestCase, RelationRequestEquivalenceTestCase, RepeatedLookupTestCase, ShowResultCapTestCase, @@ -30,11 +34,13 @@ SpeculativePrefetchTestCase, ) from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + FakeRelation, OfflineSnowflakeAdapter, RecordingSnowflakeWarehouse, build_inspection_catalog_relations, build_offline_snowflake, build_wide_schema_relations, + build_wide_schemas_relations, sorted_relation_reprs, ) @@ -433,9 +439,10 @@ def test_given_no_database_when_listing_schemas_then_show_uses_the_session_datab description="a missing database fails instead of listing nothing", database="restricted_catalog", expected_relation_count=0, - expected_error_fragment="Database 'RESTRICTED_CATALOG' does not exist", + expected_error_fragment="database restricted_catalog does not exist", expected_attempted_sql=( 'SHOW TABLES IN SCHEMA "RESTRICTED_CATALOG"."STAGING" LIMIT 10000', + "SHOW SCHEMAS LIKE 'STAGING' IN DATABASE \"RESTRICTED_CATALOG\"", ), ) ], @@ -449,7 +456,7 @@ def test_given_missing_database_when_listing_schema_then_inspection_fails( ) catalog: InspectionCatalog = InspectionCatalog(adapter=adapter, connection=connection) - with pytest.raises(RuntimeError, match=test_case.expected_error_fragment or ""): + with pytest.raises(AdapterUserError, match=test_case.expected_error_fragment or ""): _ = catalog.list_relations(database=test_case.database, schemas=("staging",)) assert tuple(warehouse.attempted_sql) == test_case.expected_attempted_sql @@ -547,5 +554,113 @@ def test_given_cap_sized_column_listing_when_reading_then_only_needed_relations_ ) +@pytest.mark.parametrize( + "test_case", + [ + MissingSchemaTestCase( + description="a schema not created yet lists empty after one existence check", + schema="fresh_target", + forbidden_schemas=frozenset(), + expected_error_fragment="", + expected_schema_checks=1, + ) + ], + ids=lambda case: case.description, +) +def test_given_missing_schema_when_listing_then_relations_and_columns_are_empty( + test_case: MissingSchemaTestCase, +) -> None: + adapter, connection, warehouse = build_offline_snowflake( + relations=build_inspection_catalog_relations() + ) + + with open_inspection_catalog(adapter=adapter, connection=connection) as catalog: + relations: tuple[RelationInfo, ...] = catalog.list_relations( + database="analytics", schemas=(test_case.schema,) + ) + listing: SchemaColumnListing = adapter.read_schema_column_listing( + connection=connection, + database="analytics", + schema=test_case.schema, + stored_names=frozenset(), + ) + + assert relations == () + assert listing.columns_by_stored_name == {} + assert len(warehouse.queries_of_kind("show_schemas")) == test_case.expected_schema_checks + + +@pytest.mark.parametrize( + "test_case", + [ + MissingSchemaTestCase( + description="an existing schema the role cannot list re-raises the SHOW error", + schema="staging", + forbidden_schemas=frozenset({"STAGING"}), + expected_error_fragment="Object does not exist, or operation cannot be performed", + expected_schema_checks=1, + ) + ], + ids=lambda case: case.description, +) +def test_given_existing_schema_when_show_reports_missing_then_original_error_is_raised( + test_case: MissingSchemaTestCase, +) -> None: + warehouse: RecordingSnowflakeWarehouse = RecordingSnowflakeWarehouse( + relations=build_inspection_catalog_relations(), + forbidden_schemas=test_case.forbidden_schemas, + ) + adapter: OfflineSnowflakeAdapter = OfflineSnowflakeAdapter(warehouse=warehouse) + connection: _SnowflakeConnection = adapter.connect({}) + + with ( + open_inspection_catalog(adapter=adapter, connection=connection) as catalog, + pytest.raises(RuntimeError, match=test_case.expected_error_fragment), + ): + _ = catalog.list_relations(database="analytics", schemas=(test_case.schema,)) + + assert len(warehouse.queries_of_kind("show_schemas")) == test_case.expected_schema_checks + + +@pytest.mark.parametrize( + "test_case", + [ + NestedConcurrencyTestCase( + description="fallback chunks inside schema workers share the outer bound", + schema_count=4, + relations_per_schema=125, + columns_per_relation=80, + statement_latency_seconds=0.02, + expected_max_concurrent=4, + expected_column_reads=12, + ) + ], + ids=lambda case: case.description, +) +def test_given_capped_schemas_when_reading_columns_then_nested_reads_do_not_multiply_bound( + test_case: NestedConcurrencyTestCase, +) -> None: + relations: tuple[FakeRelation, ...] = build_wide_schemas_relations( + schema_count=test_case.schema_count, + relation_count=test_case.relations_per_schema, + columns_per_relation=test_case.columns_per_relation, + ) + adapter, connection, warehouse = build_offline_snowflake( + relations=relations, statement_latency_seconds=test_case.statement_latency_seconds + ) + listed: tuple[RelationInfo, ...] = adapter.list_relations( + connection=connection, + database="analytics", + schemas=tuple(f"warehouse_{index}" for index in range(test_case.schema_count)), + ) + warehouse.reset() + + with open_inspection_catalog(adapter=adapter, connection=connection) as catalog: + _ = catalog.get_columns_for_relations(relations=listed) + + assert warehouse.max_concurrent_metadata <= test_case.expected_max_concurrent + assert len(warehouse.queries_of_kind("columns")) == test_case.expected_column_reads + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py index cfc2cc5dbb..2e9f051b38 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py @@ -10,7 +10,6 @@ from sqlbuild.adapter.relations.models import ListedRelation from sqlbuild.adapters.snowflake._helpers.show_metadata import ( is_missing_object_error, - is_missing_schema_error, listed_relation_from_show_table, listed_relation_from_show_view, ) @@ -165,21 +164,18 @@ def test_given_show_row_when_mapping_then_matches_information_schema_relation( errno=2003, message="SQL compilation error", expected_missing=True, - expected_missing_schema=False, ), MissingObjectErrorTestCase( description="missing schema message", errno=None, message="Schema 'ANALYTICS.STAGING' does not exist or not authorized.", expected_missing=True, - expected_missing_schema=True, ), MissingObjectErrorTestCase( description="missing database fails the plan", errno=2003, message="SQL compilation error: Database 'ANALYTICS' does not exist or not authorized.", expected_missing=True, - expected_missing_schema=False, ), MissingObjectErrorTestCase( description="missing table", @@ -187,19 +183,24 @@ def test_given_show_row_when_mapping_then_matches_information_schema_relation( message="SQL compilation error: Table 'ANALYTICS.RAW.ORDERS' does not exist or not " "authorized.", expected_missing=True, - expected_missing_schema=False, + ), + MissingObjectErrorTestCase( + description="SHOW TABLES for a missing schema or database", + errno=2043, + message="002043 (02000): SQL compilation error: Object does not exist, or operation " + "cannot be performed.", + expected_missing=True, ), MissingObjectErrorTestCase( description="other failure", errno=390, message="Metadata service unavailable", expected_missing=False, - expected_missing_schema=False, ), ], ids=lambda case: case.description, ) -def test_given_show_failure_when_classifying_then_only_missing_schemas_list_nothing( +def test_given_show_failure_when_classifying_then_missing_objects_are_recognised( test_case: MissingObjectErrorTestCase, ) -> None: error: Exception = ( @@ -208,7 +209,6 @@ def test_given_show_failure_when_classifying_then_only_missing_schemas_list_noth )[test_case.errno is not None] assert is_missing_object_error(error) is test_case.expected_missing - assert is_missing_schema_error(error) is test_case.expected_missing_schema if __name__ == "__main__": diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py index fbb16d62de..31dc7138a7 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/_test_types.py @@ -38,6 +38,6 @@ class SemanticSourceSchemaCaseCase: @dataclass(frozen=True) class SnowflakeSourceCandidateCase: description: str - declared_schema: str + declared_schemas: tuple[str, ...] expected_candidates: tuple[tuple[str | None, str | None, str], ...] expected_query_kinds: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_semantic_sources.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_semantic_sources.py index 6329dc1072..ffb78f706f 100644 --- a/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_semantic_sources.py +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/test_semantic_sources.py @@ -145,18 +145,32 @@ def test_given_source_schema_case_differs_when_inspecting_then_columns_are_found [ SnowflakeSourceCandidateCase( description="unquoted schema is answered from its SHOW listing", - declared_schema="raw", + declared_schemas=("raw",), expected_candidates=(("analytics", "raw", "orders"),), expected_query_kinds=("show_tables", "show_views"), ), SnowflakeSourceCandidateCase( description="quoted lowercase schema falls back to a database-wide name lookup", - declared_schema="landing", + declared_schemas=("landing",), expected_candidates=( ("analytics", "raw", "orders"), ("analytics", "landing", "orders"), ), - expected_query_kinds=("tables",), + expected_query_kinds=("show_schemas", "tables"), + ), + SnowflakeSourceCandidateCase( + description="a relation found by SHOW and by the name lookup is listed once", + declared_schemas=("raw", "landing"), + expected_candidates=( + ("analytics", "raw", "orders"), + ("analytics", "landing", "orders"), + ), + expected_query_kinds=( + "show_tables", + "show_views", + "show_schemas", + "tables", + ), ), ], ids=lambda case: case.description, @@ -176,18 +190,18 @@ def test_given_snowflake_source_schema_when_listing_candidates_then_finds_relati adapter=adapter, connection=connection, database="analytics", - entries=( + entries=tuple( SourceEntry( - name="orders", - database="analytics", - schema=test_case.declared_schema, - table="orders", - ), + name=f"orders_{schema}", database="analytics", schema=schema, table="orders" + ) + for schema in test_case.declared_schemas ), ) assert tuple(relation.identity for relation in candidates) == test_case.expected_candidates - assert tuple(query.kind for query in warehouse.queries) == test_case.expected_query_kinds + assert sorted(query.kind for query in warehouse.queries) == sorted( + test_case.expected_query_kinds + ) if __name__ == "__main__":