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..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 @@ -41,6 +41,19 @@ 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 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. + ## Shared connections across targets Use multiple targets to reuse one Snowflake connection while selecting different authoritative 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..b05668c885 --- /dev/null +++ b/src/sqlbuild/adapter/relations/_helpers/inspection_context.py @@ -0,0 +1,87 @@ +"""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 +) +_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 +) + + +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) + + +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 new file mode 100644 index 0000000000..8ef8562216 --- /dev/null +++ b/src/sqlbuild/adapter/relations/classes/inspection_catalog.py @@ -0,0 +1,384 @@ +"""Invocation-scoped warehouse metadata shared by every planning inspection phase.""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Callable +from typing import Any, cast + +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._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.""" + + return self._reader is not None + + @property + def concurrency(self) -> int: + """Return the bounded number of concurrent reads this adapter supports.""" + + return self._concurrency + + 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: + """Use exact per-relation reads for small requests and one read per schema otherwise.""" + + unscoped: list[RelationInfo] = [] + remaining: list[RelationInfo] = [] + relation: RelationInfo + for relation in relations: + if relation.schema is None: + unscoped.append(relation) + continue + 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)) + 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 "", + stored_names=self._stored_names(reader=reader, relations=scoped), + known_capped=schema_key in self._capped_column_schemas, + ) + for schema_key, scoped in by_schema.items() + ), + concurrency=self._concurrency, + ) + 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( + reader.metadata_name_key(relation.name) + ) + ) + + @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, + 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, + stored_names=stored_names, + known_capped=known_capped, + ) + + return read_listing + + def _bind_relation_columns( + self, *, reader: SchemaScopedMetadataReader, relation: RelationInfo + ) -> Callable[[], tuple[ColumnInfo, ...] | None]: + def read_columns() -> tuple[ColumnInfo, ...] | None: + 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..2617df8230 --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/run_bounded_inspections.py @@ -0,0 +1,82 @@ +"""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._helpers.inspection_context import ( + inside_inspection_worker, + mark_inspection_worker, +) +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 ordered results, raising the lowest failure; nested calls run serially.""" + + if concurrency <= 1 or len(tasks) <= 1 or inside_inspection_worker(): + 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 + 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, _run_as_worker, task)): index + for index, task in enumerate(tasks) + } + completed: dict[int, ResultT] = {} + failures: dict[int, BaseException] = {} + 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() + 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 + + +def _run_as_worker[ResultT](task: Callable[[], ResultT]) -> ResultT: + _ = mark_inspection_worker() + return task() 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..6c7fba5721 --- /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, ...] + + +@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, ...]] + complete: bool = True + + +@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..abce8d457f 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,46 @@ 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, + stored_names: frozenset[str], + known_capped: bool = False, + ) -> SchemaColumnListing: + """Read one schema's columns; an incomplete listing covers only ``stored_names``.""" + ... + + def read_relation_columns( + self, *, connection: Any, relation: RelationInfo + ) -> tuple[ColumnInfo, ...] | None: + """Read one relation's columns, or None when it no longer exists.""" + ... + + 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/_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..8ed213a0d5 --- /dev/null +++ b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py @@ -0,0 +1,98 @@ +"""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_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"}) +_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=None, + retention_days=None, + ) + + +def is_missing_object_error(error: Exception) -> bool: + """Return whether SHOW failed because an object is absent or invisible to the role.""" + + message: str = str(error).lower() + return getattr(error, "errno", None) in _MISSING_OBJECT_ERRNOS or any( + marker in message for marker in _MISSING_OBJECT_MESSAGES + ) + + +def _listed_relation( + *, + row: Mapping[str, object], + database: str | None, + relation_type: str, + is_transient: bool | None, + 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 7d1f9964d0..f24ff5da00 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -5,11 +5,12 @@ import csv import json import logging +from collections.abc import Callable from dataclasses import replace 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, @@ -80,7 +81,20 @@ StatementSizeLimit, 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, +) +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, ) @@ -96,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, ) @@ -124,6 +145,15 @@ from sqlbuild.sql_values.models import SqlValue _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"}) +_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" +) _BULK_COLUMN_RELATION_CHUNK_SIZE: int = 2_000 @@ -133,6 +163,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 +389,70 @@ 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, ...]]: + """Read LAST_ALTERED for only the requested tables, in bounded IN-list chunks.""" + + names: tuple[str, ...] = tuple( + dict.fromkeys( + self._information_schema_identifier(request.name) for request in scoped_requests + ) + ) + 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: + scope_clauses.append("table_catalog = %s") + scope_params.append(database) + if schema is not None: + 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) + ) + 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), + ) + 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( *, database: str | None, schema: str | None, name: str @@ -703,34 +799,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 +1775,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 +1796,334 @@ 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 + @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 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}", + 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}", + show_database=show_database, + schema=schema, + ) + 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()) + ) + + def read_schema_column_listing( + 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; a cap-sized result rereads only names.""" + + if not known_capped: + scope: str = self._show_schema_scope( + connection=connection, database=database, schema=schema + ) + 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 ), - 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]), + schema=schema, ) - for row in rows + 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, + 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 " + + 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)) + 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 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, *, connection: _SnowflakeConnection, database: str | None, schema: str + ) -> 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) + ) + if resolved is None: + raise AdapterUserError( + message=f"Snowflake schema {schema} needs a database; the session has none" + ) + return resolved + + @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, 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_schema_show_rows( + connection=connection, query=page_query, show_database=show_database, schema=schema + ) + 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 and return its rows keyed by lowercase output column.""" + + def fetch() -> list[dict[str, object]]: + cursor: Any = connection.cursor() + try: + 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: + cursor.close() + + 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: + """Read one relation's columns with SHOW COLUMNS; None once it no longer exists.""" + + 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.""" + + 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( @@ -1898,12 +2283,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 ) @@ -1924,16 +2311,21 @@ 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(), - type=self._build_show_columns_type(raw_data_type=row[3]), + type=show_columns_type(row[3]), ) for row in rows ) @@ -2011,29 +2403,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.""" @@ -3059,16 +3428,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/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/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..5ab3346d96 100644 --- a/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py +++ b/src/sqlbuild/compiler/planner/_helpers/warehouse/semantic_sources.py @@ -6,6 +6,12 @@ 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, +) from sqlbuild.compiler.compile.models import CompiledObjectKey, CompiledProject from sqlbuild.compiler.planner.types import ContractPolicy from sqlbuild.compiler.references.types import SqlReferenceKind @@ -44,15 +50,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 +64,75 @@ 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 candidates by name in any schema; schema listings first answer scoped catalogs.""" + + pending: tuple[SourceEntry, ...] = entries + relations: list[RelationInfo] = [] + 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, + 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})), + ) + ) + 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 pending})), + ) + ) + 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: + 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/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/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/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..14c1073dfd 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,62 @@ 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 + expected_freshness_reads: 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_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 + + +@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 ca25110464..87c094dd49 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,13 @@ FormatPlanIntegrationTestCase, SourceCursorInputPlanErrorTestCase, ) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + FakeRelation, + OfflineSnowflakeAdapter, + RecordedQuery, + RecordingSnowflakeWarehouse, + SyntheticSnowflakeProject, +) def build_project_from_test_case( @@ -447,3 +461,93 @@ def grouped(model: type) -> dict[str, Any]: policies=policies, **kwargs, ) + + +_OFFLINE_METADATA_KINDS: frozenset[str] = frozenset( + { + "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) + + +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 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) + ) + ) + + +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, ...]: + """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") + ) + + +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 new file mode 100644 index 0000000000..30a4a9bcd0 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/planner/main/test_snowflake_inspection_budget.py @@ -0,0 +1,278 @@ +"""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.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 ( + offline_cursor_bound_relations, + offline_information_schema_queries, + 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, + RecordedQuery, + RecordingSnowflakeWarehouse, + SyntheticSnowflakeProject, + write_synthetic_snowflake_project, +) + +_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( + "test_case", + [ + SnowflakeInspectionBudgetTestCase( + description="about 2,500 relations", + unmanaged_relations_per_schema=450, + expected_schema_reads=_ONE_READ_PER_SCHEMA, + 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=_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, +) +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=_SOURCES, + 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) + warehouse_queries: tuple[RecordedQuery, ...] = offline_information_schema_queries(warehouse) + assert ( + offline_metadata_reads_by_schema(warehouse=warehouse, kind="show_tables") + == test_case.expected_schema_reads + ) + assert ( + 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 + + +@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_listing_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("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") == () + + +@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/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/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..1315c6f7ac --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/_test_types.py @@ -0,0 +1,49 @@ +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 + + +@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 new file mode 100644 index 0000000000..197739c60c --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/helpers.py @@ -0,0 +1,111 @@ +"""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 + + +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 new file mode 100644 index 0000000000..6c01d68136 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/main/run_bounded_inspections/test_run_bounded_inspections.py @@ -0,0 +1,176 @@ +"""Bounded inspection reads keep input order and fail deterministically.""" + +from __future__ import annotations + +import threading +import time + +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, + 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, +) + + +@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, + ) + + +@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/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..f762fe4391 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/test_metadata_sql_ownership.py @@ -0,0 +1,53 @@ +"""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/_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", + "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/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..60318f9ea7 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py @@ -0,0 +1,171 @@ +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 + + +@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 | None + 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, ...] + + +@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, ...] + + +@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 new file mode 100644 index 0000000000..b4fc0d1b2d --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -0,0 +1,955 @@ +"""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, replace +from datetime import UTC, datetime +from itertools import chain, 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_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*$", + re.IGNORECASE | re.DOTALL, +) +_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) +_QUERY_KINDS: tuple[tuple[re.Pattern[str], str], ...] = ( + (_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"), + (_CURRENT_DATABASE_PATTERN, "session"), + (_SHOW_SCHEMAS_PATTERN, "show_schemas"), + (_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(") +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) +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() + 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) + _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: + _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") + 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), + (_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), + (_SHOW_SCHEMAS_PATTERN, self._answer_show_schemas), + (_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( + 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)) + 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 _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 _answer_current_database( + self, *, sql: str, params: tuple[object, ...] + ) -> tuple[list[tuple[Any, ...]], Any]: + del sql, params + return [(self.current_database,)], (("CURRENT_DATABASE()",),) + + 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) + 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, + ) + names: list[str] = sorted( + schema + for name, schema in filter( + lambda entry: entry[0] == database and pattern.fullmatch(entry[1]) is not None, + existing, + ) + ) + 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, ...] + ) -> 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"), reports_missing_database=False), + ), + 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, reports_missing_database=True), + 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, tuple((f,) for f in _SHOW_SCHEMA_COLUMN_FIELDS) + + 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 _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_object(scope: str) -> None: + del scope + raise FakeSnowflakeProgrammingError( + "002043 (02000): SQL compilation error: Object does not exist, or operation cannot be " + "performed.", + errno=2043, + ) + + +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"002003 (02000): SQL compilation error: Database '{database}' does not exist or not " + "authorized. Your primary role must have USAGE on the database.", + errno=2003, + ) + + +_MISSING_DATABASE_ACTIONS: dict[bool, Callable[[str], None]] = { + False: _accept_relation, + True: _raise_missing_database, +} +_OBJECT_MISSING_ACTIONS: dict[bool, Callable[[str], None]] = { + False: _accept_relation, + True: _raise_missing_object, +} +_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]: + 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 _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], 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], + } + + +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, 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)) + + +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(60) + ), + 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 form of listed relations without SHOW-absent LAST_ALTERED.""" + + return sorted(repr(replace(relation, last_altered_at=None)) 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" + + +def build_wide_schema_relations( + *, relation_count: int, columns_per_relation: int, schema: str = "INVENTORY" +) -> 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=schema, + name=f"STOCK_{index:05d}", + columns=columns, + ) + for index in range(relation_count) + ) + + +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.""" + + 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 new file mode 100644 index 0000000000..4ecaa7a34f --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_inspection_catalog.py @@ -0,0 +1,666 @@ +"""Catalog-served Snowflake metadata must equal what the adapter returns directly.""" + +from __future__ import annotations + +from collections import Counter +from dataclasses import replace + +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 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, + ShowScopeTestCase, + 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, +) + +_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=67, + ), + 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 only the current catalog", + database=None, + schemas=("staging",), + names=("orders",), + expected_relation_count=1, + ), + 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=62, + ), + ], + 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=("show_schema_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 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=3, + ) + ], + ids=lambda case: case.description, +) +def test_given_open_catalog_when_reading_source_freshness_then_reads_requested_tables_only( + 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) + 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( + "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, 67), + expected_query_kinds=("show_tables", "show_views"), + ) + ], + 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="unavailable", + 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("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="exactly 10,000 columns is treated as truncated and reread", + relation_count=40, + 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, +) +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 + + +@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', + "SHOW SCHEMAS LIKE 'STAGING' IN DATABASE \"RESTRICTED_CATALOG\"", + ), + ) + ], + 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(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 + + +@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 + ) + + +@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_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..2e9f051b38 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_show_metadata.py @@ -0,0 +1,215 @@ +"""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=None, + 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=None, + 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=None, + 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 alone", + errno=2003, + message="SQL compilation error", + expected_missing=True, + ), + MissingObjectErrorTestCase( + description="missing schema message", + errno=None, + message="Schema 'ANALYTICS.STAGING' does not exist or not authorized.", + expected_missing=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, + ), + MissingObjectErrorTestCase( + description="missing table", + errno=2003, + message="SQL compilation error: Table 'ANALYTICS.RAW.ORDERS' does not exist or not " + "authorized.", + expected_missing=True, + ), + 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, + ), + ], + ids=lambda case: case.description, +) +def test_given_show_failure_when_classifying_then_missing_objects_are_recognised( + 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/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"]) 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..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 @@ -8,3 +8,36 @@ 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 + + +@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_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/helpers.py b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py new file mode 100644 index 0000000000..a5c2f7b300 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/planner/_helpers/warehouse/helpers.py @@ -0,0 +1,152 @@ +"""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.adapter.contract.models import ColumnInfo +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) + ] + + +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..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 @@ -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,128 @@ 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_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_schemas=("landing",), + expected_candidates=( + ("analytics", "raw", "orders"), + ("analytics", "landing", "orders"), + ), + 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, +) +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=tuple( + SourceEntry( + 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 sorted(query.kind for query in warehouse.queries) == sorted( + test_case.expected_query_kinds + ) + + if __name__ == "__main__": pytest.main([__file__, "-vv"]) 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"]) diff --git a/website/src/content/docs/docs/concepts/adapters/snowflake.mdx b/website/src/content/docs/docs/concepts/adapters/snowflake.mdx index f3a12b355d..ff8a6787ae 100644 --- a/website/src/content/docs/docs/concepts/adapters/snowflake.mdx +++ b/website/src/content/docs/docs/concepts/adapters/snowflake.mdx @@ -38,6 +38,19 @@ 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 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. + ## Shared connections across targets Use multiple targets to reuse one Snowflake connection while selecting different authoritative