From 7ecf596309e377a031e3db85683cb1a498895f30 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 15:00:43 +0100 Subject: [PATCH 1/9] perf(snowflake): read single relations with SHOW instead of INFORMATION_SCHEMA Build-time relation_exists and get_columns ran one INFORMATION_SCHEMA query per lookup, which needs a running warehouse and queues behind model work. They now use SHOW TERSE TABLES/VIEWS LIKE ... IN SCHEMA (exact stored-name match on the result) and SHOW COLUMNS IN TABLE, then IN VIEW, with the SHOW COLUMNS type translation shared with planning. A missing relation or schema still reports absent/no columns; a missing database still fails, disambiguated with SHOW SCHEMAS as in planning. Unqualified-schema lookups keep the INFORMATION_SCHEMA query. Refs CHI-455 --- .../snowflake/_helpers/show_metadata.py | 6 + .../snowflake/classes/snowflake_adapter.py | 77 +++++++ src/sqlbuild/adapters/snowflake/constants.py | 1 + .../snowflake/inspection/_test_types.py | 35 +++ .../adapters/snowflake/inspection/helpers.py | 63 +++++- .../test_single_relation_metadata.py | 211 ++++++++++++++++++ .../adapters/snowflake/test_client.py | 33 --- 7 files changed, 381 insertions(+), 45 deletions(-) create mode 100644 tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_single_relation_metadata.py diff --git a/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py index 8ed213a0d5..feea46c827 100644 --- a/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py +++ b/src/sqlbuild/adapters/snowflake/_helpers/show_metadata.py @@ -63,6 +63,12 @@ def is_missing_object_error(error: Exception) -> bool: ) +def show_like_pattern(stored_name: str) -> str: + """Quote a stored name as a SHOW LIKE literal; ``_`` matches neighbours, so compare exactly.""" + + return stored_name.replace("'", "''") + + def _listed_relation( *, row: Mapping[str, object], diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index f24ff5da00..dbd00df1a0 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -118,6 +118,7 @@ is_missing_object_error, listed_relation_from_show_table, listed_relation_from_show_view, + show_like_pattern, ) from sqlbuild.adapters.snowflake.classes.snowflake_connection import _SnowflakeConnection from sqlbuild.adapters.snowflake.constants import ( @@ -128,6 +129,7 @@ OAUTH_AUTHORIZATION_CODE_AUTHENTICATOR, SECONDARY_ROLES_ALL, SECONDARY_ROLES_NONE, + SHOW_LIKE_ESCAPE_CHARACTER, STATUS_COLUMN_NAME, SUCCESS_STATUS_TOKENS, TRUE_METADATA_VALUE, @@ -1744,6 +1746,40 @@ def relation_exists( database: str | None, schema: str | None, name: str, + ) -> bool: + """Look one relation up with SHOW, which needs no running warehouse.""" + + if schema is None or SHOW_LIKE_ESCAPE_CHARACTER in name: + return self._information_schema_relation_exists( + connection=connection, database=database, schema=schema, name=name + ) + show_database: str = self._show_database( + connection=connection, database=database, schema=schema + ) + scope: str = self._show_schema_scope( + connection=connection, database=database, schema=schema + ) + stored_name: str = self._information_schema_identifier(name) + pattern: str = show_like_pattern(stored_name) + kind: str + for kind in ("TABLES", "VIEWS"): + rows: list[dict[str, object]] = self._fetch_schema_show_rows( + connection=connection, + query=f"SHOW TERSE {kind} LIKE '{pattern}' IN SCHEMA {scope}", + show_database=show_database, + schema=schema, + ) + if any(str(row.get("name")) == stored_name for row in rows): + return True + return False + + def _information_schema_relation_exists( + self, + *, + connection: _SnowflakeConnection, + database: str | None, + schema: str | None, + name: str, ) -> bool: clauses: list[str] = ["table_name = %s"] params: list[str] = [self._information_schema_identifier(name)] @@ -2183,6 +2219,47 @@ def get_columns( database: str | None, schema: str | None, name: str, + ) -> tuple[ColumnInfo, ...]: + """Read one relation's columns with SHOW COLUMNS, which needs no running warehouse.""" + + if schema is None: + return self._information_schema_columns( + connection=connection, database=database, schema=schema, name=name + ) + show_database: str = self._show_database( + connection=connection, database=database, schema=schema + ) + relation: str = ".".join( + self.render_identifier(part) for part in (show_database, schema, name) + ) + kind: str + for kind in ("TABLE", "VIEW"): + try: + rows: list[dict[str, object]] = self._fetch_show_rows( + connection=connection, query=f"SHOW COLUMNS IN {kind} {relation}" + ) + except Exception as error: + if not is_missing_object_error(error): + raise + continue + return tuple( + ColumnInfo( + name=str(row["column_name"]).lower(), + type=show_columns_type(row["data_type"]), + ) + for row in rows + if str(row.get("kind") or _SHOW_COLUMN_KIND).upper() in _SHOW_COLUMN_KINDS + ) + _ = self._schema_exists(connection=connection, show_database=show_database, schema=schema) + return () + + def _information_schema_columns( + self, + *, + connection: _SnowflakeConnection, + database: str | None, + schema: str | None, + name: str, ) -> tuple[ColumnInfo, ...]: query: str = ( "SELECT column_name, data_type, numeric_precision, numeric_scale, " diff --git a/src/sqlbuild/adapters/snowflake/constants.py b/src/sqlbuild/adapters/snowflake/constants.py index 01c94d449e..33c5ca7c56 100644 --- a/src/sqlbuild/adapters/snowflake/constants.py +++ b/src/sqlbuild/adapters/snowflake/constants.py @@ -22,3 +22,4 @@ TEXT_TYPE_NAME: str = "TEXT" UNBOUNDED_TEXT_TYPE_NAMES: frozenset[str] = frozenset({"STRING", "TEXT", "VARCHAR"}) TRUE_METADATA_VALUE: str = "YES" +SHOW_LIKE_ESCAPE_CHARACTER: str = "\\" diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py index 60318f9ea7..ec19a6d018 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py @@ -169,3 +169,38 @@ class NestedConcurrencyTestCase: statement_latency_seconds: float expected_max_concurrent: int expected_column_reads: int + + +@dataclass(frozen=True) +class SingleRelationLookupTestCase: + """One build-path existence and column lookup answered by SHOW.""" + + description: str + database: str | None + schema: str + name: str + expected_exists: bool + expected_column_names: tuple[str, ...] + expected_query_kinds: tuple[str, ...] + + +@dataclass(frozen=True) +class SingleRelationErrorTestCase: + """A build-path lookup whose database does not exist.""" + + description: str + database: str + schema: str + name: str + expected_error_fragment: str + + +@dataclass(frozen=True) +class ShowColumnsEquivalenceTestCase: + """One relation whose SHOW COLUMNS read must equal its INFORMATION_SCHEMA read.""" + + description: str + database: str + schema: str + name: str + expected_column_count: int diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py index b4fc0d1b2d..5fcec5acd7 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -36,7 +36,8 @@ 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"^SHOW (?PTERSE )?(?PTABLES|VIEWS)(?: LIKE '(?P(?:[^']|'')*)')? " + r"IN SCHEMA (?P\S+)" r"(?: LIMIT (?P\d+)(?: FROM '(?P(?:[^']|'')*)')?)?$" ) _SHOW_SCHEMA_COLUMNS_PATTERN: re.Pattern[str] = re.compile( @@ -59,6 +60,7 @@ (_TABLES_PATTERN, "tables"), (_COLUMNS_PATTERN, "columns"), (_SHOW_COLUMNS_PATTERN, "show_columns"), + (re.compile(r"^SHOW TERSE (?:TABLES|VIEWS) LIKE "), "show_relation_lookup"), (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"), @@ -108,6 +110,13 @@ "is_secure", "is_materialized", ) +_SHOW_TERSE_FIELDS: tuple[str, ...] = ( + "created_on", + "name", + "kind", + "database_name", + "schema_name", +) _SHOW_SCHEMA_COLUMN_FIELDS: tuple[str, ...] = ( "table_name", "schema_name", @@ -378,12 +387,7 @@ def _answer_show_schemas( 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, - ) + pattern: re.Pattern[str] = _like_pattern(match.group("pattern")) names: list[str] = sorted( schema for name, schema in filter( @@ -404,21 +408,28 @@ def _answer_show_schema_relations( 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("''", "'") + like: re.Pattern[str] = _like_pattern(match.group("like") or "%") selected: list[FakeRelation] = sorted( filter( lambda fake: ( - (fake.table_type in _VIEW_TABLE_TYPES) is wants_views and fake.name > after + (fake.table_type in _VIEW_TABLE_TYPES) is wants_views + and fake.name > after + and like.fullmatch(fake.name) is not None ), 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) + terse: bool = match.group("terse") is not None 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] + (_show_table_row, _show_view_row)[wants_views], + _show_terse_row, + )[terse] + fields: tuple[str, ...] = ( + (_SHOW_TABLE_FIELDS, _SHOW_VIEW_FIELDS)[wants_views], + _SHOW_TERSE_FIELDS, + )[terse] return [row_builder(fake) for fake in selected[:limit]], tuple((f,) for f in fields) def _answer_show_schema_columns( @@ -463,6 +474,13 @@ def _answer_cursor_bounds( return [values], tuple((f"_{index}",) for index in range(len(values))) +def _like_pattern(pattern: str) -> re.Pattern[str]: + return re.compile( + re.escape(pattern.replace("''", "'")).replace("%", ".*").replace("_", "."), + re.IGNORECASE, + ) + + def _accept_relation(relation: str) -> None: del relation @@ -550,6 +568,16 @@ def _show_table_row(relation: FakeRelation) -> tuple[Any, ...]: ) +def _show_terse_row(relation: FakeRelation) -> tuple[Any, ...]: + return ( + _FIXED_CREATED_AT, + relation.name, + relation.table_type, + relation.database, + relation.schema, + ) + + def _show_view_row(relation: FakeRelation) -> tuple[Any, ...]: return ( _FIXED_CREATED_AT, @@ -953,3 +981,14 @@ def show_relation_row(*, created_on: datetime, **fields: object) -> dict[str, ob """Return one SHOW TABLES or SHOW VIEWS row for a STAGING.ORDERS relation.""" return {"created_on": created_on, "name": "ORDERS", "schema_name": "STAGING", **fields} + + +def attempted_query_kinds(warehouse: RecordingSnowflakeWarehouse) -> tuple[str, ...]: + """Classify every statement sent to the warehouse, including ones that failed.""" + + return tuple( + RecordedQuery( + sql=sql, params=(), thread_id=0, started_at=0.0, finished_at=0.0, row_count=0 + ).kind + for sql in warehouse.attempted_sql + ) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_single_relation_metadata.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_single_relation_metadata.py new file mode 100644 index 0000000000..d2119d36f9 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_single_relation_metadata.py @@ -0,0 +1,211 @@ +"""Build-path single-relation existence and column reads use SHOW, never INFORMATION_SCHEMA.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.adapter.contract.exceptions import AdapterUserError +from sqlbuild.adapter.contract.models import ColumnInfo +from sqlbuild.adapters.snowflake.classes.snowflake_connection import _SnowflakeConnection +from tests.unit.src.sqlbuild.adapters.snowflake.inspection._test_types import ( + ShowColumnsEquivalenceTestCase, + SingleRelationErrorTestCase, + SingleRelationLookupTestCase, +) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + FakeColumn, + FakeRelation, + OfflineSnowflakeAdapter, + RecordingSnowflakeWarehouse, + attempted_query_kinds, + build_offline_snowflake, +) + +_TYPED_COLUMNS: tuple[FakeColumn, ...] = ( + FakeColumn(name="ORDER_ID", data_type="NUMBER", numeric_precision=38, numeric_scale=0), + FakeColumn(name="AMOUNT", data_type="NUMBER", numeric_precision=10, numeric_scale=2), + FakeColumn(name="STATUS", data_type="TEXT", character_maximum_length=40), + FakeColumn(name="RATE", data_type="FLOAT", character_maximum_length=None), + FakeColumn(name="ORDERED_AT", data_type="TIMESTAMP_NTZ", character_maximum_length=None), + FakeColumn(name="PAYLOAD", data_type="VARIANT", character_maximum_length=None), +) +_RELATIONS: tuple[FakeRelation, ...] = ( + FakeRelation(database="ANALYTICS", schema="STAGING", name="ORDERS", columns=_TYPED_COLUMNS), + FakeRelation( + database="ANALYTICS", + schema="STAGING", + name="ORDERS_V", + table_type="VIEW", + columns=_TYPED_COLUMNS[:2], + ), + FakeRelation(database="ANALYTICS", schema="STAGING", name="ORDERSXDELTA"), +) + + +@pytest.mark.parametrize( + "test_case", + [ + SingleRelationLookupTestCase( + description="table answered by SHOW TABLES and SHOW COLUMNS IN TABLE", + database="analytics", + schema="staging", + name="orders", + expected_exists=True, + expected_column_names=("order_id", "amount", "status", "rate", "ordered_at", "payload"), + expected_query_kinds=("show_relation_lookup", "show_columns"), + ), + SingleRelationLookupTestCase( + description="view falls through to SHOW VIEWS", + database="analytics", + schema="staging", + name="orders_v", + expected_exists=True, + expected_column_names=("order_id", "amount"), + expected_query_kinds=("show_relation_lookup", "show_relation_lookup", "show_columns"), + ), + SingleRelationLookupTestCase( + description="session database qualifies an unqualified lookup", + database=None, + schema="staging", + name="orders", + expected_exists=True, + expected_column_names=("order_id", "amount", "status", "rate", "ordered_at", "payload"), + expected_query_kinds=("session", "show_relation_lookup", "show_columns"), + ), + SingleRelationLookupTestCase( + description="LIKE wildcard neighbour is not an exact match", + database="analytics", + schema="staging", + name="orders_delta", + expected_exists=False, + expected_column_names=(), + expected_query_kinds=( + "show_relation_lookup", + "show_relation_lookup", + "show_columns", + "show_columns", + "show_schemas", + ), + ), + SingleRelationLookupTestCase( + description="missing schema behaves as an absent relation", + database="analytics", + schema="fresh_target", + name="orders", + expected_exists=False, + expected_column_names=(), + expected_query_kinds=( + "show_relation_lookup", + "show_schemas", + "show_relation_lookup", + "show_schemas", + "show_columns", + "show_columns", + "show_schemas", + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_relation_when_looking_up_one_then_show_answers_without_information_schema( + test_case: SingleRelationLookupTestCase, +) -> None: + adapter: OfflineSnowflakeAdapter + connection: _SnowflakeConnection + warehouse: RecordingSnowflakeWarehouse + adapter, connection, warehouse = build_offline_snowflake(relations=_RELATIONS) + + exists: bool = adapter.relation_exists( + connection=connection, + database=test_case.database, + schema=test_case.schema, + name=test_case.name, + ) + columns: tuple[ColumnInfo, ...] = adapter.get_columns( + connection=connection, + database=test_case.database, + schema=test_case.schema, + name=test_case.name, + ) + + assert exists is test_case.expected_exists + assert tuple(column.name for column in columns) == test_case.expected_column_names + assert attempted_query_kinds(warehouse) == test_case.expected_query_kinds + + +@pytest.mark.parametrize( + "test_case", + [ + ShowColumnsEquivalenceTestCase( + description=relation.name, + database=relation.database.lower(), + schema=relation.schema.lower(), + name=relation.name.lower(), + expected_column_count=len(relation.columns), + ) + for relation in _RELATIONS + ], + ids=lambda case: case.description, +) +def test_given_relation_when_reading_columns_then_show_matches_information_schema( + test_case: ShowColumnsEquivalenceTestCase, +) -> None: + adapter: OfflineSnowflakeAdapter + connection: _SnowflakeConnection + adapter, connection, _ = build_offline_snowflake(relations=_RELATIONS) + + shown: tuple[ColumnInfo, ...] = adapter.get_columns( + connection=connection, + database=test_case.database, + schema=test_case.schema, + name=test_case.name, + ) + legacy: tuple[ColumnInfo, ...] = adapter._information_schema_columns( + connection=connection, + database=test_case.database, + schema=test_case.schema, + name=test_case.name, + ) + + assert len(shown) == test_case.expected_column_count + assert shown == legacy + + +@pytest.mark.parametrize( + "test_case", + [ + SingleRelationErrorTestCase( + description="missing database fails instead of reporting an absent relation", + database="retired", + schema="staging", + name="orders", + expected_error_fragment="database retired does not exist", + ) + ], + ids=lambda case: case.description, +) +def test_given_missing_database_when_looking_up_one_relation_then_raises_user_error( + test_case: SingleRelationErrorTestCase, +) -> None: + adapter: OfflineSnowflakeAdapter + connection: _SnowflakeConnection + adapter, connection, _ = build_offline_snowflake(relations=_RELATIONS) + + with pytest.raises(AdapterUserError, match=test_case.expected_error_fragment): + _ = adapter.relation_exists( + connection=connection, + database=test_case.database, + schema=test_case.schema, + name=test_case.name, + ) + with pytest.raises(AdapterUserError, match=test_case.expected_error_fragment): + _ = adapter.get_columns( + connection=connection, + database=test_case.database, + schema=test_case.schema, + name=test_case.name, + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py b/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py index d4b7cd9a7e..43e7e24ec8 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py @@ -407,39 +407,6 @@ def test_given_lowercase_schema_when_listing_relations_then_uppercases_filter_bi assert "created, last_altered" in str(cursor.executed_sql) -@pytest.mark.parametrize( - "test_case", - [ - SnowflakeInformationSchemaFilterTestCase( - description="uppercases relation existence filter bind values", - database="analytics", - schemas=("staging",), - names=("commerce__stg_order",), - expected_params=("COMMERCE__STG_ORDER", "STAGING", "ANALYTICS"), - ) - ], - ids=lambda case: case.description, -) -def test_given_lowercase_schema_when_checking_relation_exists_then_uses_sargable_filters( - test_case: SnowflakeInformationSchemaFilterTestCase, -) -> None: - cursor: FakeSnowflakeMetadataCursor = FakeSnowflakeMetadataCursor(row=(1,)) - connection: FakeSnowflakeMetadataConnection = FakeSnowflakeMetadataConnection(cursor) - adapter: SnowflakeAdapter = SnowflakeAdapter() - - exists: bool = adapter.relation_exists( - connection=cast(Any, connection), - database=test_case.database, - schema=test_case.schemas[0], - name=test_case.names[0], - ) - - assert exists is True - assert cursor.executed_sql is not None - assert "UPPER(table_" not in cursor.executed_sql - assert cursor.executed_params == test_case.expected_params - - @pytest.mark.parametrize( "test_case", [ From 58a3c49409019490db99a684d9494da35ad4dcf6 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 15:20:54 +0100 Subject: [PATCH 2/9] perf(executor): cache build-time relation metadata per invocation Materializations looked the same relation up repeatedly: an incremental read its delta's columns three times and its target's twice, and a microbatch read delta and target columns in every batch phase. Each build now opens one relation metadata cache. Runtime column and existence lookups return the previous identical answer until a statement may have changed that relation. Every finished outermost statement is reported to the cache through an in-process StatementLifecycle listener: CREATE/DROP/ALTER/COPY and SELECT ... INTO evict entries with the same unqualified name (RENAME TO and SWAP WITH evict both sides); reads, DML, and transaction control keep them; USE, ALTER SESSION, and anything unrecognized evict everything. A read that overlaps an eviction of its relation is not stored, and on adapters with transactional DDL a COMMIT or ROLLBACK replays that thread's evictions. Nothing is persisted across invocations; migrations keep reading live columns directly. Refs CHI-455 --- .../adapter/relations/_helpers/ddl_effects.py | 165 +++++++++++++ .../_helpers/relation_metadata_context.py | 55 +++++ .../classes/relation_metadata_cache.py | 96 ++++++++ src/sqlbuild/adapter/relations/constants.py | 6 + .../relations/main/cached_relation_columns.py | 30 +++ .../relations/main/cached_relation_exists.py | 33 +++ .../main/relation_metadata_cache_scope.py | 37 +++ src/sqlbuild/adapter/relations/models.py | 9 + src/sqlbuild/adapter/relations/types.py | 4 + src/sqlbuild/executor/build/main/_execute.py | 31 +++ .../run/_helpers/execution/promotion.py | 4 +- .../executor/run/_helpers/execution/schema.py | 5 +- .../_helpers/materializations/full_refresh.py | 5 +- .../_helpers/materializations/microbatch.py | 7 +- .../run/_helpers/materializations/snapshot.py | 4 +- .../run/_helpers/validation/cursor_bounds.py | 4 +- .../_helpers/statement_listeners.py | 32 +++ .../classes/statement_lifecycle.py | 19 +- src/sqlbuild/runtime/observability/types.py | 10 +- .../_helpers/ddl_effects/_test_types.py | 14 ++ .../test_statement_metadata_effect.py | 207 ++++++++++++++++ .../relation_metadata_cache/_test_types.py | 33 +++ .../relation_metadata_cache/helpers.py | 37 +++ .../test_relation_metadata_cache.py | 231 ++++++++++++++++++ 24 files changed, 1065 insertions(+), 13 deletions(-) create mode 100644 src/sqlbuild/adapter/relations/_helpers/ddl_effects.py create mode 100644 src/sqlbuild/adapter/relations/_helpers/relation_metadata_context.py create mode 100644 src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py create mode 100644 src/sqlbuild/adapter/relations/main/cached_relation_columns.py create mode 100644 src/sqlbuild/adapter/relations/main/cached_relation_exists.py create mode 100644 src/sqlbuild/adapter/relations/main/relation_metadata_cache_scope.py create mode 100644 src/sqlbuild/runtime/observability/_helpers/statement_listeners.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/_test_types.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/_test_types.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/helpers.py create mode 100644 tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/test_relation_metadata_cache.py diff --git a/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py b/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py new file mode 100644 index 0000000000..54d5e8dc13 --- /dev/null +++ b/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py @@ -0,0 +1,165 @@ +"""Classify executed SQL by its effect on cached relation existence and columns.""" + +from __future__ import annotations + +import re + +from sqlbuild.adapter.relations.constants import ( + DOUBLE_QUOTE, + IDENTIFIER_QUOTE_PAIRS, + QUOTED_IDENTIFIER_MIN_LENGTH, + STATEMENT_SEPARATOR, +) +from sqlbuild.adapter.relations.models import StatementMetadataEffect + +_LINE_COMMENT: re.Pattern[str] = re.compile(r"--[^\n]*") +_BLOCK_COMMENT: re.Pattern[str] = re.compile(r"/\*.*?\*/", re.DOTALL) +_FIRST_WORD: re.Pattern[str] = re.compile(r"[A-Za-z]+") +_IDENTIFIER_PART: str = r'(?:"(?:[^"]|"")*"|`[^`]*`|\[[^\]]*\]|[A-Za-z0-9_$@#-]+)' +_QUALIFIED_NAME: str = rf"(?P{_IDENTIFIER_PART}(?:\s*\.\s*{_IDENTIFIER_PART})*)" +_RELATION_KIND: str = ( + r"(?:(?:OR\s+REPLACE|LOCAL|GLOBAL|TEMP|TEMPORARY|TRANSIENT|VOLATILE|SECURE|RECURSIVE|" + r"MATERIALIZED|EXTERNAL|DYNAMIC|ICEBERG|HYBRID|UNLOGGED|FOREIGN|STREAMING)\s+)*" + r"(?:TABLE|VIEW)\s+" +) +_CREATE_RELATION: re.Pattern[str] = re.compile( + rf"^CREATE\s+{_RELATION_KIND}(?:IF\s+NOT\s+EXISTS\s+)?{_QUALIFIED_NAME}", + re.IGNORECASE, +) +_DROP_RELATION: re.Pattern[str] = re.compile( + rf"^DROP\s+{_RELATION_KIND}(?:IF\s+EXISTS\s+)?{_QUALIFIED_NAME}", re.IGNORECASE +) +_ALTER_RELATION: re.Pattern[str] = re.compile( + rf"^ALTER\s+{_RELATION_KIND}(?:IF\s+EXISTS\s+)?{_QUALIFIED_NAME}", re.IGNORECASE +) +_ALTER_SECOND_RELATION: re.Pattern[str] = re.compile( + rf"\b(?:RENAME\s+TO|SWAP\s+WITH)\s+{_QUALIFIED_NAME}", re.IGNORECASE +) +_SELECT_INTO_RELATION: re.Pattern[str] = re.compile(rf"\bINTO\s+{_QUALIFIED_NAME}", re.IGNORECASE) +_COPY_INTO_RELATION: re.Pattern[str] = re.compile( + rf"^COPY\s+INTO\s+{_QUALIFIED_NAME}", re.IGNORECASE +) +_CREATE_SCHEMA_IF_NOT_EXISTS: re.Pattern[str] = re.compile( + r"^CREATE\s+SCHEMA\s+IF\s+NOT\s+EXISTS\s+[^\s;]+\s*;?\s*$", re.IGNORECASE +) +_RELATION_NEUTRAL_OBJECT: re.Pattern[str] = re.compile( + r"^(?:CREATE|DROP|ALTER)\s+(?:OR\s+REPLACE\s+)?(?:(?:SECURE|TEMP|TEMPORARY|AGGREGATE)\s+)*" + r"(?:FUNCTION|PROCEDURE|SEQUENCE|STAGE|FILE\s+FORMAT|INDEX|UNIQUE\s+INDEX|WAREHOUSE|MACRO)\b", + re.IGNORECASE, +) +_QUERY_STATEMENTS: frozenset[str] = frozenset({"SELECT", "WITH"}) +_READ_OR_DATA_STATEMENTS: frozenset[str] = frozenset( + { + "VALUES", + "SHOW", + "DESCRIBE", + "DESC", + "EXPLAIN", + "INSERT", + "UPDATE", + "DELETE", + "MERGE", + "TRUNCATE", + "BEGIN", + "START", + "SAVEPOINT", + "RELEASE", + "SET", + "UNSET", + "GRANT", + "REVOKE", + "COMMENT", + "PUT", + "GET", + "LIST", + "LS", + "REMOVE", + "RM", + "ANALYZE", + } +) +_TRANSACTION_END_STATEMENTS: frozenset[str] = frozenset({"COMMIT", "ROLLBACK", "END", "ABORT"}) +_RELATION_DDL_STATEMENTS: frozenset[str] = frozenset({"CREATE", "DROP", "ALTER", "COPY"}) + + +_NO_EFFECT: StatementMetadataEffect = StatementMetadataEffect() +_ALL_RELATIONS: StatementMetadataEffect = StatementMetadataEffect(invalidates_all=True) +_TRANSACTION_END: StatementMetadataEffect = StatementMetadataEffect(ends_transaction=True) + + +def statement_metadata_effect(sql: str) -> StatementMetadataEffect: + """Return which cached relation names ``sql`` may have changed.""" + + statement: str = ( + _BLOCK_COMMENT.sub(" ", _LINE_COMMENT.sub(" ", sql)) + .strip() + .rstrip(STATEMENT_SEPARATOR) + .strip() + .lstrip("(") + ) + first: re.Match[str] | None = _FIRST_WORD.match(statement) + if first is None or STATEMENT_SEPARATOR in statement: + return _ALL_RELATIONS + keyword: str = first.group(0).upper() + if keyword in _QUERY_STATEMENTS: + return _select_into_effect(statement=statement) + if keyword in _READ_OR_DATA_STATEMENTS: + return _NO_EFFECT + if keyword in _TRANSACTION_END_STATEMENTS: + return _TRANSACTION_END + if keyword not in _RELATION_DDL_STATEMENTS: + return _ALL_RELATIONS + if ( + _CREATE_SCHEMA_IF_NOT_EXISTS.match(statement) is not None + or _RELATION_NEUTRAL_OBJECT.match(statement) is not None + ): + return _NO_EFFECT + return _relation_ddl_effect(statement=statement) + + +def relation_name_key(name: str) -> str: + """Return the invalidation key for one unquoted or quoted unqualified relation name.""" + + return _strip_identifier_quotes(name).casefold() + + +def _select_into_effect(*, statement: str) -> StatementMetadataEffect: + created: re.Match[str] | None = _SELECT_INTO_RELATION.search(statement) + if created is None: + return _NO_EFFECT + return StatementMetadataEffect( + relation_names=frozenset({_unqualified_name_key(created.group("name"))}) + ) + + +def _relation_ddl_effect(*, statement: str) -> StatementMetadataEffect: + pattern: re.Pattern[str] + for pattern in (_CREATE_RELATION, _DROP_RELATION, _ALTER_RELATION, _COPY_INTO_RELATION): + match: re.Match[str] | None = pattern.match(statement) + if match is None: + continue + names: set[str] = {_unqualified_name_key(match.group("name"))} + if pattern is _ALTER_RELATION: + names.update( + _unqualified_name_key(second.group("name")) + for second in _ALTER_SECOND_RELATION.finditer(statement, match.end()) + ) + return StatementMetadataEffect(relation_names=frozenset(names)) + return _ALL_RELATIONS + + +def _unqualified_name_key(qualified: str) -> str: + parts: list[str] = re.findall(_IDENTIFIER_PART, qualified) + last: str = _strip_identifier_quotes(parts[-1]) + return last.rsplit(".", 1)[-1].casefold() + + +def _strip_identifier_quotes(part: str) -> str: + if len(part) < QUOTED_IDENTIFIER_MIN_LENGTH: + return part + edges: tuple[str, str] = (part[0], part[-1]) + if edges == (DOUBLE_QUOTE, DOUBLE_QUOTE): + return part[1:-1].replace(DOUBLE_QUOTE * 2, DOUBLE_QUOTE) + if edges in IDENTIFIER_QUOTE_PAIRS: + return part[1:-1] + return part diff --git a/src/sqlbuild/adapter/relations/_helpers/relation_metadata_context.py b/src/sqlbuild/adapter/relations/_helpers/relation_metadata_context.py new file mode 100644 index 0000000000..52885f6a29 --- /dev/null +++ b/src/sqlbuild/adapter/relations/_helpers/relation_metadata_context.py @@ -0,0 +1,55 @@ +"""Invocation-local relation metadata cache state.""" + +from __future__ import annotations + +from collections.abc import Callable +from contextvars import ContextVar, Token +from typing import Any + +from sqlbuild.adapter.relations.classes.relation_metadata_cache import RelationMetadataCache + +_ACTIVE_CACHE: ContextVar[tuple[Any, RelationMetadataCache] | None] = ContextVar( + "sqlbuild_relation_metadata_cache", default=None +) + + +def current_relation_metadata_cache(*, adapter: Any) -> RelationMetadataCache | None: + """Return the cache open for exactly this adapter, if any.""" + + active: tuple[Any, RelationMetadataCache] | None = _ACTIVE_CACHE.get() + if active is None or active[0] is not adapter: + return None + return active[1] + + +def activate_relation_metadata_cache( + *, adapter: Any, cache: RelationMetadataCache +) -> Token[tuple[Any, RelationMetadataCache] | None]: + """Make ``cache`` the open cache for ``adapter`` in this context.""" + + return _ACTIVE_CACHE.set((adapter, cache)) + + +def deactivate_relation_metadata_cache( + token: Token[tuple[Any, RelationMetadataCache] | None], +) -> None: + """Restore the cache that was open before ``activate_relation_metadata_cache``.""" + + _ACTIVE_CACHE.reset(token) + + +def read_through_relation_metadata_cache[ValueT]( + *, + adapter: Any, + kind: str, + database: str | None, + schema: str | None, + name: str, + read: Callable[[], ValueT], +) -> ValueT: + """Answer one lookup from the cache open for ``adapter``, or read it when none is open.""" + + cache: RelationMetadataCache | None = current_relation_metadata_cache(adapter=adapter) + if cache is None: + return read() + return cache.lookup(kind=kind, database=database, schema=schema, name=name, read=read) diff --git a/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py b/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py new file mode 100644 index 0000000000..4a757a20aa --- /dev/null +++ b/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py @@ -0,0 +1,96 @@ +"""Invocation-scoped relation existence and column metadata shared by build workers.""" + +from __future__ import annotations + +import threading +from collections.abc import Callable +from typing import cast + +from sqlbuild.adapter.relations._helpers.ddl_effects import ( + relation_name_key, + statement_metadata_effect, +) +from sqlbuild.adapter.relations.models import StatementMetadataEffect +from sqlbuild.adapter.relations.types import RelationCacheKey, RelationCacheVersion + + +class RelationMetadataCache: + """Serve repeated identical metadata lookups until SQL may have changed the relation.""" + + def __init__(self, *, transactional_ddl: bool) -> None: + self._transactional_ddl: bool = transactional_ddl + self._lock: threading.Lock = threading.Lock() + self._entries: dict[RelationCacheKey, object] = {} + self._name_versions: dict[str, int] = {} + self._global_version: int = 0 + self._pending_names: dict[int, set[str]] = {} + self._pending_all: set[int] = set() + + def lookup[ValueT]( + self, + *, + kind: str, + database: str | None, + schema: str | None, + name: str, + read: Callable[[], ValueT], + ) -> ValueT: + """Return the cached answer for this exact lookup, reading it on a miss.""" + + return cast(ValueT, self._lookup(key=(kind, database, schema, name), read=read)) + + def observe_statement(self, *, sql: str) -> None: + """Invalidate entries the finished statement ``sql`` may have changed.""" + + effect: StatementMetadataEffect = statement_metadata_effect(sql) + thread_id: int = threading.get_ident() + with self._lock: + if effect.ends_transaction and self._transactional_ddl: + names: set[str] = self._pending_names.pop(thread_id, set()) + if thread_id in self._pending_all: + self._pending_all.discard(thread_id) + self._invalidate_all() + else: + self._invalidate_names(names=frozenset(names)) + return + if effect.invalidates_all: + if self._transactional_ddl: + self._pending_all.add(thread_id) + self._invalidate_all() + return + if effect.relation_names: + if self._transactional_ddl: + self._pending_names.setdefault(thread_id, set()).update(effect.relation_names) + self._invalidate_names(names=effect.relation_names) + + def _lookup(self, *, key: RelationCacheKey, read: Callable[[], object]) -> object: + name_key: str = relation_name_key(key[3]) + with self._lock: + if key in self._entries: + return self._entries[key] + version: RelationCacheVersion = self._version(name_key=name_key) + value: object = read() + with self._lock: + if self._version(name_key=name_key) == version: + self._entries[key] = value + return value + + def _version(self, *, name_key: str) -> RelationCacheVersion: + return self._global_version, self._name_versions.get(name_key, 0) + + def _invalidate_names(self, *, names: frozenset[str]) -> None: + if not names: + return + name: str + for name in names: + self._name_versions[name] = self._name_versions.get(name, 0) + 1 + stale: list[RelationCacheKey] = [ + key for key in self._entries if relation_name_key(key[3]) in names + ] + key: RelationCacheKey + for key in stale: + del self._entries[key] + + def _invalidate_all(self) -> None: + self._global_version += 1 + self._entries.clear() diff --git a/src/sqlbuild/adapter/relations/constants.py b/src/sqlbuild/adapter/relations/constants.py index c0d775459f..750edf868b 100644 --- a/src/sqlbuild/adapter/relations/constants.py +++ b/src/sqlbuild/adapter/relations/constants.py @@ -4,3 +4,9 @@ INSPECTION_IN_LIST_LIMIT: int = 50 METADATA_NAME_FILTER_LIMIT: int = 250 SHORTENED_LOGICAL_NAME_HASH_LENGTH: int = 8 +STATEMENT_SEPARATOR: str = ";" +QUOTED_IDENTIFIER_MIN_LENGTH: int = 2 +DOUBLE_QUOTE: str = '"' +IDENTIFIER_QUOTE_PAIRS: frozenset[tuple[str, str]] = frozenset({("`", "`"), ("[", "]")}) +RELATION_COLUMNS_LOOKUP: str = "columns" +RELATION_EXISTS_LOOKUP: str = "exists" diff --git a/src/sqlbuild/adapter/relations/main/cached_relation_columns.py b/src/sqlbuild/adapter/relations/main/cached_relation_columns.py new file mode 100644 index 0000000000..bb9c28fa80 --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/cached_relation_columns.py @@ -0,0 +1,30 @@ +"""Read one relation's columns through the open relation metadata cache.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.adapter.contract.models import ColumnInfo +from sqlbuild.adapter.relations._helpers.relation_metadata_context import ( + read_through_relation_metadata_cache, +) +from sqlbuild.adapter.relations.constants import RELATION_COLUMNS_LOOKUP + + +def cached_relation_columns( + *, adapter: BaseAdapter, connection: Any, database: str | None, schema: str | None, name: str +) -> tuple[ColumnInfo, ...]: + """Return ``adapter.get_columns`` for this lookup, reusing a cached identical read.""" + + return read_through_relation_metadata_cache( + adapter=adapter, + kind=RELATION_COLUMNS_LOOKUP, + database=database, + schema=schema, + name=name, + read=partial( + adapter.get_columns, connection=connection, database=database, schema=schema, name=name + ), + ) diff --git a/src/sqlbuild/adapter/relations/main/cached_relation_exists.py b/src/sqlbuild/adapter/relations/main/cached_relation_exists.py new file mode 100644 index 0000000000..00dc8080ed --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/cached_relation_exists.py @@ -0,0 +1,33 @@ +"""Check one relation's existence through the open relation metadata cache.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.adapter.relations._helpers.relation_metadata_context import ( + read_through_relation_metadata_cache, +) +from sqlbuild.adapter.relations.constants import RELATION_EXISTS_LOOKUP + + +def cached_relation_exists( + *, adapter: BaseAdapter, connection: Any, database: str | None, schema: str | None, name: str +) -> bool: + """Return ``adapter.relation_exists`` for this lookup, reusing a cached identical read.""" + + return read_through_relation_metadata_cache( + adapter=adapter, + kind=RELATION_EXISTS_LOOKUP, + database=database, + schema=schema, + name=name, + read=partial( + adapter.relation_exists, + connection=connection, + database=database, + schema=schema, + name=name, + ), + ) diff --git a/src/sqlbuild/adapter/relations/main/relation_metadata_cache_scope.py b/src/sqlbuild/adapter/relations/main/relation_metadata_cache_scope.py new file mode 100644 index 0000000000..64ebe51ca8 --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/relation_metadata_cache_scope.py @@ -0,0 +1,37 @@ +"""Scope one relation metadata cache to a build 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.relation_metadata_context import ( + activate_relation_metadata_cache, + current_relation_metadata_cache, + deactivate_relation_metadata_cache, +) +from sqlbuild.adapter.relations.classes.relation_metadata_cache import RelationMetadataCache +from sqlbuild.runtime.observability.classes.statement_lifecycle import StatementLifecycle + + +@contextmanager +def relation_metadata_cache_scope(*, adapter: Any) -> Iterator[RelationMetadataCache]: + """Share relation metadata reads until the scope exits; nested scopes reuse the open cache.""" + + existing: RelationMetadataCache | None = current_relation_metadata_cache(adapter=adapter) + if existing is not None: + yield existing + return + cache: RelationMetadataCache = RelationMetadataCache( + transactional_ddl=bool(adapter.supports_transactional_ddl()) + ) + token: Token[tuple[Any, RelationMetadataCache] | None] = activate_relation_metadata_cache( + adapter=adapter, cache=cache + ) + try: + with StatementLifecycle.listener_scope(cache.observe_statement): + yield cache + finally: + _ = deactivate_relation_metadata_cache(token) diff --git a/src/sqlbuild/adapter/relations/models.py b/src/sqlbuild/adapter/relations/models.py index 6c7fba5721..40ed29e19e 100644 --- a/src/sqlbuild/adapter/relations/models.py +++ b/src/sqlbuild/adapter/relations/models.py @@ -42,3 +42,12 @@ class InspectionQueryRecord: elapsed_seconds: float row_count: int | None error: str | None = None + + +@dataclass(frozen=True) +class StatementMetadataEffect: + """What one executed statement may have changed in cached relation metadata.""" + + relation_names: frozenset[str] = frozenset() + invalidates_all: bool = False + ends_transaction: bool = False diff --git a/src/sqlbuild/adapter/relations/types.py b/src/sqlbuild/adapter/relations/types.py index abce8d457f..acdad6b3b2 100644 --- a/src/sqlbuild/adapter/relations/types.py +++ b/src/sqlbuild/adapter/relations/types.py @@ -9,6 +9,10 @@ from sqlbuild.adapter.relations.models import SchemaColumnListing, SchemaRelationListing +type RelationCacheKey = tuple[str, str | None, str | None, str] +type RelationCacheVersion = tuple[int, int] + + class RelationLocation(Protocol): """Structural relation location accepted by adapter naming.""" diff --git a/src/sqlbuild/executor/build/main/_execute.py b/src/sqlbuild/executor/build/main/_execute.py index 310d6e4a92..a9fe682714 100644 --- a/src/sqlbuild/executor/build/main/_execute.py +++ b/src/sqlbuild/executor/build/main/_execute.py @@ -6,6 +6,9 @@ from typing import Any from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.adapter.relations.main.relation_metadata_cache_scope import ( + relation_metadata_cache_scope, +) from sqlbuild.compiler.planner.models import PlanOutput from sqlbuild.compiler.planner.types import RetentionPlanPhase from sqlbuild.executor.auditing.main._project_results import ( @@ -54,6 +57,34 @@ def execute_build_plan( ) -> BuildExecutionResult: """Execute a full build plan using the DAG scheduler.""" + with relation_metadata_cache_scope(adapter=adapter): + return _execute_build_plan( + plan=plan, + adapter=adapter, + connection_config=connection_config, + connections=connections, + scheduler_connection=scheduler_connection, + runtime=runtime, + callbacks=callbacks, + customizations=customizations, + initial_state=initial_state, + schema_prepared=schema_prepared, + ) + + +def _execute_build_plan( + *, + plan: PlanOutput, + adapter: BaseAdapter, + connection_config: dict[str, object], + connections: tuple[Any, ...], + scheduler_connection: Any, + runtime: BuildRuntimeParams, + callbacks: BuildCallbacks | None, + customizations: BuildCustomizations | None, + initial_state: BuildInitialState | None, + schema_prepared: bool, +) -> BuildExecutionResult: resolved_callbacks: BuildCallbacks = callbacks if callbacks is not None else BuildCallbacks() resolved_customizations: BuildCustomizations = ( customizations if customizations is not None else BuildCustomizations() diff --git a/src/sqlbuild/executor/run/_helpers/execution/promotion.py b/src/sqlbuild/executor/run/_helpers/execution/promotion.py index 99f041e51a..c3c20f17d8 100644 --- a/src/sqlbuild/executor/run/_helpers/execution/promotion.py +++ b/src/sqlbuild/executor/run/_helpers/execution/promotion.py @@ -7,6 +7,7 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.types import PromotionStrategy +from sqlbuild.adapter.relations.main.cached_relation_exists import cached_relation_exists from sqlbuild.errors.contracts.exceptions import ExecutorInputError from sqlbuild.executor.run.models import BoundViewGuard from sqlbuild.runtime.observability.classes.operation_lifecycle import ( @@ -47,7 +48,8 @@ def promote_relation_to_destination( target_kind="relation", ), ) as lifecycle: - existing: bool = adapter.relation_exists( + existing: bool = cached_relation_exists( + adapter=adapter, connection=connection, database=destination_database, schema=destination_schema, diff --git a/src/sqlbuild/executor/run/_helpers/execution/schema.py b/src/sqlbuild/executor/run/_helpers/execution/schema.py index 6fdc0ac1d7..607a1b3fd5 100644 --- a/src/sqlbuild/executor/run/_helpers/execution/schema.py +++ b/src/sqlbuild/executor/run/_helpers/execution/schema.py @@ -6,6 +6,7 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.models import ColumnInfo +from sqlbuild.adapter.relations.main.cached_relation_columns import cached_relation_columns from sqlbuild.runtime.observability.classes.operation_lifecycle import ( OperationAttributes, OperationLifecycle, @@ -34,8 +35,8 @@ def inspect_runtime_relation_schema( target_kind="relation", ), ) as lifecycle: - columns: tuple[ColumnInfo, ...] = adapter.get_columns( - connection=connection, database=database, schema=schema, name=name + columns: tuple[ColumnInfo, ...] = cached_relation_columns( + adapter=adapter, connection=connection, database=database, schema=schema, name=name ) lifecycle.completed(metadata={"item_count": len(columns)}) return columns diff --git a/src/sqlbuild/executor/run/_helpers/materializations/full_refresh.py b/src/sqlbuild/executor/run/_helpers/materializations/full_refresh.py index aeb3029caa..012aa5470d 100644 --- a/src/sqlbuild/executor/run/_helpers/materializations/full_refresh.py +++ b/src/sqlbuild/executor/run/_helpers/materializations/full_refresh.py @@ -8,6 +8,7 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.types import BuiltinAdapter +from sqlbuild.adapter.relations.main.cached_relation_exists import cached_relation_exists from sqlbuild.adapter.relations.main.resolve_qualified_name_parts import ( resolve_qualified_name_parts, ) @@ -70,8 +71,8 @@ def relation_exists( ) -> bool: """Inspect whether one full-refresh relation exists.""" - return adapter.relation_exists( - connection=connection, database=database, schema=schema, name=name + return cached_relation_exists( + adapter=adapter, connection=connection, database=database, schema=schema, name=name ) diff --git a/src/sqlbuild/executor/run/_helpers/materializations/microbatch.py b/src/sqlbuild/executor/run/_helpers/materializations/microbatch.py index b57dd613fe..6580032c99 100644 --- a/src/sqlbuild/executor/run/_helpers/materializations/microbatch.py +++ b/src/sqlbuild/executor/run/_helpers/materializations/microbatch.py @@ -14,6 +14,7 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.models import ColumnInfo +from sqlbuild.adapter.relations.main.cached_relation_exists import cached_relation_exists from sqlbuild.adapter.relations.main.resolve_qualified_name_parts import ( resolve_qualified_name_parts, ) @@ -654,7 +655,8 @@ def _no_work_microbatch_result( ) -> ModelExecutionResult | None: if batch_plan.batches: return None - destination_exists: bool = context.adapter.relation_exists( + destination_exists: bool = cached_relation_exists( + adapter=context.adapter, connection=context.connection, database=context.entry.destination.database, schema=context.entry.destination.schema, @@ -2345,7 +2347,8 @@ def _physical_cursor_envelope( state: MicrobatchLifecycleState, targets: MicrobatchTargets, ) -> CursorBounds | None | ModelExecutionResult: - if not context.adapter.relation_exists( + if not cached_relation_exists( + adapter=context.adapter, connection=context.connection, database=targets.target_database, schema=targets.target_schema, diff --git a/src/sqlbuild/executor/run/_helpers/materializations/snapshot.py b/src/sqlbuild/executor/run/_helpers/materializations/snapshot.py index 1572b3bfaa..0e6f458515 100644 --- a/src/sqlbuild/executor/run/_helpers/materializations/snapshot.py +++ b/src/sqlbuild/executor/run/_helpers/materializations/snapshot.py @@ -12,6 +12,7 @@ SnapshotChangeTarget, ) from sqlbuild.adapter.contract.types import TypeFamily +from sqlbuild.adapter.relations.main.cached_relation_exists import cached_relation_exists from sqlbuild.adapter.relations.main.resolve_qualified_name_parts import ( resolve_qualified_name_parts, ) @@ -443,7 +444,8 @@ def _apply_snapshot_phase( check_columns=check_columns, statement_recorder=statement_recorder, ) - target_exists: bool = context.adapter.relation_exists( + target_exists: bool = cached_relation_exists( + adapter=context.adapter, connection=context.connection, database=entry.destination.database, schema=entry.destination.schema, diff --git a/src/sqlbuild/executor/run/_helpers/validation/cursor_bounds.py b/src/sqlbuild/executor/run/_helpers/validation/cursor_bounds.py index 11f1e97d25..c07bbbf56e 100644 --- a/src/sqlbuild/executor/run/_helpers/validation/cursor_bounds.py +++ b/src/sqlbuild/executor/run/_helpers/validation/cursor_bounds.py @@ -9,6 +9,7 @@ from typing import Any from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter +from sqlbuild.adapter.relations.main.cached_relation_exists import cached_relation_exists from sqlbuild.compiler.compile.main.cursor_intrinsics import resolve_cursor_intrinsics from sqlbuild.compiler.planner.constants import MICROBATCH_END_SENTINEL, MICROBATCH_START_SENTINEL from sqlbuild.compiler.planner.main.execution.cursor_replay_policy import ( @@ -625,7 +626,8 @@ def _query_target_max( ) -> CursorScalar | None: """Read the target cursor high-water mark or None when the target does not exist.""" - if not adapter.relation_exists( + if not cached_relation_exists( + adapter=adapter, connection=connection, database=target_database, schema=target_schema, diff --git a/src/sqlbuild/runtime/observability/_helpers/statement_listeners.py b/src/sqlbuild/runtime/observability/_helpers/statement_listeners.py new file mode 100644 index 0000000000..0060834c4a --- /dev/null +++ b/src/sqlbuild/runtime/observability/_helpers/statement_listeners.py @@ -0,0 +1,32 @@ +"""Scoped in-process listeners told about every finished outermost statement.""" + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar, Token + +from sqlbuild.runtime.observability.types import StatementListener + +_STATEMENT_LISTENERS: ContextVar[tuple[StatementListener, ...]] = ContextVar( + "sqlbuild_statement_listeners", default=() +) + + +def notify_statement_finished(*, sql: str) -> None: + """Tell every listener in this context that ``sql`` finished, successfully or not.""" + + listener: StatementListener + for listener in _STATEMENT_LISTENERS.get(): + listener(sql=sql) + + +@contextmanager +def statement_listener_scope(listener: StatementListener) -> Iterator[None]: + """Deliver each finished statement's SQL text to ``listener`` while the scope is open.""" + + token: Token[tuple[StatementListener, ...]] = _STATEMENT_LISTENERS.set( + (*_STATEMENT_LISTENERS.get(), listener) + ) + try: + yield + finally: + _STATEMENT_LISTENERS.reset(token) diff --git a/src/sqlbuild/runtime/observability/classes/statement_lifecycle.py b/src/sqlbuild/runtime/observability/classes/statement_lifecycle.py index 16d9e7a2bf..b4e034e447 100644 --- a/src/sqlbuild/runtime/observability/classes/statement_lifecycle.py +++ b/src/sqlbuild/runtime/observability/classes/statement_lifecycle.py @@ -7,8 +7,8 @@ import re import threading import time -from collections.abc import Callable, Mapping -from contextlib import ExitStack +from collections.abc import Callable, Iterator, Mapping +from contextlib import ExitStack, contextmanager from contextvars import ContextVar, Token from types import TracebackType @@ -22,6 +22,10 @@ invocation_scope, statement_scope, ) +from sqlbuild.runtime.observability._helpers.statement_listeners import ( + notify_statement_finished, + statement_listener_scope, +) from sqlbuild.runtime.observability.classes.event_dispatcher import EventDispatcher from sqlbuild.runtime.observability.classes.statement_monitor import ( StatementMonitor, @@ -29,7 +33,7 @@ from sqlbuild.runtime.observability.constants import STATEMENT_HEARTBEAT_THRESHOLD_SECONDS from sqlbuild.runtime.observability.exceptions import ObservabilityValidationError from sqlbuild.runtime.observability.models import ExecutionIdentity -from sqlbuild.runtime.observability.types import JSONValue +from sqlbuild.runtime.observability.types import JSONValue, StatementListener _MAX_ERROR_VALUE_LENGTH: int = 128 _STATEMENT_KIND_PATTERN: re.Pattern[str] = re.compile(r"^([A-Za-z]{1,64})(?![A-Za-z0-9_])") @@ -78,6 +82,14 @@ def __init__( ) self._pending_failure: tuple[BaseException, str | None, str | None] | None = None + @staticmethod + @contextmanager + def listener_scope(listener: StatementListener) -> Iterator[None]: + """Deliver each finished outermost statement's SQL text to an in-process listener.""" + + with statement_listener_scope(listener): + yield + def __enter__(self) -> StatementLifecycle: execution_owner_key: _ExecutionOwnerKey = _current_execution_owner_key() owner: StatementLifecycle | None = _STATEMENT_LIFECYCLE_OWNER.get() @@ -253,6 +265,7 @@ def __exit__( _STATEMENT_LIFECYCLE_OWNER.reset(self._owner_token) self._owner_token = None self._stack.close() + notify_statement_finished(sql=self._sql) def _base_payload(self) -> dict[str, JSONValue]: payload: dict[str, JSONValue] = { diff --git a/src/sqlbuild/runtime/observability/types.py b/src/sqlbuild/runtime/observability/types.py index 8232b53cdf..94cb1e76e9 100644 --- a/src/sqlbuild/runtime/observability/types.py +++ b/src/sqlbuild/runtime/observability/types.py @@ -1,7 +1,7 @@ """Type declarations for runtime observability contracts.""" from collections.abc import Callable, Mapping -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Protocol if TYPE_CHECKING: from sqlbuild.runtime.observability.models import ( @@ -22,3 +22,11 @@ object, KnownLifecycleSubscriber | OpaqueLifecycleSubscriber, bool ] type DiagnosticRegistration = tuple[object, DiagnosticSubscriber] + + +class StatementListener(Protocol): + """In-process callback receiving the SQL text of each finished outermost statement.""" + + def __call__(self, *, sql: str) -> None: + """Handle one statement that finished, successfully or not.""" + ... diff --git a/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/_test_types.py b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/_test_types.py new file mode 100644 index 0000000000..93508cddb1 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/_test_types.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class StatementMetadataEffectTestCase: + """One executed statement and the cached relation names it may have changed.""" + + description: str + sql: str + expected_relation_names: frozenset[str] + expected_invalidates_all: bool + expected_ends_transaction: bool = False diff --git a/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py new file mode 100644 index 0000000000..b908b44d05 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py @@ -0,0 +1,207 @@ +"""Executed SQL is classified conservatively by its effect on cached relation metadata.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.adapter.relations._helpers.ddl_effects import statement_metadata_effect +from sqlbuild.adapter.relations.models import StatementMetadataEffect +from tests.unit.src.sqlbuild.adapter.relations._helpers.ddl_effects._test_types import ( + StatementMetadataEffectTestCase, +) + +_NONE: frozenset[str] = frozenset() + + +@pytest.mark.parametrize( + "test_case", + [ + StatementMetadataEffectTestCase( + description="create or replace transient table as select", + sql="CREATE OR REPLACE TRANSIENT TABLE analytics.marts.orders__delta AS SELECT 1 AS id", + expected_relation_names=frozenset({"orders__delta"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="quoted three-part view", + sql='CREATE OR REPLACE VIEW "ANALYTICS"."MARTS"."ORDERS_V" AS SELECT 1 AS id', + expected_relation_names=frozenset({"orders_v"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="clone names only the new relation", + sql="CREATE TABLE marts.orders__prev CLONE marts.orders", + expected_relation_names=frozenset({"orders__prev"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="drop table if exists", + sql="DROP TABLE IF EXISTS analytics.marts.orders__delta", + expected_relation_names=frozenset({"orders__delta"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="rename names both relations", + sql="ALTER TABLE marts.orders__stage RENAME TO marts.orders", + expected_relation_names=frozenset({"orders__stage", "orders"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="swap names both relations", + sql="ALTER TABLE analytics.marts.orders SWAP WITH analytics.marts.orders__stage", + expected_relation_names=frozenset({"orders", "orders__stage"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="add column after a leading comment", + sql='-- schema change\nALTER TABLE marts.orders ADD COLUMN "NOTE" VARCHAR', + expected_relation_names=frozenset({"orders"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="column rename is not a relation rename", + sql='ALTER TABLE marts.orders RENAME COLUMN "OLD" TO "NEW"', + expected_relation_names=frozenset({"orders"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="backtick qualified name", + sql="CREATE OR REPLACE TABLE `project.dataset.orders` AS SELECT 1 AS id", + expected_relation_names=frozenset({"orders"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="bracket qualified name", + sql="DROP TABLE IF EXISTS [dbo].[orders__delta]", + expected_relation_names=frozenset({"orders__delta"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="select into creates a table", + sql="SELECT * INTO [dbo].[orders__staging] FROM (SELECT 1 AS id) AS __create_source", + expected_relation_names=frozenset({"orders__staging"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="common table expression select into creates a table", + sql="WITH v AS (SELECT 1 AS id) SELECT id INTO dbo.orders_history FROM v", + expected_relation_names=frozenset({"orders_history"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="copy into a table may evolve its columns", + sql="COPY INTO raw.orders FROM @landing/orders", + expected_relation_names=frozenset({"orders"}), + expected_invalidates_all=False, + ), + ], + ids=lambda case: case.description, +) +def test_given_targeted_ddl_statement_when_classifying_then_reports_possible_metadata_changes( + test_case: StatementMetadataEffectTestCase, +) -> None: + effect: StatementMetadataEffect = statement_metadata_effect(test_case.sql) + + assert effect.relation_names == test_case.expected_relation_names + assert effect.invalidates_all is test_case.expected_invalidates_all + assert effect.ends_transaction is test_case.expected_ends_transaction + + +@pytest.mark.parametrize( + "test_case", + [ + StatementMetadataEffectTestCase( + description=sql.split()[0].lower() + " " + str(index), + sql=sql, + expected_relation_names=_NONE, + expected_invalidates_all=False, + ) + for index, sql in enumerate( + ( + "SELECT COUNT(*) FROM marts.orders__delta", + "WITH totals AS (SELECT 1 AS id) SELECT * FROM totals", + "SHOW COLUMNS IN TABLE analytics.marts.orders", + "INSERT INTO marts.orders (id) SELECT id FROM marts.orders__delta", + "MERGE INTO marts.orders AS t USING (SELECT 1 AS id) AS s ON t.id = s.id", + "DELETE FROM marts.orders WHERE id = 1", + "BEGIN", + "CREATE SCHEMA IF NOT EXISTS analytics.marts", + "CREATE OR REPLACE FUNCTION marts.add_one(x NUMBER) RETURNS NUMBER AS $$ x + 1 $$", + "TRUNCATE TABLE marts.orders", + ) + ) + ], + ids=lambda case: case.description, +) +def test_given_read_or_data_statement_when_classifying_then_reports_possible_metadata_changes( + test_case: StatementMetadataEffectTestCase, +) -> None: + effect: StatementMetadataEffect = statement_metadata_effect(test_case.sql) + + assert effect.relation_names == test_case.expected_relation_names + assert effect.invalidates_all is test_case.expected_invalidates_all + assert effect.ends_transaction is test_case.expected_ends_transaction + + +@pytest.mark.parametrize( + "test_case", + [ + StatementMetadataEffectTestCase( + description=description, + sql=sql, + expected_relation_names=_NONE, + expected_invalidates_all=True, + ) + for description, sql in ( + ("procedure call", "CALL marts.rebuild_orders()"), + ("replace schema drops its relations", "CREATE OR REPLACE SCHEMA analytics.marts"), + ("drop schema", "DROP SCHEMA IF EXISTS analytics.marts"), + ("multi statement", "INSERT INTO marts.orders VALUES (1); DROP TABLE marts.orders"), + ("dynamic sql", "EXECUTE IMMEDIATE 'DROP TABLE marts.orders'"), + ( + "session schema change re-resolves unqualified lookups", + "USE SCHEMA analytics.staging", + ), + ("session parameters may change identifier resolution", "ALTER SESSION SET X = 1"), + ("unparseable relation name", 'CREATE TABLE "unterminated AS SELECT 1'), + ("empty statement", " "), + ) + ], + ids=lambda case: case.description, +) +def test_given_unbounded_statement_when_classifying_then_reports_possible_metadata_changes( + test_case: StatementMetadataEffectTestCase, +) -> None: + effect: StatementMetadataEffect = statement_metadata_effect(test_case.sql) + + assert effect.relation_names == test_case.expected_relation_names + assert effect.invalidates_all is test_case.expected_invalidates_all + assert effect.ends_transaction is test_case.expected_ends_transaction + + +@pytest.mark.parametrize( + "test_case", + [ + StatementMetadataEffectTestCase( + description=sql.lower(), + sql=sql, + expected_relation_names=_NONE, + expected_invalidates_all=False, + expected_ends_transaction=True, + ) + for sql in ("COMMIT", "ROLLBACK") + ], + ids=lambda case: case.description, +) +def test_given_transaction_end_statement_when_classifying_then_reports_possible_metadata_changes( + test_case: StatementMetadataEffectTestCase, +) -> None: + effect: StatementMetadataEffect = statement_metadata_effect(test_case.sql) + + assert effect.relation_names == test_case.expected_relation_names + assert effect.invalidates_all is test_case.expected_invalidates_all + assert effect.ends_transaction is test_case.expected_ends_transaction + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/_test_types.py b/tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/_test_types.py new file mode 100644 index 0000000000..eb198fc12f --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/_test_types.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class CacheInvalidationTestCase: + """Statements executed between two identical column lookups of ``orders``.""" + + description: str + transactional_ddl: bool + statements_between: tuple[str, ...] + expected_reads: int + + +@dataclass(frozen=True) +class OverlappingReadTestCase: + """A statement that finishes while a lookup of ``orders`` is in flight.""" + + description: str + statement_during_read: str + expected_reads: int + + +@dataclass(frozen=True) +class LiveInvalidationTestCase: + """DDL executed through the adapter between two lookups on a real DuckDB connection.""" + + description: str + ddl: str + expected_first_columns: tuple[str, ...] + expected_second_columns: tuple[str, ...] + expected_second_exists: bool diff --git a/tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/helpers.py b/tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/helpers.py new file mode 100644 index 0000000000..459e4d5fd3 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/helpers.py @@ -0,0 +1,37 @@ +"""Helpers for relation metadata cache tests.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from sqlbuild.adapter.contract.models import ColumnInfo +from sqlbuild.adapter.relations.classes.relation_metadata_cache import RelationMetadataCache +from sqlbuild.adapter.relations.constants import RELATION_COLUMNS_LOOKUP + + +@dataclass +class CountingColumnsRead: + """Column read double that counts calls and can run a statement while reading.""" + + cache: RelationMetadataCache + statement_during_read: str | None = None + reads: int = 0 + columns: tuple[ColumnInfo, ...] = field( + default_factory=lambda: (ColumnInfo(name="order_id", type="NUMBER(38,0)"),) + ) + + def __call__(self) -> tuple[ColumnInfo, ...]: + self.reads += 1 + statements: tuple[str, ...] = tuple(filter(None, (self.statement_during_read,))) + statement: str + for statement in statements: + self.cache.observe_statement(sql=statement) + return self.columns + + +def read_orders(*, cache: RelationMetadataCache, read: CountingColumnsRead) -> None: + """Look up the ``orders`` columns once through the cache.""" + + _ = cache.lookup( + kind=RELATION_COLUMNS_LOOKUP, database="analytics", schema="marts", name="orders", read=read + ) diff --git a/tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/test_relation_metadata_cache.py b/tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/test_relation_metadata_cache.py new file mode 100644 index 0000000000..99aa2b0f47 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/classes/relation_metadata_cache/test_relation_metadata_cache.py @@ -0,0 +1,231 @@ +"""Identical metadata lookups are served once until SQL may have changed the relation.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest + +from sqlbuild.adapter.contract.models import ColumnInfo +from sqlbuild.adapter.relations.classes.relation_metadata_cache import RelationMetadataCache +from sqlbuild.adapter.relations.main.cached_relation_columns import cached_relation_columns +from sqlbuild.adapter.relations.main.cached_relation_exists import cached_relation_exists +from sqlbuild.adapter.relations.main.relation_metadata_cache_scope import ( + relation_metadata_cache_scope, +) +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from tests.unit.src.sqlbuild.adapter.relations.classes.relation_metadata_cache._test_types import ( + CacheInvalidationTestCase, + LiveInvalidationTestCase, + OverlappingReadTestCase, +) +from tests.unit.src.sqlbuild.adapter.relations.classes.relation_metadata_cache.helpers import ( + CountingColumnsRead, + read_orders, +) + + +@pytest.mark.parametrize( + "test_case", + [ + CacheInvalidationTestCase( + description="no statement serves the cached read", + transactional_ddl=False, + statements_between=(), + expected_reads=1, + ), + CacheInvalidationTestCase( + description="data statements keep the entry", + transactional_ddl=False, + statements_between=( + "INSERT INTO analytics.marts.orders SELECT * FROM analytics.marts.orders__delta", + "SELECT COUNT(*) FROM analytics.marts.orders", + ), + expected_reads=1, + ), + CacheInvalidationTestCase( + description="DDL on another relation keeps the entry", + transactional_ddl=False, + statements_between=( + "CREATE OR REPLACE TABLE analytics.marts.customers AS SELECT 1 AS id", + ), + expected_reads=1, + ), + CacheInvalidationTestCase( + description="replace of the relation refreshes the entry", + transactional_ddl=False, + statements_between=( + "CREATE OR REPLACE TABLE ANALYTICS.MARTS.ORDERS AS SELECT 1 AS id", + ), + expected_reads=2, + ), + CacheInvalidationTestCase( + description="alter of the relation refreshes the entry", + transactional_ddl=False, + statements_between=('ALTER TABLE "ANALYTICS"."MARTS"."ORDERS" ADD COLUMN "NOTE" TEXT',), + expected_reads=2, + ), + CacheInvalidationTestCase( + description="swap into the relation refreshes the entry", + transactional_ddl=False, + statements_between=("ALTER TABLE marts.orders__stage SWAP WITH marts.orders",), + expected_reads=2, + ), + CacheInvalidationTestCase( + description="unknown statement refreshes everything", + transactional_ddl=False, + statements_between=("CALL marts.rebuild()",), + expected_reads=2, + ), + CacheInvalidationTestCase( + description="commit without transactional DDL keeps the entry", + transactional_ddl=False, + statements_between=("COMMIT",), + expected_reads=1, + ), + ], + ids=lambda case: case.description, +) +def test_given_statements_between_lookups_when_reading_columns_then_rereads_only_after_changes( + test_case: CacheInvalidationTestCase, +) -> None: + cache: RelationMetadataCache = RelationMetadataCache( + transactional_ddl=test_case.transactional_ddl + ) + read: CountingColumnsRead = CountingColumnsRead(cache=cache) + + read_orders(cache=cache, read=read) + statement: str + for statement in test_case.statements_between: + cache.observe_statement(sql=statement) + read_orders(cache=cache, read=read) + + assert read.reads == test_case.expected_reads + + +@pytest.mark.parametrize( + "test_case", + [ + CacheInvalidationTestCase( + description=description, + transactional_ddl=transactional_ddl, + statements_between=(statement,), + expected_reads=expected_reads, + ) + for description, transactional_ddl, statement, expected_reads in ( + ("commit replays transactional DDL", True, "COMMIT", 2), + ("rollback replays transactional DDL", True, "ROLLBACK", 2), + ("commit is inert for autocommitted DDL", False, "COMMIT", 1), + ) + ], + ids=lambda case: case.description, +) +def test_given_transaction_end_after_ddl_when_reading_columns_then_follows_ddl_visibility( + test_case: CacheInvalidationTestCase, +) -> None: + cache: RelationMetadataCache = RelationMetadataCache( + transactional_ddl=test_case.transactional_ddl + ) + read: CountingColumnsRead = CountingColumnsRead(cache=cache) + + cache.observe_statement(sql="ALTER TABLE analytics.marts.orders ADD COLUMN note TEXT") + read_orders(cache=cache, read=read) + statement: str + for statement in test_case.statements_between: + cache.observe_statement(sql=statement) + read_orders(cache=cache, read=read) + + assert read.reads == test_case.expected_reads + + +@pytest.mark.parametrize( + "test_case", + [ + OverlappingReadTestCase( + description="DDL during the read is not cached over", + statement_during_read="DROP TABLE IF EXISTS analytics.marts.orders", + expected_reads=2, + ), + OverlappingReadTestCase( + description="unrelated DDL during the read still caches", + statement_during_read="DROP TABLE IF EXISTS analytics.marts.customers", + expected_reads=1, + ), + ], + ids=lambda case: case.description, +) +def test_given_statement_during_read_when_reading_again_then_never_serves_a_pre_ddl_read( + test_case: OverlappingReadTestCase, +) -> None: + cache: RelationMetadataCache = RelationMetadataCache(transactional_ddl=False) + overlapping: CountingColumnsRead = CountingColumnsRead( + cache=cache, statement_during_read=test_case.statement_during_read + ) + read_orders(cache=cache, read=overlapping) + following: CountingColumnsRead = CountingColumnsRead(cache=cache) + + read_orders(cache=cache, read=following) + + assert overlapping.reads + following.reads == test_case.expected_reads + + +@pytest.mark.parametrize( + "test_case", + [ + LiveInvalidationTestCase( + description="added column is visible on the next lookup", + ddl="ALTER TABLE main.orders ADD COLUMN note VARCHAR", + expected_first_columns=("order_id",), + expected_second_columns=("order_id", "note"), + expected_second_exists=True, + ), + LiveInvalidationTestCase( + description="replaced table is read again", + ddl="CREATE OR REPLACE TABLE main.orders AS SELECT 1 AS order_id, 2 AS customer_id", + expected_first_columns=("order_id",), + expected_second_columns=("order_id", "customer_id"), + expected_second_exists=True, + ), + LiveInvalidationTestCase( + description="dropped table is reported missing", + ddl="DROP TABLE main.orders", + expected_first_columns=("order_id",), + expected_second_columns=(), + expected_second_exists=False, + ), + ], + ids=lambda case: case.description, +) +def test_given_adapter_ddl_when_looking_up_again_then_returns_fresh_metadata( + test_case: LiveInvalidationTestCase, tmp_path: Path +) -> None: + adapter: DuckDbAdapter = DuckDbAdapter() + connection: Any = adapter.connect({"database": str(tmp_path / "warehouse.duckdb")}) + _ = adapter.execute( + connection=connection, sql="CREATE TABLE main.orders AS SELECT 1 AS order_id" + ) + + with relation_metadata_cache_scope(adapter=adapter): + first: tuple[ColumnInfo, ...] = cached_relation_columns( + adapter=adapter, connection=connection, database=None, schema="main", name="orders" + ) + _ = cached_relation_exists( + adapter=adapter, connection=connection, database=None, schema="main", name="orders" + ) + _ = adapter.execute(connection=connection, sql=test_case.ddl) + second: tuple[ColumnInfo, ...] = cached_relation_columns( + adapter=adapter, connection=connection, database=None, schema="main", name="orders" + ) + second_exists: bool = cached_relation_exists( + adapter=adapter, connection=connection, database=None, schema="main", name="orders" + ) + adapter.close(connection) + + assert tuple(column.name for column in first) == test_case.expected_first_columns + assert tuple(column.name for column in second) == test_case.expected_second_columns + assert second_exists is test_case.expected_second_exists + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) From d7de75051a8a534a5b28ff98744dd9b9dce29506 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 15:21:01 +0100 Subject: [PATCH 3/9] test(executor): guard Snowflake build metadata round trips offline Compile, plan, and build a table, a view, a merge incremental, and a three-day microbatch model twice through the real build pipeline against a simulated Snowflake whose catalog follows the DDL the build issues. Each build must send no INFORMATION_SCHEMA query, must never repeat a metadata read without DDL on that relation in between, and must read columns exactly once per relation version, with one and four workers. A delta that gains a column must reach the merge through a fresh read of the altered target. Refs CHI-455 --- .../build/snowflake_metadata/__init__.py | 0 .../build/snowflake_metadata/_test_types.py | 23 + .../build/snowflake_metadata/helpers.py | 515 ++++++++++++++++++ .../test_build_metadata_budget.py | 130 +++++ .../adapters/snowflake/inspection/helpers.py | 13 + 5 files changed, 681 insertions(+) create mode 100644 tests/integration/src/sqlbuild/executor/build/snowflake_metadata/__init__.py create mode 100644 tests/integration/src/sqlbuild/executor/build/snowflake_metadata/_test_types.py create mode 100644 tests/integration/src/sqlbuild/executor/build/snowflake_metadata/helpers.py create mode 100644 tests/integration/src/sqlbuild/executor/build/snowflake_metadata/test_build_metadata_budget.py diff --git a/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/__init__.py b/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/_test_types.py b/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/_test_types.py new file mode 100644 index 0000000000..c0fc8e91cc --- /dev/null +++ b/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/_test_types.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class BuildMetadataBudgetTestCase: + """Two consecutive offline Snowflake builds and the column reads each may issue.""" + + description: str + max_concurrency: int + expected_first_column_reads: dict[str, int] + expected_second_column_reads: dict[str, int] + + +@dataclass(frozen=True) +class SchemaEvolutionTestCase: + """A merge model whose delta gains a column between builds.""" + + description: str + added_column: str + expected_target_reads: int + expected_merge_fragment: str diff --git a/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/helpers.py b/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/helpers.py new file mode 100644 index 0000000000..e0e158de04 --- /dev/null +++ b/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/helpers.py @@ -0,0 +1,515 @@ +"""Offline Snowflake build harness whose catalog follows the DDL the build issues. + +Data statements are accepted and ignored; only relation existence and columns are simulated, so +metadata round trips during a real compile, plan, and build can be counted without a warehouse. +""" + +from __future__ import annotations + +import re +from collections import Counter +from collections.abc import Callable +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any + +from sqlbuild.adapter.contract.types import TablePromotionMode +from sqlbuild.compiler.discovery.main.discover import discover_project_inputs +from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs +from sqlbuild.compiler.pipeline.main.compile import run_compile_pipeline +from sqlbuild.compiler.pipeline.models import CompilePipelineOptions, CompilePipelineResult +from sqlbuild.executor.build.models import BuildExecutionResult, BuildRuntimeParams +from sqlbuild.executor.pipeline.main.run import run_build_pipeline +from sqlbuild.executor.scheduling.types import ExecutionStatus +from sqlbuild.runtime.contracts.models import ConnectionHooks +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + FakeColumn, + FakeRelation, + OfflineSnowflakeAdapter, + RecordingSnowflakeWarehouse, +) + +_NAME: str = r'(?P<{group}>(?:"[^"]*"|[A-Za-z0-9_$]+)(?:\.(?:"[^"]*"|[A-Za-z0-9_$]+)){{0,2}})' +_CREATE_AS_PATTERN: re.Pattern[str] = re.compile( + r"^\s*CREATE\s+(?:OR\s+REPLACE\s+)?(?:TRANSIENT\s+|TEMPORARY\s+)?(?PTABLE|VIEW)\s+" + + _NAME.format(group="name") + + r"\s+AS\s+(?P.*)$", + re.IGNORECASE | re.DOTALL, +) +_CREATE_DEFINED_PATTERN: re.Pattern[str] = re.compile( + r"^\s*CREATE\s+(?:OR\s+REPLACE\s+)?(?:TRANSIENT\s+)?TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?" + + _NAME.format(group="name") + + r"\s*\(", + re.IGNORECASE | re.DOTALL, +) +_DROP_PATTERN: re.Pattern[str] = re.compile( + r"^\s*DROP\s+(?:TABLE|VIEW)\s+(?:IF\s+EXISTS\s+)?" + _NAME.format(group="name") + r"\s*$", + re.IGNORECASE, +) +_RENAME_PATTERN: re.Pattern[str] = re.compile( + r"^\s*ALTER\s+(?:TABLE|VIEW)\s+" + + _NAME.format(group="name") + + r"\s+RENAME\s+TO\s+" + + _NAME.format(group="other") + + r"\s*$", + re.IGNORECASE, +) +_SWAP_PATTERN: re.Pattern[str] = re.compile( + r"^\s*ALTER\s+TABLE\s+" + + _NAME.format(group="name") + + r"\s+SWAP\s+WITH\s+" + + _NAME.format(group="other") + + r"\s*$", + re.IGNORECASE, +) +_ADD_COLUMN_PATTERN: re.Pattern[str] = re.compile( + r"^\s*ALTER\s+TABLE\s+" + + _NAME.format(group="name") + + r'\s+ADD\s+COLUMN\s+"?(?P[A-Za-z0-9_]+)"?\s+(?P.+?)\s*$', + re.IGNORECASE, +) +_CREATE_SCHEMA_PATTERN: re.Pattern[str] = re.compile( + r"^\s*CREATE\s+SCHEMA\s+IF\s+NOT\s+EXISTS\s+" + _NAME.format(group="name") + r"\s*$", + re.IGNORECASE, +) +_DESCRIBE_QUERY_PATTERN: re.Pattern[str] = re.compile( + r"^SELECT \* FROM \(SELECT (?P.+?) FROM \S+\) AS __describe_source LIMIT 0$", + re.DOTALL, +) +_RUNTIME_BOUNDS_PATTERN: re.Pattern[str] = re.compile(r"^SELECT MIN\([^)]*\), MAX\([^)]*\) FROM ") +_COPY_SOURCE_PATTERN: re.Pattern[str] = re.compile( + r"^\s*SELECT\s+\*\s+FROM\s+" + _NAME.format(group="source") + r"\s*$", re.IGNORECASE +) +_DEFAULT_COLUMNS: tuple[FakeColumn, ...] = (FakeColumn(name="ID", data_type="NUMBER"),) +_STATUS_DESCRIPTION: tuple[tuple[str], ...] = (("status",),) + + +type _Identity = tuple[str, str, str] +type _DdlAction = Callable[[re.Match[str]], frozenset[_Identity]] + + +@dataclass +class SimulatedSnowflakeWarehouse(RecordingSnowflakeWarehouse): + """Recording warehouse whose relations follow CREATE, DROP, RENAME, SWAP, and ADD COLUMN.""" + + model_columns: dict[str, tuple[FakeColumn, ...]] = field(default_factory=dict) + ddl_events: list[tuple[int, frozenset[_Identity]]] = field(default_factory=list) + + def reset(self) -> None: + """Forget recorded statements and simulated DDL events, keeping the catalog.""" + + super().reset() + with self._lock: + self.ddl_events.clear() + + def set_model_columns(self, *, model: str, columns: tuple[FakeColumn, ...]) -> None: + """Change the columns a model's CREATE ... AS statements produce.""" + + with self._lock: + self.model_columns[model.upper()] = columns + + def _answer(self, *, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[Any, ...]], Any]: + routes: tuple[tuple[re.Pattern[str], _DdlAction], ...] = ( + (_CREATE_AS_PATTERN, self._create_as), + (_CREATE_DEFINED_PATTERN, self._create_defined), + (_DROP_PATTERN, self._drop), + (_RENAME_PATTERN, self._rename), + (_SWAP_PATTERN, self._swap), + (_ADD_COLUMN_PATTERN, self._add_column), + (_CREATE_SCHEMA_PATTERN, self._create_schema), + ) + pattern: re.Pattern[str] + apply: _DdlAction + for pattern, apply in routes: + match: re.Match[str] + for match in filter(None, (pattern.match(sql),)): + with self._lock: + changed: frozenset[_Identity] = apply(match) + self.ddl_events.append((len(self.attempted_sql) - 1, changed)) + return [], _STATUS_DESCRIPTION + described: tuple[re.Match[str], ...] = tuple( + filter(None, (_DESCRIBE_QUERY_PATTERN.match(sql),)) + ) + description: re.Match[str] + for description in described: + return [], tuple( + (column.strip().strip('"'),) for column in description.group("columns").split(",") + ) + bounds: tuple[re.Match[str], ...] = tuple( + filter(None, (_RUNTIME_BOUNDS_PATTERN.match(sql),)) + ) + for _ in bounds: + return [self.cursor_values], (("min",), ("max",)) + return super()._answer(sql=sql, params=params) + + def _create_as(self, match: re.Match[str]) -> frozenset[_Identity]: + identity: _Identity = self._parts(match.group("name")) + copied: tuple[re.Match[str], ...] = tuple( + filter(None, (_COPY_SOURCE_PATTERN.match(match.group("query")),)) + ) + sources: frozenset[_Identity] = frozenset( + self._parts(copy.group("source")) for copy in copied + ) + source_columns: tuple[tuple[FakeColumn, ...], ...] = tuple( + fake.columns for fake in filter(lambda fake: _identity(fake) in sources, self.relations) + ) + columns: tuple[FakeColumn, ...] = next( + iter(source_columns), self._model_columns_for(name=identity[2]) + ) + table_type: str = ("BASE TABLE", "VIEW")[match.group("kind").upper() == "VIEW"] + self._put( + relation=FakeRelation( + database=identity[0], + schema=identity[1], + name=identity[2], + table_type=table_type, + is_transient=True, + columns=columns, + ) + ) + return frozenset({identity}) + + def _create_defined(self, match: re.Match[str]) -> frozenset[_Identity]: + identity: _Identity = self._parts(match.group("name")) + existing: tuple[FakeRelation, ...] = tuple( + filter(lambda fake: _identity(fake) == identity, self.relations) + ) + self._put( + relation=next( + iter(existing), + FakeRelation( + database=identity[0], + schema=identity[1], + name=identity[2], + is_transient=True, + columns=_DEFAULT_COLUMNS, + ), + ) + ) + return frozenset({identity}) + + def _drop(self, match: re.Match[str]) -> frozenset[_Identity]: + identity: _Identity = self._parts(match.group("name")) + self.replace_relations( + relations=tuple(filter(lambda fake: _identity(fake) != identity, self.relations)) + ) + return frozenset({identity}) + + def _rename(self, match: re.Match[str]) -> frozenset[_Identity]: + origin: _Identity = self._parts(match.group("name")) + destination: _Identity = self._parts(match.group("other")) + self.replace_relations( + relations=tuple( + (fake, _moved(fake=fake, identity=destination))[_identity(fake) == origin] + for fake in self.relations + ) + ) + return frozenset({origin, destination}) + + def _swap(self, match: re.Match[str]) -> frozenset[_Identity]: + left: _Identity = self._parts(match.group("name")) + right: _Identity = self._parts(match.group("other")) + targets: dict[_Identity, _Identity] = {left: right, right: left} + self.replace_relations( + relations=tuple( + _moved(fake=fake, identity=targets.get(_identity(fake), _identity(fake))) + for fake in self.relations + ) + ) + return frozenset({left, right}) + + def _add_column(self, match: re.Match[str]) -> frozenset[_Identity]: + identity: _Identity = self._parts(match.group("name")) + added: FakeColumn = FakeColumn(name=match.group("column").upper(), data_type="TEXT") + self.replace_relations( + relations=tuple( + (fake, replace(fake, columns=(*fake.columns, added)))[_identity(fake) == identity] + for fake in self.relations + ) + ) + return frozenset({identity}) + + def _create_schema(self, match: re.Match[str]) -> frozenset[_Identity]: + parts: tuple[str, ...] = tuple( + part.strip('"').upper() for part in match.group("name").split(".") + ) + scope: tuple[str, ...] = (self.current_database, *parts)[-2:] + self.extra_schemas = self.extra_schemas | {(scope[0], scope[1])} + return frozenset({(scope[0], scope[1], "")}) + + def _put(self, *, relation: FakeRelation) -> None: + identity: _Identity = _identity(relation) + self.replace_relations( + relations=(*filter(lambda fake: _identity(fake) != identity, self.relations), relation) + ) + + def _model_columns_for(self, *, name: str) -> tuple[FakeColumn, ...]: + candidates: tuple[str, ...] = tuple( + filter(lambda model: _model_owns(model=model, name=name), self.model_columns) + ) + models: list[str] = sorted(candidates, key=_name_length, reverse=True) + return next((self.model_columns[model] for model in models), _DEFAULT_COLUMNS) + + def _parts(self, name: str) -> _Identity: + parts: tuple[str, ...] = tuple(part.strip('"').upper() for part in name.split(".")) + full: tuple[str, ...] = (self.current_database, self.current_database, *parts)[-3:] + return full[0], full[1], full[2] + + +def _name_length(name: str) -> int: + return len(name) + + +def _model_owns(*, model: str, name: str) -> bool: + return name == model or name.startswith(f"{model}__") + + +def _identity(fake: FakeRelation) -> _Identity: + return fake.database, fake.schema, fake.name + + +def _moved(*, fake: FakeRelation, identity: _Identity) -> FakeRelation: + return replace(fake, database=identity[0], schema=identity[1], name=identity[2]) + + +@dataclass(frozen=True) +class OfflineBuild: + """One compiled, planned, and executed offline Snowflake build.""" + + result: BuildExecutionResult + statements: tuple[str, ...] + ddl_events: tuple[tuple[int, frozenset[_Identity]], ...] + + +def build_offline_snowflake_project( + *, project_dir: Path, warehouse: SimulatedSnowflakeWarehouse, max_concurrency: int +) -> OfflineBuild: + """Compile, plan, and build ``project_dir`` through the build pipeline, offline.""" + + adapter: OfflineSnowflakeAdapter = OfflineSnowflakeAdapter(warehouse=warehouse) + inputs: DiscoveredProjectInputs = discover_project_inputs(project_dir=project_dir) + config: dict[str, object] = dict(inputs.project_config.connection) + planned: CompilePipelineResult = run_compile_pipeline( + discovered_inputs=inputs, + adapter=adapter, + options=CompilePipelineOptions( + no_sql_validation=True, no_cache=True, connection_config=config + ), + hooks=ConnectionHooks(), + ) + warehouse.reset() + result: BuildExecutionResult = run_build_pipeline( + plan=planned.plan_output, + connection_config=config, + adapter=adapter, + settings=inputs.project_config.settings, + runtime=BuildRuntimeParams( + run_id="offline_build", + runtime_dir=project_dir / "target", + promotion_mode=TablePromotionMode(adapter.default_table_promotion_mode()), + max_concurrency=max_concurrency, + ), + ) + return OfflineBuild( + result=result, + statements=tuple(warehouse.attempted_sql), + ddl_events=tuple(warehouse.ddl_events), + ) + + +def information_schema_statements(build: OfflineBuild) -> tuple[str, ...]: + """Return build statements that read INFORMATION_SCHEMA, which needs a running warehouse.""" + + return tuple( + filter( + lambda statement: _INFORMATION_SCHEMA_PATTERN.search(statement) is not None, + build.statements, + ) + ) + + +def column_reads_by_relation(build: OfflineBuild) -> dict[str, int]: + """Count SHOW COLUMNS reads per relation the build inspected.""" + + reads: tuple[re.Match[str], ...] = tuple( + filter(None, (_SHOW_COLUMNS_RELATION.match(statement) for statement in build.statements)) + ) + return dict(Counter(read.group(1).replace('"', "") for read in reads)) + + +def repeated_metadata_reads(build: OfflineBuild) -> tuple[str, ...]: + """Return metadata statements re-issued with no DDL on their relation since the last issue.""" + + last_read: dict[str, int] = {} + repeated: list[str] = [] + index: int + statement: str + for index, statement in enumerate(build.statements): + identities: tuple[_Identity, ...] = _metadata_identities(statement) + identity: _Identity + for identity in identities: + previous: int = last_read.get(statement, -1) + changed: bool = any( + previous < event_index < index and identity in names + for event_index, names in build.ddl_events + ) + repeated.extend(statement for _ in range(int(previous >= 0 and not changed))) + last_read[statement] = index + return tuple(repeated) + + +def _metadata_identities(statement: str) -> tuple[_Identity, ...]: + identities: list[_Identity] = [] + pattern: re.Pattern[str] + to_identity: Callable[[re.Match[str]], _Identity] + for pattern, to_identity in _METADATA_STATEMENTS: + identities.extend(map(to_identity, filter(None, (pattern.match(statement),)))) + return tuple(identities) + + +def _show_columns_identity(match: re.Match[str]) -> _Identity: + parts: list[str] = match.group("relation").replace('"', "").split(".") + return parts[0], parts[1], parts[2] + + +def _show_relation_identity(match: re.Match[str]) -> _Identity: + return match.group("database"), match.group("schema"), match.group("name") + + +def _show_schema_identity(match: re.Match[str]) -> _Identity: + return match.group("database"), match.group("schema"), "" + + +_INFORMATION_SCHEMA_PATTERN: re.Pattern[str] = re.compile(r"information_schema", re.IGNORECASE) +_SHOW_COLUMNS_RELATION: re.Pattern[str] = re.compile(r"^SHOW COLUMNS IN (?:TABLE|VIEW) (\S+)$") +_METADATA_STATEMENTS: tuple[tuple[re.Pattern[str], Callable[[re.Match[str]], _Identity]], ...] = ( + ( + re.compile(r"^SHOW COLUMNS IN (?:TABLE|VIEW) (?P\S+)$"), + _show_columns_identity, + ), + ( + re.compile( + r"^SHOW TERSE (?:TABLES|VIEWS) LIKE '(?P[^']*)' " + r'IN SCHEMA "(?P[^"]*)"\."(?P[^"]*)"$' + ), + _show_relation_identity, + ), + ( + re.compile(r"^SHOW SCHEMAS LIKE '(?P[^']*)' IN DATABASE \"(?P[^\"]*)\"$"), + _show_schema_identity, + ), +) + + +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", character_maximum_length=None), + FakeColumn(name="AMOUNT", data_type="NUMBER", numeric_precision=10, numeric_scale=2), +) +ORDER_MODELS: tuple[str, ...] = ("ORDERS_TABLE", "ORDERS_VIEW", "ORDERS_MERGE", "ORDERS_DAILY") +_PROJECT_FILES: dict[str, str] = { + "sqlbuild_project.toml": ( + 'name = "orders_platform"\n' + 'adapter = "snowflake"\n\n' + "[connection]\n" + 'account = "example"\n' + 'user = "example"\n' + 'database = "analytics"\n' + 'schema = "marts"\n\n' + "[connection.session_parameters]\n" + "QUOTED_IDENTIFIERS_IGNORE_CASE = true\n\n" + "[defaults]\n" + 'database = "analytics"\n' + 'materialized = "table"\n\n' + "[path_defaults.marts]\n" + 'schema = "marts"\n' + ), + "sources/raw.yml": ( + "sources:\n" + " - name: raw_orders\n" + " database: analytics\n" + " schema: raw\n" + " table: orders\n" + " columns:\n" + " - name: order_id\n" + " type: NUMBER(38, 0)\n" + " - name: ordered_at\n" + " type: TIMESTAMP_NTZ\n" + " - name: amount\n" + " type: NUMBER(10, 2)\n" + ), + "models/marts/orders_table.sql": ( + "MODEL (\n materialized table,\n);\n\n" + 'SELECT order_id, ordered_at, amount FROM __source("raw_orders")\n' + ), + "models/marts/orders_view.sql": ( + 'MODEL (\n materialized view,\n);\n\nSELECT order_id, amount FROM __ref("orders_table")\n' + ), + "models/marts/orders_merge.sql": ( + "MODEL (\n" + " materialized incremental,\n" + " incremental_strategy merge,\n" + " unique_key [order_id],\n" + " cursor ordered_at,\n" + " cursor_type timestamp,\n" + " cursor_grain day,\n" + " cursor_inputs (\n raw_orders ordered_at,\n ),\n" + " on_schema_change append_new_columns,\n" + ");\n\n" + 'SELECT order_id, ordered_at, amount FROM __source("raw_orders")\n' + ), + "models/marts/orders_daily.sql": ( + "MODEL (\n" + " materialized incremental,\n" + " incremental_strategy delete_insert,\n" + " incremental_mode microbatch,\n" + " microbatch_strategy watermark,\n" + " cursor ordered_at,\n" + " cursor_type timestamp,\n" + " cursor_grain day,\n" + " cursor_start '2026-01-01',\n" + " cursor_end '2026-01-04',\n" + " cursor_watermark_mode all,\n" + " cursor_inputs (\n raw_orders (column ordered_at, roles [filter, watermark]),\n ),\n" + " batch_size 1d,\n" + ");\n\n" + 'SELECT order_id, ordered_at, amount FROM __source("raw_orders")\n' + ), +} + + +def write_orders_project(*, project_dir: Path) -> None: + """Write a table, a view, a merge incremental, and a three-day microbatch model.""" + + relative: str + content: str + for relative, content in _PROJECT_FILES.items(): + path: Path = project_dir / relative + path.parent.mkdir(parents=True, exist_ok=True) + _ = path.write_text(content, encoding="utf-8") + + +def orders_warehouse() -> SimulatedSnowflakeWarehouse: + """Return a simulated warehouse holding the raw orders source.""" + + warehouse: SimulatedSnowflakeWarehouse = SimulatedSnowflakeWarehouse( + relations=( + FakeRelation(database="ANALYTICS", schema="RAW", name="ORDERS", columns=ORDER_COLUMNS), + ), + cursor_values=("2026-01-01 00:00:00.000", "2026-01-03 12:00:00.000"), + ) + model: str + for model in ORDER_MODELS: + warehouse.set_model_columns(model=model, columns=ORDER_COLUMNS) + return warehouse + + +def failed_models(build: OfflineBuild) -> tuple[str, ...]: + """Return models that did not build successfully.""" + + return tuple( + result.model_name + for result in filter( + lambda result: result.status != ExecutionStatus.SUCCESS, build.result.model_results + ) + ) diff --git a/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/test_build_metadata_budget.py b/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/test_build_metadata_budget.py new file mode 100644 index 0000000000..d71f563a96 --- /dev/null +++ b/tests/integration/src/sqlbuild/executor/build/snowflake_metadata/test_build_metadata_budget.py @@ -0,0 +1,130 @@ +"""Build-time metadata budget for a Snowflake project, compiled, planned, and built offline. + +Only the network is replaced: a simulated warehouse answers SHOW and INFORMATION_SCHEMA reads +from a catalog that follows the DDL the build itself issues. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.integration.src.sqlbuild.executor.build.snowflake_metadata._test_types import ( + BuildMetadataBudgetTestCase, + SchemaEvolutionTestCase, +) +from tests.integration.src.sqlbuild.executor.build.snowflake_metadata.helpers import ( + ORDER_COLUMNS, + OfflineBuild, + SimulatedSnowflakeWarehouse, + build_offline_snowflake_project, + column_reads_by_relation, + failed_models, + information_schema_statements, + orders_warehouse, + repeated_metadata_reads, + write_orders_project, +) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import FakeColumn + +_FIRST_BUILD_COLUMN_READS: dict[str, int] = { + "ANALYTICS.MARTS.ORDERS_TABLE__STAGING": 1, + "ANALYTICS.MARTS.ORDERS_MERGE__STAGING": 1, + "ANALYTICS.MARTS.ORDERS_DAILY__DELTA": 3, + "ANALYTICS.MARTS.__SQB_REBUILD__ORDERS_DAILY": 1, +} +_SECOND_BUILD_COLUMN_READS: dict[str, int] = { + "ANALYTICS.MARTS.ORDERS_TABLE__STAGING": 1, + "ANALYTICS.MARTS.ORDERS_MERGE__DELTA": 1, + "ANALYTICS.MARTS.ORDERS_MERGE": 1, + "ANALYTICS.MARTS.ORDERS_DAILY__DELTA": 2, + "ANALYTICS.MARTS.ORDERS_DAILY": 1, +} + + +@pytest.mark.parametrize( + "test_case", + [ + BuildMetadataBudgetTestCase( + description="sequential workers", + max_concurrency=1, + expected_first_column_reads=_FIRST_BUILD_COLUMN_READS, + expected_second_column_reads=_SECOND_BUILD_COLUMN_READS, + ), + BuildMetadataBudgetTestCase( + description="concurrent workers", + max_concurrency=4, + expected_first_column_reads=_FIRST_BUILD_COLUMN_READS, + expected_second_column_reads=_SECOND_BUILD_COLUMN_READS, + ), + ], + ids=lambda case: case.description, +) +def test_given_table_view_merge_and_microbatch_when_building_then_reads_each_relation_version_once( + test_case: BuildMetadataBudgetTestCase, tmp_path: Path +) -> None: + write_orders_project(project_dir=tmp_path) + warehouse: SimulatedSnowflakeWarehouse = orders_warehouse() + + first: OfflineBuild = build_offline_snowflake_project( + project_dir=tmp_path, warehouse=warehouse, max_concurrency=test_case.max_concurrency + ) + second: OfflineBuild = build_offline_snowflake_project( + project_dir=tmp_path, warehouse=warehouse, max_concurrency=test_case.max_concurrency + ) + + assert failed_models(first) == () + assert failed_models(second) == () + assert information_schema_statements(first) == () + assert information_schema_statements(second) == () + assert repeated_metadata_reads(first) == () + assert repeated_metadata_reads(second) == () + assert column_reads_by_relation(first) == test_case.expected_first_column_reads + assert column_reads_by_relation(second) == test_case.expected_second_column_reads + + +@pytest.mark.parametrize( + "test_case", + [ + SchemaEvolutionTestCase( + description="appended delta column reaches the merge", + added_column="NOTE", + expected_target_reads=2, + expected_merge_fragment='INSERT ("ORDER_ID", "ORDERED_AT", "AMOUNT", "NOTE")', + ) + ], + ids=lambda case: case.description, +) +def test_given_delta_gains_column_when_building_then_merge_reads_altered_target_again( + test_case: SchemaEvolutionTestCase, tmp_path: Path +) -> None: + write_orders_project(project_dir=tmp_path) + warehouse: SimulatedSnowflakeWarehouse = orders_warehouse() + _ = build_offline_snowflake_project( + project_dir=tmp_path, warehouse=warehouse, max_concurrency=1 + ) + warehouse.set_model_columns( + model="ORDERS_MERGE", + columns=(*ORDER_COLUMNS, FakeColumn(name=test_case.added_column, data_type="TEXT")), + ) + + evolved: OfflineBuild = build_offline_snowflake_project( + project_dir=tmp_path, warehouse=warehouse, max_concurrency=1 + ) + + merges: tuple[str, ...] = tuple( + filter(lambda statement: statement.startswith("MERGE INTO"), evolved.statements) + ) + assert failed_models(evolved) == () + assert repeated_metadata_reads(evolved) == () + assert ( + column_reads_by_relation(evolved)["ANALYTICS.MARTS.ORDERS_MERGE"] + == test_case.expected_target_reads + ) + assert len(merges) == 1 + assert test_case.expected_merge_fragment in merges[0] + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py index 5fcec5acd7..0a5689440d 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -248,6 +248,19 @@ def queries_of_kind(self, kind: str) -> tuple[RecordedQuery, ...]: with self._lock: return tuple(filter(lambda query: query.kind == kind, self.queries)) + def replace_relations(self, *, relations: tuple[FakeRelation, ...]) -> None: + """Replace the catalog that SHOW and INFORMATION_SCHEMA reads answer from.""" + + self.relations = relations + self._database.execute("DELETE FROM fake_tables") + self._database.execute("DELETE FROM fake_columns") + self._load(table="fake_tables", rows=[_table_row(relation) for relation in relations]) + column_rows: list[dict[str, object]] = [] + relation: FakeRelation + for relation in relations: + column_rows.extend(_column_rows(relation)) + self._load(table="fake_columns", rows=column_rows) + def reset(self) -> None: """Forget recorded statements and concurrency peaks.""" From 7379701b3a87f020bea3769b5115890dd73f4171 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 15:23:09 +0100 Subject: [PATCH 4/9] docs(snowflake): describe build-time SHOW metadata reads and reuse Refs CHI-455 --- .../docs/concepts/adapters/snowflake.md | 16 ++++++++++++++++ .../docs/docs/concepts/adapters/snowflake.mdx | 7 +++++++ 2 files changed, 23 insertions(+) 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 6672418ab2..4789033457 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 @@ -6,6 +6,15 @@ Online: https://sqlbuild.com/docs/concepts/adapters/snowflake/ +## Contents + +- Connection config +- Session initialization +- Planning metadata +- Build metadata +- Shared connections across targets +- Cost estimates + Snowflake requires the optional `snowflake-connector-python` dependency: ```bash @@ -54,6 +63,13 @@ and need a running warehouse: Run `sqb plan -v` to see each metadata query with its elapsed time and row count. +## Build metadata + +During a build, checks for whether one relation exists use `SHOW TERSE TABLES` and +`SHOW TERSE VIEWS`, and reads of one relation's columns use `SHOW COLUMNS`, so they also don't +need a running warehouse. Each build reuses these answers until it changes that relation, for +example by replacing a delta table or adding a column, and then reads it again. + ## Shared connections across targets Use multiple targets to reuse one Snowflake connection while selecting different authoritative diff --git a/website/src/content/docs/docs/concepts/adapters/snowflake.mdx b/website/src/content/docs/docs/concepts/adapters/snowflake.mdx index ff8a6787ae..389c7149cb 100644 --- a/website/src/content/docs/docs/concepts/adapters/snowflake.mdx +++ b/website/src/content/docs/docs/concepts/adapters/snowflake.mdx @@ -51,6 +51,13 @@ and need a running warehouse: Run `sqb plan -v` to see each metadata query with its elapsed time and row count. +## Build metadata + +During a build, checks for whether one relation exists use `SHOW TERSE TABLES` and +`SHOW TERSE VIEWS`, and reads of one relation's columns use `SHOW COLUMNS`, so they also don't +need a running warehouse. Each build reuses these answers until it changes that relation, for +example by replacing a delta table or adding a column, and then reads it again. + ## Shared connections across targets Use multiple targets to reuse one Snowflake connection while selecting different authoritative From 2a03ba3da9746f7a95045f7db81cec9b75d9c186 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:08:16 +0100 Subject: [PATCH 5/9] fix(executor): invalidate all cached metadata when DDL names are ambiguous The statement classifier now lexes literals, quoted identifiers, comments, and dollar-quoted bodies before classifying, so a comment marker or separator inside a literal no longer hides a second statement. It names a relation only when the name is followed by a clause it expects (DROP of several relations, ALTER TABLE ONLY, unquoted non-ASCII names, and unexpected clauses fall back to evicting everything), classifies under standard and backslash-escaped literal rules and evicts everything when they disagree, and keys quoted dotted names the same way lookups do. Refs CHI-455 --- .../adapter/relations/_helpers/ddl_effects.py | 182 ++++++++++++++---- src/sqlbuild/adapter/relations/constants.py | 1 + .../test_statement_metadata_effect.py | 35 ++++ 3 files changed, 176 insertions(+), 42 deletions(-) diff --git a/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py b/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py index 54d5e8dc13..b129ff6db0 100644 --- a/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py +++ b/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py @@ -5,6 +5,7 @@ import re from sqlbuild.adapter.relations.constants import ( + BACKTICK, DOUBLE_QUOTE, IDENTIFIER_QUOTE_PAIRS, QUOTED_IDENTIFIER_MIN_LENGTH, @@ -12,19 +13,58 @@ ) from sqlbuild.adapter.relations.models import StatementMetadataEffect -_LINE_COMMENT: re.Pattern[str] = re.compile(r"--[^\n]*") -_BLOCK_COMMENT: re.Pattern[str] = re.compile(r"/\*.*?\*/", re.DOTALL) +_LEXEME_ALTERNATIVES: tuple[str, ...] = ( + r"(?P--[^\n]*|/\*.*?\*/)", + r"(?P\$(?P[A-Za-z_]*)\$.*?\$(?P=tag)\$)", + "{literal}", + r'(?P"(?:[^"]|"")*"|`[^`]*`{bracket})', + r"(?P['\"`]|/\*|\$[A-Za-z_]*\${bracket_open})", +) +_STANDARD_LITERAL: str = r"(?P'(?:[^']|'')*')" +_BACKSLASH_LITERAL: str = r"(?P'(?:[^'\\]|\\.|'')*')" +_BRACKET_IDENTIFIER: str = r"|\[[^\]]*\]" +_BRACKET_OPEN: str = r"|\[" + +_LEXEME_PATTERN: str = "|".join(_LEXEME_ALTERNATIVES) +_LEXERS: tuple[re.Pattern[str], ...] = ( + re.compile( + _LEXEME_PATTERN.format(literal=_STANDARD_LITERAL, bracket="", bracket_open=""), re.DOTALL + ), + re.compile( + _LEXEME_PATTERN.format( + literal=_STANDARD_LITERAL, bracket=_BRACKET_IDENTIFIER, bracket_open=_BRACKET_OPEN + ), + re.DOTALL, + ), + re.compile( + _LEXEME_PATTERN.format(literal=_BACKSLASH_LITERAL, bracket="", bracket_open=""), re.DOTALL + ), + re.compile( + _LEXEME_PATTERN.format( + literal=_BACKSLASH_LITERAL, bracket=_BRACKET_IDENTIFIER, bracket_open=_BRACKET_OPEN + ), + re.DOTALL, + ), +) +_UNTERMINATED_MARKER: str = "\x00" +_LEXEME_REPLACEMENTS: dict[str, str] = { + "comment": " ", + "dollar": "''", + "literal": "''", + "unterminated": _UNTERMINATED_MARKER, +} _FIRST_WORD: re.Pattern[str] = re.compile(r"[A-Za-z]+") -_IDENTIFIER_PART: str = r'(?:"(?:[^"]|"")*"|`[^`]*`|\[[^\]]*\]|[A-Za-z0-9_$@#-]+)' -_QUALIFIED_NAME: str = rf"(?P{_IDENTIFIER_PART}(?:\s*\.\s*{_IDENTIFIER_PART})*)" +_IDENTIFIER_PART: str = r'(?:"(?:[^"]|"")*"|`[^`]*`|\[[^\]]*\]|[A-Za-z_][A-Za-z0-9_$]*)' +_QUALIFIED_NAME: str = ( + rf"(?P{_IDENTIFIER_PART}(?:\s*\.\s*{_IDENTIFIER_PART})*)(?P(?=\s|\()(?s:.*)|$)" +) _RELATION_KIND: str = ( r"(?:(?:OR\s+REPLACE|LOCAL|GLOBAL|TEMP|TEMPORARY|TRANSIENT|VOLATILE|SECURE|RECURSIVE|" r"MATERIALIZED|EXTERNAL|DYNAMIC|ICEBERG|HYBRID|UNLOGGED|FOREIGN|STREAMING)\s+)*" r"(?:TABLE|VIEW)\s+" ) _CREATE_RELATION: re.Pattern[str] = re.compile( - rf"^CREATE\s+{_RELATION_KIND}(?:IF\s+NOT\s+EXISTS\s+)?{_QUALIFIED_NAME}", - re.IGNORECASE, + rf"^CREATE\s+{_RELATION_KIND}(?:IF\s+NOT\s+EXISTS\s+)?{_QUALIFIED_NAME}", re.IGNORECASE ) _DROP_RELATION: re.Pattern[str] = re.compile( rf"^DROP\s+{_RELATION_KIND}(?:IF\s+EXISTS\s+)?{_QUALIFIED_NAME}", re.IGNORECASE @@ -32,21 +72,39 @@ _ALTER_RELATION: re.Pattern[str] = re.compile( rf"^ALTER\s+{_RELATION_KIND}(?:IF\s+EXISTS\s+)?{_QUALIFIED_NAME}", re.IGNORECASE ) -_ALTER_SECOND_RELATION: re.Pattern[str] = re.compile( - rf"\b(?:RENAME\s+TO|SWAP\s+WITH)\s+{_QUALIFIED_NAME}", re.IGNORECASE -) -_SELECT_INTO_RELATION: re.Pattern[str] = re.compile(rf"\bINTO\s+{_QUALIFIED_NAME}", re.IGNORECASE) _COPY_INTO_RELATION: re.Pattern[str] = re.compile( rf"^COPY\s+INTO\s+{_QUALIFIED_NAME}", re.IGNORECASE ) +_SELECT_INTO_RELATION: re.Pattern[str] = re.compile(rf"\bINTO\s+{_QUALIFIED_NAME}", re.IGNORECASE) +_SELECT_INTO_KEYWORD: re.Pattern[str] = re.compile(r"\bINTO\b", re.IGNORECASE) +_SECOND_RELATION: re.Pattern[str] = re.compile( + rf"^\s+(?:RENAME\s+TO|SWAP\s+WITH)\s+{_QUALIFIED_NAME}", re.IGNORECASE +) +_FOLLOWING_WORDS: dict[re.Pattern[str], re.Pattern[str]] = { + _CREATE_RELATION: re.compile( + r"^\s*(?:$|\(|(?:AS|CLONE|LIKE|COPY|USING|COMMENT|CLUSTER|PARTITION|OPTIONS|WITH|" + r"TBLPROPERTIES|LOCATION|DATA_RETENTION_TIME_IN_DAYS|CHANGE_TRACKING)\b)", + re.IGNORECASE, + ), + _DROP_RELATION: re.compile(r"^\s*(?:CASCADE|RESTRICT|PURGE)?\s*$", re.IGNORECASE), + _ALTER_RELATION: re.compile( + r"^\s+(?:ADD|DROP|ALTER|MODIFY|RENAME|SWAP|SET|UNSET|CLUSTER|RECLUSTER|OWNER|" + r"SUSPEND|RESUME|REFRESH)\b", + re.IGNORECASE, + ), + _COPY_INTO_RELATION: re.compile(r"^\s+FROM\b", re.IGNORECASE), + _SELECT_INTO_RELATION: re.compile(r"^\s+FROM\b", re.IGNORECASE), + _SECOND_RELATION: re.compile(r"^\s*$"), +} _CREATE_SCHEMA_IF_NOT_EXISTS: re.Pattern[str] = re.compile( - r"^CREATE\s+SCHEMA\s+IF\s+NOT\s+EXISTS\s+[^\s;]+\s*;?\s*$", re.IGNORECASE + r"^CREATE\s+SCHEMA\s+IF\s+NOT\s+EXISTS\s+[^\s;]+\s*$", re.IGNORECASE ) _RELATION_NEUTRAL_OBJECT: re.Pattern[str] = re.compile( r"^(?:CREATE|DROP|ALTER)\s+(?:OR\s+REPLACE\s+)?(?:(?:SECURE|TEMP|TEMPORARY|AGGREGATE)\s+)*" r"(?:FUNCTION|PROCEDURE|SEQUENCE|STAGE|FILE\s+FORMAT|INDEX|UNIQUE\s+INDEX|WAREHOUSE|MACRO)\b", re.IGNORECASE, ) +_RENAMING_ALTER: re.Pattern[str] = re.compile(r"^\s+(?:RENAME\s+TO|SWAP\s+WITH)\b", re.IGNORECASE) _QUERY_STATEMENTS: frozenset[str] = frozenset({"SELECT", "WITH"}) _READ_OR_DATA_STATEMENTS: frozenset[str] = frozenset( { @@ -88,15 +146,45 @@ def statement_metadata_effect(sql: str) -> StatementMetadataEffect: - """Return which cached relation names ``sql`` may have changed.""" - - statement: str = ( - _BLOCK_COMMENT.sub(" ", _LINE_COMMENT.sub(" ", sql)) - .strip() - .rstrip(STATEMENT_SEPARATOR) - .strip() - .lstrip("(") - ) + """Return which cached relation names ``sql`` may have changed; ambiguity changes all.""" + + effects: set[StatementMetadataEffect] = { + _code_effect(code=_code_text(sql=sql, lexer=lexer)) for lexer in _LEXERS + } + if len(effects) != 1: + return _ALL_RELATIONS + return effects.pop() + + +def relation_name_key(name: str) -> str: + """Return the invalidation key for one unquoted or quoted unqualified relation name.""" + + return _strip_identifier_quotes(name).casefold() + + +def qualified_relation_name_key(qualified: str) -> str: + """Return the invalidation key for the last part of a possibly qualified relation name.""" + + return _unqualified_name_key(qualified) + + +def _code_text(*, sql: str, lexer: re.Pattern[str]) -> str | None: + if _UNTERMINATED_MARKER in sql: + return None + code: str = lexer.sub(_replace_lexeme, sql) + if _UNTERMINATED_MARKER in code: + return None + return code + + +def _replace_lexeme(match: re.Match[str]) -> str: + return _LEXEME_REPLACEMENTS.get(match.lastgroup or "", match.group(0)) + + +def _code_effect(*, code: str | None) -> StatementMetadataEffect: + if code is None: + return _ALL_RELATIONS + statement: str = code.strip().rstrip(STATEMENT_SEPARATOR).strip().lstrip("(") first: re.Match[str] | None = _FIRST_WORD.match(statement) if first is None or STATEMENT_SEPARATOR in statement: return _ALL_RELATIONS @@ -117,41 +205,51 @@ def statement_metadata_effect(sql: str) -> StatementMetadataEffect: return _relation_ddl_effect(statement=statement) -def relation_name_key(name: str) -> str: - """Return the invalidation key for one unquoted or quoted unqualified relation name.""" - - return _strip_identifier_quotes(name).casefold() - - def _select_into_effect(*, statement: str) -> StatementMetadataEffect: - created: re.Match[str] | None = _SELECT_INTO_RELATION.search(statement) - if created is None: + if _SELECT_INTO_KEYWORD.search(statement) is None: return _NO_EFFECT - return StatementMetadataEffect( - relation_names=frozenset({_unqualified_name_key(created.group("name"))}) - ) + created: re.Match[str] | None = _SELECT_INTO_RELATION.search(statement) + names: frozenset[str] | None = _unambiguous_names(pattern=_SELECT_INTO_RELATION, match=created) + if names is None: + return _ALL_RELATIONS + return StatementMetadataEffect(relation_names=names) def _relation_ddl_effect(*, statement: str) -> StatementMetadataEffect: pattern: re.Pattern[str] for pattern in (_CREATE_RELATION, _DROP_RELATION, _ALTER_RELATION, _COPY_INTO_RELATION): match: re.Match[str] | None = pattern.match(statement) - if match is None: - continue - names: set[str] = {_unqualified_name_key(match.group("name"))} - if pattern is _ALTER_RELATION: - names.update( - _unqualified_name_key(second.group("name")) - for second in _ALTER_SECOND_RELATION.finditer(statement, match.end()) - ) - return StatementMetadataEffect(relation_names=frozenset(names)) + names: frozenset[str] | None = _unambiguous_names(pattern=pattern, match=match) + if names is not None: + return StatementMetadataEffect(relation_names=names) return _ALL_RELATIONS +def _unambiguous_names( + *, pattern: re.Pattern[str], match: re.Match[str] | None +) -> frozenset[str] | None: + if match is None: + return None + rest: str = match.group("rest") + if _FOLLOWING_WORDS[pattern].match(rest) is None: + return None + name: str = _unqualified_name_key(match.group("name")) + if pattern is not _ALTER_RELATION or _RENAMING_ALTER.match(rest) is None: + return frozenset({name}) + second: frozenset[str] | None = _unambiguous_names( + pattern=_SECOND_RELATION, match=_SECOND_RELATION.match(rest) + ) + if second is None: + return None + return frozenset({name, *second}) + + def _unqualified_name_key(qualified: str) -> str: parts: list[str] = re.findall(_IDENTIFIER_PART, qualified) - last: str = _strip_identifier_quotes(parts[-1]) - return last.rsplit(".", 1)[-1].casefold() + last: str = parts[-1] + if last.startswith(BACKTICK): + return _strip_identifier_quotes(last).rsplit(".", 1)[-1].casefold() + return relation_name_key(last) def _strip_identifier_quotes(part: str) -> str: diff --git a/src/sqlbuild/adapter/relations/constants.py b/src/sqlbuild/adapter/relations/constants.py index 750edf868b..53d4ad02c5 100644 --- a/src/sqlbuild/adapter/relations/constants.py +++ b/src/sqlbuild/adapter/relations/constants.py @@ -10,3 +10,4 @@ IDENTIFIER_QUOTE_PAIRS: frozenset[tuple[str, str]] = frozenset({("`", "`"), ("[", "]")}) RELATION_COLUMNS_LOOKUP: str = "columns" RELATION_EXISTS_LOOKUP: str = "exists" +BACKTICK: str = "`" diff --git a/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py index b908b44d05..676cc41e61 100644 --- a/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py +++ b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py @@ -76,6 +76,24 @@ expected_relation_names=frozenset({"orders__delta"}), expected_invalidates_all=False, ), + StatementMetadataEffectTestCase( + description="quoted dotted name keeps its dot like a lookup name", + sql='CREATE TABLE "s"."my.table" (id INT)', + expected_relation_names=frozenset({"my.table"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="statement separator inside a literal is data", + sql="CREATE OR REPLACE TABLE marts.notes AS SELECT 'a;b -- c' AS note", + expected_relation_names=frozenset({"notes"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="escaped backslash literal reads the same either way", + sql="CREATE OR REPLACE TABLE marts.codes AS SELECT REGEXP_LIKE(code, '\\\\d+') AS ok", + expected_relation_names=frozenset({"codes"}), + expected_invalidates_all=False, + ), StatementMetadataEffectTestCase( description="select into creates a table", sql="SELECT * INTO [dbo].[orders__staging] FROM (SELECT 1 AS id) AS __create_source", @@ -128,6 +146,8 @@ def test_given_targeted_ddl_statement_when_classifying_then_reports_possible_met "CREATE SCHEMA IF NOT EXISTS analytics.marts", "CREATE OR REPLACE FUNCTION marts.add_one(x NUMBER) RETURNS NUMBER AS $$ x + 1 $$", "TRUNCATE TABLE marts.orders", + "CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$ SELECT 1; SELECT 2 $$", + "INSERT INTO marts.notes VALUES ('drop table u; --')", ) ) ], @@ -165,6 +185,21 @@ def test_given_read_or_data_statement_when_classifying_then_reports_possible_met ("session parameters may change identifier resolution", "ALTER SESSION SET X = 1"), ("unparseable relation name", 'CREATE TABLE "unterminated AS SELECT 1'), ("empty statement", " "), + ("keyword before the relation name", 'ALTER TABLE ONLY "s"."t" ADD COLUMN c INT'), + ("several dropped relations", "DROP TABLE IF EXISTS a, b"), + ("unquoted non-ascii name", "CREATE TABLE s.café (id INT)"), + ("comment marker inside a literal hides nothing", "SELECT 'a--b'; DROP TABLE u"), + ("separator after a literal", "INSERT INTO t VALUES ('x'); DROP TABLE u"), + ( + "backslash-escaped quote is ambiguous", + "SELECT 'a\\' AS x FROM t WHERE y = '; DROP TABLE u'", + ), + ("unterminated literal", "CREATE TABLE t AS SELECT 'open"), + ( + "unexpected clause after the name", + "CREATE DYNAMIC TABLE t TARGET_LAG = '1 hour' AS SELECT 1", + ), + ("rename target followed by more text", "ALTER TABLE a RENAME TO b c"), ) ], ids=lambda case: case.description, From 2b21e4dd178ba43db0cae49f9a16d0f89eb09918 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:08:17 +0100 Subject: [PATCH 6/9] fix(bigquery): evict cached metadata after copy and load jobs replace_table_from_relation and load_seed change relations through client jobs that no statement listener sees; they now evict the destination (and the copy origin) from the build's relation metadata cache. Refs CHI-455 --- .../classes/relation_metadata_cache.py | 8 +++ .../main/invalidate_cached_relations.py | 18 +++++ .../bigquery/classes/bigquery_adapter.py | 33 +++++---- .../sqlbuild/adapters/bigquery/_test_types.py | 9 +++ .../src/sqlbuild/adapters/bigquery/helpers.py | 53 ++++++++++++++ .../test_relation_metadata_invalidation.py | 71 +++++++++++++++++++ 6 files changed, 179 insertions(+), 13 deletions(-) create mode 100644 src/sqlbuild/adapter/relations/main/invalidate_cached_relations.py create mode 100644 tests/unit/src/sqlbuild/adapters/bigquery/test_relation_metadata_invalidation.py diff --git a/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py b/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py index 4a757a20aa..d1b56cfd4c 100644 --- a/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py +++ b/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py @@ -7,6 +7,7 @@ from typing import cast from sqlbuild.adapter.relations._helpers.ddl_effects import ( + qualified_relation_name_key, relation_name_key, statement_metadata_effect, ) @@ -39,6 +40,13 @@ def lookup[ValueT]( return cast(ValueT, self._lookup(key=(kind, database, schema, name), read=read)) + def invalidate_relations(self, *, qualified_names: tuple[str, ...]) -> None: + """Evict relations changed outside SQL statements, such as by a warehouse copy job.""" + + names: frozenset[str] = frozenset(map(qualified_relation_name_key, qualified_names)) + with self._lock: + self._invalidate_names(names=names) + def observe_statement(self, *, sql: str) -> None: """Invalidate entries the finished statement ``sql`` may have changed.""" diff --git a/src/sqlbuild/adapter/relations/main/invalidate_cached_relations.py b/src/sqlbuild/adapter/relations/main/invalidate_cached_relations.py new file mode 100644 index 0000000000..a0b3a3e81d --- /dev/null +++ b/src/sqlbuild/adapter/relations/main/invalidate_cached_relations.py @@ -0,0 +1,18 @@ +"""Evict relations an adapter changed without a SQL statement from the open cache.""" + +from __future__ import annotations + +from typing import Any + +from sqlbuild.adapter.relations._helpers.relation_metadata_context import ( + current_relation_metadata_cache, +) +from sqlbuild.adapter.relations.classes.relation_metadata_cache import RelationMetadataCache + + +def invalidate_cached_relations(*, adapter: Any, qualified_names: tuple[str, ...]) -> None: + """Make the next lookup of each named relation read the warehouse again.""" + + cache: RelationMetadataCache | None = current_relation_metadata_cache(adapter=adapter) + if cache is not None: + cache.invalidate_relations(qualified_names=qualified_names) diff --git a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py index 846a84416e..595820a0f5 100644 --- a/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py +++ b/src/sqlbuild/adapters/bigquery/classes/bigquery_adapter.py @@ -88,6 +88,7 @@ from sqlbuild.adapter.relations.main.get_columns_for_relations import ( get_columns_for_relations_bulk, ) +from sqlbuild.adapter.relations.main.invalidate_cached_relations import invalidate_cached_relations from sqlbuild.adapter.relations.main.relation_age_timestamp import relation_age_timestamp_utc from sqlbuild.adapter.state_sql.main.render_insert_source_freshness_records_sql import ( render_insert_source_freshness_records_sql, @@ -2010,12 +2011,15 @@ def replace_table_from_relation( write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE, ) statement_recorder.record(f"COPY WRITE_TRUNCATE {origin} TO {destination}") - connection.client.copy_table( - origin_table, - destination_table, - job_config=job_config, - location=connection.location, - ).result() + try: + connection.client.copy_table( + origin_table, + destination_table, + job_config=job_config, + location=connection.location, + ).result() + finally: + invalidate_cached_relations(adapter=self, qualified_names=(destination, origin)) def move_or_copy_relation( self, @@ -2313,13 +2317,16 @@ def load_seed( statement_recorder.record( f"LOAD CSV {file_path} INTO {destination} ({', '.join(col.name for col in columns)})" ) - with file_path.open("rb") as seed_file: - connection.client.load_table_from_file( - seed_file, - self._strip_identifier_quotes(destination), - job_config=job_config, - location=connection.location, - ).result() + try: + with file_path.open("rb") as seed_file: + connection.client.load_table_from_file( + seed_file, + self._strip_identifier_quotes(destination), + job_config=job_config, + location=connection.location, + ).result() + finally: + invalidate_cached_relations(adapter=self, qualified_names=(destination,)) def add_columns( self, diff --git a/tests/unit/src/sqlbuild/adapters/bigquery/_test_types.py b/tests/unit/src/sqlbuild/adapters/bigquery/_test_types.py index 7e4d7d9111..49945ac720 100644 --- a/tests/unit/src/sqlbuild/adapters/bigquery/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/bigquery/_test_types.py @@ -210,3 +210,12 @@ class BigQueryCloneRefusalTestCase: description: str driver_error: Exception expected_refusal: bool + + +@dataclass(frozen=True) +class BigQueryJobInvalidationTestCase: + """A BigQuery job that changes relations without a SQL statement.""" + + description: str + operation: str + expected_reads: dict[str, int] diff --git a/tests/unit/src/sqlbuild/adapters/bigquery/helpers.py b/tests/unit/src/sqlbuild/adapters/bigquery/helpers.py index 1305c9a277..d5b7e30583 100644 --- a/tests/unit/src/sqlbuild/adapters/bigquery/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/bigquery/helpers.py @@ -1,11 +1,14 @@ from __future__ import annotations from collections.abc import Callable +from pathlib import Path from types import MappingProxyType from typing import Any +from sqlbuild.adapter.contract.classes.statement_recorder import StatementRecorder from sqlbuild.adapter.contract.models import ColumnInfo, RetentionRequest from sqlbuild.adapter.contract.types import RetentionScope +from sqlbuild.adapters.bigquery.classes.bigquery_adapter import BigQueryAdapter def build_retention_request(*, desired_days: int) -> RetentionRequest: @@ -201,3 +204,53 @@ class FakeBigQueryBadRequest(Exception): def __init__(self, message: str, *, errors: list[dict[str, object]]) -> None: super().__init__(message) self.errors: list[dict[str, object]] = errors + + +class FakeBigQueryJobClient(FakeBigQueryClient): + """Client double accepting copy and load jobs.""" + + def copy_table(self, origin: str, destination: str, **kwargs: Any) -> FakeBigQueryJob: + del origin, destination, kwargs + return FakeBigQueryJob(self.rows) + + def load_table_from_file(self, file: Any, destination: str, **kwargs: Any) -> FakeBigQueryJob: + del file, destination, kwargs + return FakeBigQueryJob(self.rows) + + +class CountingColumnsBigQueryAdapter(BigQueryAdapter): + """BigQuery adapter whose column reads are counted instead of sent.""" + + def __init__(self) -> None: + super().__init__() + self.column_reads: dict[str, int] = {} + + def get_columns( + self, *, connection: Any, database: str | None, schema: str | None, name: str + ) -> tuple[ColumnInfo, ...]: + del connection, database, schema + self.column_reads[name] = self.column_reads.get(name, 0) + 1 + return (ColumnInfo(name="order_id", type="INT64"),) + + +def run_bigquery_job( + *, operation: str, adapter: BigQueryAdapter, connection: Any, seed_path: Path +) -> None: + """Run one BigQuery copy or load job between the orders relations.""" + + jobs: dict[str, Callable[[], None]] = { + "copy": lambda: adapter.replace_table_from_relation( + connection=connection, + destination="example-project.marts.orders", + origin="example-project.marts.orders__stage", + statement_recorder=StatementRecorder(), + ), + "load": lambda: adapter.load_seed( + connection=connection, + destination="example-project.marts.orders", + file_path=seed_path, + columns=(ColumnInfo(name="order_id", type="INT64"),), + statement_recorder=StatementRecorder(), + ), + } + jobs[operation]() diff --git a/tests/unit/src/sqlbuild/adapters/bigquery/test_relation_metadata_invalidation.py b/tests/unit/src/sqlbuild/adapters/bigquery/test_relation_metadata_invalidation.py new file mode 100644 index 0000000000..ad4482cbc8 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/bigquery/test_relation_metadata_invalidation.py @@ -0,0 +1,71 @@ +"""BigQuery copy and load jobs change relations without SQL, so they evict cached metadata.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from sqlbuild.adapter.relations.main.cached_relation_columns import cached_relation_columns +from sqlbuild.adapter.relations.main.relation_metadata_cache_scope import ( + relation_metadata_cache_scope, +) +from sqlbuild.adapters.bigquery.classes.bigquery_connection import _BigQueryConnection +from tests.unit.src.sqlbuild.adapters.bigquery._test_types import BigQueryJobInvalidationTestCase +from tests.unit.src.sqlbuild.adapters.bigquery.helpers import ( + CountingColumnsBigQueryAdapter, + FakeBigQueryJobClient, + run_bigquery_job, +) + +_RELATIONS: tuple[str, ...] = ("orders", "orders__stage", "customers") + + +@pytest.mark.parametrize( + "test_case", + [ + BigQueryJobInvalidationTestCase( + description="copy job evicts destination and origin", + operation="copy", + expected_reads={"orders": 2, "orders__stage": 2, "customers": 1}, + ), + BigQueryJobInvalidationTestCase( + description="seed load job evicts its destination", + operation="load", + expected_reads={"orders": 2, "orders__stage": 1, "customers": 1}, + ), + ], + ids=lambda case: case.description, +) +def test_given_bigquery_job_when_looking_up_again_then_changed_relations_are_reread( + test_case: BigQueryJobInvalidationTestCase, tmp_path: Path +) -> None: + adapter: CountingColumnsBigQueryAdapter = CountingColumnsBigQueryAdapter() + connection: _BigQueryConnection = _BigQueryConnection( + client=FakeBigQueryJobClient(), location=None + ) + seed_path: Path = tmp_path / "orders.csv" + _ = seed_path.write_text("order_id\n1\n", encoding="utf-8") + + with relation_metadata_cache_scope(adapter=adapter): + name: str + for name in _RELATIONS: + _ = cached_relation_columns( + adapter=adapter, connection=connection, database=None, schema="marts", name=name + ) + run_bigquery_job( + operation=test_case.operation, + adapter=adapter, + connection=connection, + seed_path=seed_path, + ) + for name in _RELATIONS: + _ = cached_relation_columns( + adapter=adapter, connection=connection, database=None, schema="marts", name=name + ) + + assert adapter.column_reads == test_case.expected_reads + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) From 1880407492bc22ff766bef8f59c44fb23c4a8617 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:08:17 +0100 Subject: [PATCH 7/9] fix(snowflake): forget the remembered session database after USE The current database used to qualify unqualified SHOW reads was remembered per connection for the life of the session. A USE, CALL, or EXECUTE statement on that connection now clears it, so build and planning lookups re-resolve it. Refs CHI-455 --- .../snowflake/classes/snowflake_adapter.py | 6 +- .../snowflake/classes/snowflake_connection.py | 16 ++++- .../snowflake/classes/snowflake_cursor.py | 59 ++++++++++------ src/sqlbuild/adapters/snowflake/constants.py | 1 + .../snowflake/inspection/_test_types.py | 10 +++ .../adapters/snowflake/inspection/helpers.py | 10 +++ .../test_session_database_switch.py | 70 +++++++++++++++++++ 7 files changed, 146 insertions(+), 26 deletions(-) create mode 100644 tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index dbd00df1a0..7555d23333 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -123,6 +123,7 @@ from sqlbuild.adapters.snowflake.classes.snowflake_connection import _SnowflakeConnection from sqlbuild.adapters.snowflake.constants import ( BASE_TABLE_METADATA_TYPE, + CURRENT_DATABASE_ATTRIBUTE, EXTERNAL_BROWSER_AUTHENTICATOR, MAX_STATEMENT_TIMEOUT_SECONDS, MFA_AUTHENTICATOR, @@ -151,7 +152,6 @@ _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" @@ -2021,7 +2021,7 @@ def _show_database( @staticmethod def _current_database(*, connection: _SnowflakeConnection) -> str | None: - cached: object = getattr(connection, _CURRENT_DATABASE_ATTRIBUTE, _UNRESOLVED) + cached: object = getattr(connection, CURRENT_DATABASE_ATTRIBUTE, _UNRESOLVED) if cached is not _UNRESOLVED: return cast(str | None, cached) cursor: Any = connection.cursor() @@ -2031,7 +2031,7 @@ def _current_database(*, connection: _SnowflakeConnection) -> str | None: 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) + setattr(connection, CURRENT_DATABASE_ATTRIBUTE, current) return current def _fetch_paged_show_rows( diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py index 19349018df..419bf7b8bb 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py @@ -1,8 +1,14 @@ """Snowflake connection wrapper with statement telemetry.""" +import re from typing import Any from sqlbuild.adapters.snowflake.classes.snowflake_cursor import _SnowflakeCursor +from sqlbuild.adapters.snowflake.constants import CURRENT_DATABASE_ATTRIBUTE + +_SESSION_CONTEXT_STATEMENT: re.Pattern[str] = re.compile( + r"^(?:\s+|--[^\n]*|/\*.*?\*/)*(?:USE|CALL|EXECUTE)\b", re.IGNORECASE | re.DOTALL +) class _SnowflakeConnection: @@ -21,4 +27,12 @@ def close(self) -> None: self.raw_connection.close() def cursor(self) -> _SnowflakeCursor: - return _SnowflakeCursor(self.raw_connection.cursor()) + return _SnowflakeCursor( + self.raw_connection.cursor(), on_executed=self._forget_changed_session_context + ) + + def _forget_changed_session_context(self, *, sql: str) -> None: + """Drop the remembered session database once a statement may have switched it.""" + + if _SESSION_CONTEXT_STATEMENT.match(sql) is not None: + _ = self.__dict__.pop(CURRENT_DATABASE_ATTRIBUTE, None) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_cursor.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_cursor.py index 0d9dc57594..712fe1b1f6 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_cursor.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_cursor.py @@ -25,12 +25,15 @@ current_execution_identity, ) from sqlbuild.runtime.observability.models import ExecutionIdentity +from sqlbuild.runtime.observability.types import StatementListener _LOGGER: logging.Logger = logging.getLogger("sqlbuild.cost") _QUERY_TAG_PARAMETER: str = "QUERY_TAG" _STATEMENT_TIMEOUT_PARAMETER: str = "STATEMENT_TIMEOUT_IN_SECONDS" _STATEMENT_PARAMETER_ERROR_FRAGMENT: str = "STATEMENT PARAMETER" _RAW_CURSOR_ATTRIBUTE: str = "raw_cursor" +_ON_EXECUTED_ATTRIBUTE: str = "on_executed" +_OWN_ATTRIBUTES: frozenset[str] = frozenset({_RAW_CURSOR_ATTRIBUTE, _ON_EXECUTED_ATTRIBUTE}) _EXECUTEMANY_INTENT: str = "executemany" @@ -44,11 +47,12 @@ def _statement_id() -> str: class _SnowflakeCursor: """Delegate cursor operations while tagging and recording executed statements.""" - def __init__(self, raw_cursor: Any) -> None: + def __init__(self, raw_cursor: Any, *, on_executed: StatementListener | None = None) -> None: object.__setattr__(self, _RAW_CURSOR_ATTRIBUTE, raw_cursor) + object.__setattr__(self, _ON_EXECUTED_ATTRIBUTE, on_executed) def __setattr__(self, name: str, value: Any) -> None: - if name == _RAW_CURSOR_ATTRIBUTE: + if name in _OWN_ATTRIBUTES: object.__setattr__(self, name, value) return setattr(self.raw_cursor, name, value) @@ -60,16 +64,19 @@ def execute(self, sql: str, *args: Any, **kwargs: Any) -> Any: tagged_kwargs, query_tag_injected = _with_query_tag( kwargs=deadline_kwargs, context=context, statement_id=statement_id ) - return self._execute_with_telemetry( - operation=self.raw_cursor.execute, - sql=sql, - args=args, - kwargs=tagged_kwargs, - context=context, - statement_id=statement_id, - query_tag_injected=query_tag_injected, - intent="execute", - ) + try: + return self._execute_with_telemetry( + operation=self.raw_cursor.execute, + sql=sql, + args=args, + kwargs=tagged_kwargs, + context=context, + statement_id=statement_id, + query_tag_injected=query_tag_injected, + intent="execute", + ) + finally: + self._notify_executed(sql=sql) def executemany(self, sql: str, *args: Any, **kwargs: Any) -> Any: context: CostResourceContext | None = CostContext.current() @@ -78,16 +85,24 @@ def executemany(self, sql: str, *args: Any, **kwargs: Any) -> Any: tagged_kwargs, query_tag_injected = _with_query_tag( kwargs=deadline_kwargs, context=context, statement_id=statement_id ) - return self._execute_with_telemetry( - operation=self.raw_cursor.executemany, - sql=sql, - args=args, - kwargs=tagged_kwargs, - context=context, - statement_id=statement_id, - query_tag_injected=query_tag_injected, - intent=_EXECUTEMANY_INTENT, - ) + try: + return self._execute_with_telemetry( + operation=self.raw_cursor.executemany, + sql=sql, + args=args, + kwargs=tagged_kwargs, + context=context, + statement_id=statement_id, + query_tag_injected=query_tag_injected, + intent=_EXECUTEMANY_INTENT, + ) + finally: + self._notify_executed(sql=sql) + + def _notify_executed(self, *, sql: str) -> None: + listener: StatementListener | None = object.__getattribute__(self, _ON_EXECUTED_ATTRIBUTE) + if listener is not None: + listener(sql=sql) def _execute_with_telemetry( self, diff --git a/src/sqlbuild/adapters/snowflake/constants.py b/src/sqlbuild/adapters/snowflake/constants.py index 33c5ca7c56..69692ed103 100644 --- a/src/sqlbuild/adapters/snowflake/constants.py +++ b/src/sqlbuild/adapters/snowflake/constants.py @@ -23,3 +23,4 @@ UNBOUNDED_TEXT_TYPE_NAMES: frozenset[str] = frozenset({"STRING", "TEXT", "VARCHAR"}) TRUE_METADATA_VALUE: str = "YES" SHOW_LIKE_ESCAPE_CHARACTER: str = "\\" +CURRENT_DATABASE_ATTRIBUTE: str = "_sqlbuild_current_database" diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py index ec19a6d018..3eb644e84e 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/_test_types.py @@ -204,3 +204,13 @@ class ShowColumnsEquivalenceTestCase: schema: str name: str expected_column_count: int + + +@dataclass(frozen=True) +class SessionDatabaseSwitchTestCase: + """Unqualified lookups before and after a USE statement switches the session database.""" + + description: str + switch_sql: str + expected_scopes: tuple[str, ...] + expected_session_reads: int diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py index 0a5689440d..63541ae13b 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -50,6 +50,7 @@ ) _ANY_PATTERN: re.Pattern[str] = re.compile(r"") _CURRENT_DATABASE_PATTERN: re.Pattern[str] = re.compile(r"^SELECT CURRENT_DATABASE\(\)$") +_USE_DATABASE_PATTERN: re.Pattern[str] = re.compile(r'^USE DATABASE "?(?P[^"]+)"?$') _SHOW_SCHEMAS_PATTERN: re.Pattern[str] = re.compile( r"^SHOW SCHEMAS LIKE '(?P(?:[^']|'')*)' IN DATABASE \"(?P[^\"]*)\"$" ) @@ -317,6 +318,7 @@ def _answer(self, *, sql: str, params: tuple[object, ...]) -> tuple[list[tuple[A (_SHOW_SCHEMA_COLUMNS_PATTERN, self._answer_show_schema_columns), (_CURSOR_BOUND_PATTERN, self._answer_cursor_bounds), (_CURRENT_DATABASE_PATTERN, self._answer_current_database), + (_USE_DATABASE_PATTERN, self._answer_use_database), (_SHOW_SCHEMAS_PATTERN, self._answer_show_schemas), (_ANY_PATTERN, _answer_status), ) @@ -373,6 +375,14 @@ def _answer_current_database( del sql, params return [(self.current_database,)], (("CURRENT_DATABASE()",),) + def _answer_use_database( + self, *, sql: str, params: tuple[object, ...] + ) -> tuple[list[tuple[Any, ...]], Any]: + del params + with self._lock: + self.current_database = next(_USE_DATABASE_PATTERN.finditer(sql)).group("database") + return [], _STATUS_DESCRIPTION + def _existing_schemas(self) -> frozenset[tuple[str, str]]: return self.extra_schemas | {(fake.database, fake.schema) for fake in self.relations} diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py new file mode 100644 index 0000000000..ae6ebc34e4 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py @@ -0,0 +1,70 @@ +"""A remembered session database is forgotten once a statement may switch it.""" + +from __future__ import annotations + +import re + +import pytest + +from sqlbuild.adapters.snowflake.classes.snowflake_connection import _SnowflakeConnection +from tests.unit.src.sqlbuild.adapters.snowflake.inspection._test_types import ( + SessionDatabaseSwitchTestCase, +) +from tests.unit.src.sqlbuild.adapters.snowflake.inspection.helpers import ( + OfflineSnowflakeAdapter, + RecordingSnowflakeWarehouse, + build_inspection_catalog_relations, + build_offline_snowflake, +) + +_SCOPE_PATTERN: re.Pattern[str] = re.compile(r'IN SCHEMA ("[^"]+")\.') + + +@pytest.mark.parametrize( + "test_case", + [ + SessionDatabaseSwitchTestCase( + description="USE DATABASE re-resolves build and planning lookups", + switch_sql='USE DATABASE "ARCHIVE"', + expected_scopes=(*('"ANALYTICS"',) * 3, *('"ARCHIVE"',) * 3), + expected_session_reads=2, + ), + SessionDatabaseSwitchTestCase( + description="statements that cannot switch databases keep the remembered one", + switch_sql="SELECT 1", + expected_scopes=('"ANALYTICS"',) * 6, + expected_session_reads=1, + ), + ], + ids=lambda case: case.description, +) +def test_given_session_switch_when_looking_up_unqualified_relations_then_uses_current_database( + test_case: SessionDatabaseSwitchTestCase, +) -> None: + adapter: OfflineSnowflakeAdapter + connection: _SnowflakeConnection + warehouse: RecordingSnowflakeWarehouse + adapter, connection, warehouse = build_offline_snowflake( + relations=build_inspection_catalog_relations() + ) + + _ = adapter.relation_exists( + connection=connection, database=None, schema="staging", name="orders" + ) + _ = adapter.read_schema_relation_listing(connection=connection, database=None, schema="staging") + _ = adapter.execute(connection=connection, sql=test_case.switch_sql) + _ = adapter.relation_exists( + connection=connection, database=None, schema="staging", name="orders" + ) + _ = adapter.read_schema_relation_listing(connection=connection, database=None, schema="staging") + + scopes: list[str] = [] + sql: str + for sql in warehouse.attempted_sql: + scopes.extend(match.group(1) for match in _SCOPE_PATTERN.finditer(sql)) + assert tuple(scopes) == test_case.expected_scopes + assert len(warehouse.queries_of_kind("session")) == test_case.expected_session_reads + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) From 2a0be3db1721a0b53d08c8c601e889889ddf3f8f Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:39:24 +0100 Subject: [PATCH 8/9] fix(snowflake): detect session database switches with a linear scan The previous leading-whitespace/comment regex backtracked exponentially on leading whitespace and ran on every Snowflake execute. A small scanner now skips whitespace and comments in linear time, then clears the remembered database for USE, CALL, EXECUTE, and scripting blocks (BEGIN, DECLARE), or when a leading block comment never closes. Refs CHI-455 --- .../snowflake/_helpers/session_context.py | 42 +++++++++++++++++++ .../snowflake/classes/snowflake_connection.py | 8 +--- .../adapters/snowflake/inspection/helpers.py | 4 +- .../test_session_database_switch.py | 6 +++ 4 files changed, 53 insertions(+), 7 deletions(-) create mode 100644 src/sqlbuild/adapters/snowflake/_helpers/session_context.py diff --git a/src/sqlbuild/adapters/snowflake/_helpers/session_context.py b/src/sqlbuild/adapters/snowflake/_helpers/session_context.py new file mode 100644 index 0000000000..11d676e6f6 --- /dev/null +++ b/src/sqlbuild/adapters/snowflake/_helpers/session_context.py @@ -0,0 +1,42 @@ +"""Recognize Snowflake statements that may switch the session database.""" + +from __future__ import annotations + +import re + +_LEADING_WHITESPACE: re.Pattern[str] = re.compile(r"\s*+") +_FIRST_WORD: re.Pattern[str] = re.compile(r"[A-Za-z]++") +_LINE_COMMENT_START: str = "--" +_BLOCK_COMMENT_START: str = "/*" +_BLOCK_COMMENT_END: str = "*/" +_NEWLINE: str = "\n" +_SESSION_SWITCHING_WORDS: frozenset[str] = frozenset({"USE", "CALL", "EXECUTE", "BEGIN", "DECLARE"}) + + +def may_switch_session_database(sql: str) -> bool: + """Return whether ``sql`` may change the current database; unreadable text counts as yes.""" + + position: int | None = _first_code_position(sql=sql) + if position is None: + return True + word: re.Match[str] | None = _FIRST_WORD.match(sql, position) + return word is not None and word.group(0).upper() in _SESSION_SWITCHING_WORDS + + +def _first_code_position(*, sql: str) -> int | None: + position: int = 0 + while True: + whitespace: re.Match[str] | None = _LEADING_WHITESPACE.match(sql, position) + position = position if whitespace is None else whitespace.end() + if sql.startswith(_LINE_COMMENT_START, position): + line_end: int = sql.find(_NEWLINE, position) + if line_end < 0: + return len(sql) + position = line_end + 1 + continue + if not sql.startswith(_BLOCK_COMMENT_START, position): + return position + comment_end: int = sql.find(_BLOCK_COMMENT_END, position + len(_BLOCK_COMMENT_START)) + if comment_end < 0: + return None + position = comment_end + len(_BLOCK_COMMENT_END) diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py index 419bf7b8bb..1ee72c1948 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py @@ -1,15 +1,11 @@ """Snowflake connection wrapper with statement telemetry.""" -import re from typing import Any +from sqlbuild.adapters.snowflake._helpers.session_context import may_switch_session_database from sqlbuild.adapters.snowflake.classes.snowflake_cursor import _SnowflakeCursor from sqlbuild.adapters.snowflake.constants import CURRENT_DATABASE_ATTRIBUTE -_SESSION_CONTEXT_STATEMENT: re.Pattern[str] = re.compile( - r"^(?:\s+|--[^\n]*|/\*.*?\*/)*(?:USE|CALL|EXECUTE)\b", re.IGNORECASE | re.DOTALL -) - class _SnowflakeConnection: """Small wrapper exposing a DuckDB-like execute method for base adapter helpers.""" @@ -34,5 +30,5 @@ def cursor(self) -> _SnowflakeCursor: def _forget_changed_session_context(self, *, sql: str) -> None: """Drop the remembered session database once a statement may have switched it.""" - if _SESSION_CONTEXT_STATEMENT.match(sql) is not None: + if may_switch_session_database(sql): _ = self.__dict__.pop(CURRENT_DATABASE_ATTRIBUTE, None) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py index 63541ae13b..2e774c0b65 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/helpers.py @@ -50,7 +50,9 @@ ) _ANY_PATTERN: re.Pattern[str] = re.compile(r"") _CURRENT_DATABASE_PATTERN: re.Pattern[str] = re.compile(r"^SELECT CURRENT_DATABASE\(\)$") -_USE_DATABASE_PATTERN: re.Pattern[str] = re.compile(r'^USE DATABASE "?(?P[^"]+)"?$') +_USE_DATABASE_PATTERN: re.Pattern[str] = re.compile( + r'^(?s:.*?)USE DATABASE "?(?P[^";]+)"?(?s:.*)$' +) _SHOW_SCHEMAS_PATTERN: re.Pattern[str] = re.compile( r"^SHOW SCHEMAS LIKE '(?P(?:[^']|'')*)' IN DATABASE \"(?P[^\"]*)\"$" ) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py index ae6ebc34e4..53ad14f0bb 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py @@ -29,6 +29,12 @@ expected_scopes=(*('"ANALYTICS"',) * 3, *('"ARCHIVE"',) * 3), expected_session_reads=2, ), + SessionDatabaseSwitchTestCase( + description="commented scripting block is treated as a possible switch", + switch_sql='-- switch\n/* block */ BEGIN\n USE DATABASE "ARCHIVE";\nEND;', + expected_scopes=(*('"ANALYTICS"',) * 3, *('"ARCHIVE"',) * 3), + expected_session_reads=2, + ), SessionDatabaseSwitchTestCase( description="statements that cannot switch databases keep the remembered one", switch_sql="SELECT 1", From 111f9597d04faba17dad08d03c7cb48b883409f9 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:39:24 +0100 Subject: [PATCH 9/9] fix(executor): keep DDL classification linear on hostile SQL The literal-aware lexer is now an explicit scanner that stops at the first unterminated comment, literal, or dollar tag instead of rescanning, and the name and tail patterns use possessive quantifiers. CREATE ... SHALLOW/DEEP CLONE evicts only its target, digit-leading unquoted names keep their digits, and adapter-supplied names that cannot be parsed evict everything. A runtime guard covers adversarial inputs. Refs CHI-455 --- .../adapter/relations/_helpers/ddl_effects.py | 238 ++++++++++++------ .../classes/relation_metadata_cache.py | 6 +- .../_helpers/ddl_effects/_test_types.py | 18 ++ .../test_adversarial_input_runtime.py | 99 ++++++++ .../test_statement_metadata_effect.py | 61 ++++- 5 files changed, 341 insertions(+), 81 deletions(-) create mode 100644 tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_adversarial_input_runtime.py diff --git a/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py b/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py index b129ff6db0..f78d68623c 100644 --- a/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py +++ b/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py @@ -13,98 +13,79 @@ ) from sqlbuild.adapter.relations.models import StatementMetadataEffect -_LEXEME_ALTERNATIVES: tuple[str, ...] = ( - r"(?P--[^\n]*|/\*.*?\*/)", - r"(?P\$(?P[A-Za-z_]*)\$.*?\$(?P=tag)\$)", - "{literal}", - r'(?P"(?:[^"]|"")*"|`[^`]*`{bracket})', - r"(?P['\"`]|/\*|\$[A-Za-z_]*\${bracket_open})", -) -_STANDARD_LITERAL: str = r"(?P'(?:[^']|'')*')" -_BACKSLASH_LITERAL: str = r"(?P'(?:[^'\\]|\\.|'')*')" -_BRACKET_IDENTIFIER: str = r"|\[[^\]]*\]" -_BRACKET_OPEN: str = r"|\[" - -_LEXEME_PATTERN: str = "|".join(_LEXEME_ALTERNATIVES) -_LEXERS: tuple[re.Pattern[str], ...] = ( - re.compile( - _LEXEME_PATTERN.format(literal=_STANDARD_LITERAL, bracket="", bracket_open=""), re.DOTALL - ), - re.compile( - _LEXEME_PATTERN.format( - literal=_STANDARD_LITERAL, bracket=_BRACKET_IDENTIFIER, bracket_open=_BRACKET_OPEN - ), - re.DOTALL, - ), - re.compile( - _LEXEME_PATTERN.format(literal=_BACKSLASH_LITERAL, bracket="", bracket_open=""), re.DOTALL - ), - re.compile( - _LEXEME_PATTERN.format( - literal=_BACKSLASH_LITERAL, bracket=_BRACKET_IDENTIFIER, bracket_open=_BRACKET_OPEN - ), - re.DOTALL, - ), -) -_UNTERMINATED_MARKER: str = "\x00" -_LEXEME_REPLACEMENTS: dict[str, str] = { - "comment": " ", - "dollar": "''", - "literal": "''", - "unterminated": _UNTERMINATED_MARKER, -} +_LEXEME_START: re.Pattern[str] = re.compile(r"--|/\*|['\"`$\[]") +_BACKSLASH_LITERAL_STOP: re.Pattern[str] = re.compile(r"[\\']") +_DOLLAR_TAG: re.Pattern[str] = re.compile(r"\$[A-Za-z_]*+\$") +_LINE_COMMENT_START: str = "--" +_BLOCK_COMMENT_START: str = "/*" +_BLOCK_COMMENT_END: str = "*/" +_SINGLE_QUOTE: str = "'" +_BACKSLASH: str = "\\" +_DOLLAR: str = "$" +_OPEN_BRACKET: str = "[" +_CLOSE_BRACKET: str = "]" +_NEWLINE: str = "\n" +_EMPTY_LITERAL: str = "''" +_COMMENT_REPLACEMENT: str = " " _FIRST_WORD: re.Pattern[str] = re.compile(r"[A-Za-z]+") -_IDENTIFIER_PART: str = r'(?:"(?:[^"]|"")*"|`[^`]*`|\[[^\]]*\]|[A-Za-z_][A-Za-z0-9_$]*)' +_IDENTIFIER_PART: str = r'(?:"(?:[^"]|"")*+"|`[^`]*+`|\[[^\]]*+\]|[A-Za-z0-9_][A-Za-z0-9_$]*+)' _QUALIFIED_NAME: str = ( - rf"(?P{_IDENTIFIER_PART}(?:\s*\.\s*{_IDENTIFIER_PART})*)(?P(?=\s|\()(?s:.*)|$)" + rf"(?P{_IDENTIFIER_PART}(?:\s*+\.\s*+{_IDENTIFIER_PART})*+)" + r"(?P(?=\s|\()(?s:.*)|$)" +) +_LOOSE_IDENTIFIER_PART: str = r'(?:"(?:[^"]|"")*+"|`[^`]*+`|\[[^\]]*+\]|[A-Za-z0-9_$-]++)' +_LOOSE_QUALIFIED_NAME: re.Pattern[str] = re.compile( + rf"^\s*+{_LOOSE_IDENTIFIER_PART}(?:\s*+\.\s*+{_LOOSE_IDENTIFIER_PART})*+\s*+$" ) _RELATION_KIND: str = ( - r"(?:(?:OR\s+REPLACE|LOCAL|GLOBAL|TEMP|TEMPORARY|TRANSIENT|VOLATILE|SECURE|RECURSIVE|" - r"MATERIALIZED|EXTERNAL|DYNAMIC|ICEBERG|HYBRID|UNLOGGED|FOREIGN|STREAMING)\s+)*" - r"(?:TABLE|VIEW)\s+" + r"(?:(?:OR\s++REPLACE|LOCAL|GLOBAL|TEMP|TEMPORARY|TRANSIENT|VOLATILE|SECURE|RECURSIVE|" + r"MATERIALIZED|EXTERNAL|DYNAMIC|ICEBERG|HYBRID|UNLOGGED|FOREIGN|STREAMING)\s++)*+" + r"(?:TABLE|VIEW)\s++" ) _CREATE_RELATION: re.Pattern[str] = re.compile( - rf"^CREATE\s+{_RELATION_KIND}(?:IF\s+NOT\s+EXISTS\s+)?{_QUALIFIED_NAME}", re.IGNORECASE + rf"^CREATE\s++{_RELATION_KIND}(?:IF\s++NOT\s++EXISTS\s++)?{_QUALIFIED_NAME}", re.IGNORECASE ) _DROP_RELATION: re.Pattern[str] = re.compile( - rf"^DROP\s+{_RELATION_KIND}(?:IF\s+EXISTS\s+)?{_QUALIFIED_NAME}", re.IGNORECASE + rf"^DROP\s++{_RELATION_KIND}(?:IF\s++EXISTS\s++)?{_QUALIFIED_NAME}", re.IGNORECASE ) _ALTER_RELATION: re.Pattern[str] = re.compile( - rf"^ALTER\s+{_RELATION_KIND}(?:IF\s+EXISTS\s+)?{_QUALIFIED_NAME}", re.IGNORECASE + rf"^ALTER\s++{_RELATION_KIND}(?:IF\s++EXISTS\s++)?{_QUALIFIED_NAME}", re.IGNORECASE ) _COPY_INTO_RELATION: re.Pattern[str] = re.compile( - rf"^COPY\s+INTO\s+{_QUALIFIED_NAME}", re.IGNORECASE + rf"^COPY\s++INTO\s++{_QUALIFIED_NAME}", re.IGNORECASE ) -_SELECT_INTO_RELATION: re.Pattern[str] = re.compile(rf"\bINTO\s+{_QUALIFIED_NAME}", re.IGNORECASE) +_SELECT_INTO_RELATION: re.Pattern[str] = re.compile(rf"\bINTO\s++{_QUALIFIED_NAME}", re.IGNORECASE) _SELECT_INTO_KEYWORD: re.Pattern[str] = re.compile(r"\bINTO\b", re.IGNORECASE) _SECOND_RELATION: re.Pattern[str] = re.compile( - rf"^\s+(?:RENAME\s+TO|SWAP\s+WITH)\s+{_QUALIFIED_NAME}", re.IGNORECASE + rf"^\s++(?:RENAME\s++TO|SWAP\s++WITH)\s++{_QUALIFIED_NAME}", re.IGNORECASE ) _FOLLOWING_WORDS: dict[re.Pattern[str], re.Pattern[str]] = { _CREATE_RELATION: re.compile( - r"^\s*(?:$|\(|(?:AS|CLONE|LIKE|COPY|USING|COMMENT|CLUSTER|PARTITION|OPTIONS|WITH|" + r"^\s*+(?:$|\(|(?:AS|CLONE|SHALLOW|DEEP|LIKE|COPY|USING|COMMENT|CLUSTER|PARTITION|OPTIONS|WITH|" r"TBLPROPERTIES|LOCATION|DATA_RETENTION_TIME_IN_DAYS|CHANGE_TRACKING)\b)", re.IGNORECASE, ), - _DROP_RELATION: re.compile(r"^\s*(?:CASCADE|RESTRICT|PURGE)?\s*$", re.IGNORECASE), + _DROP_RELATION: re.compile(r"^\s*+(?:CASCADE|RESTRICT|PURGE)?+\s*+$", re.IGNORECASE), _ALTER_RELATION: re.compile( - r"^\s+(?:ADD|DROP|ALTER|MODIFY|RENAME|SWAP|SET|UNSET|CLUSTER|RECLUSTER|OWNER|" + r"^\s++(?:ADD|DROP|ALTER|MODIFY|RENAME|SWAP|SET|UNSET|CLUSTER|RECLUSTER|OWNER|" r"SUSPEND|RESUME|REFRESH)\b", re.IGNORECASE, ), - _COPY_INTO_RELATION: re.compile(r"^\s+FROM\b", re.IGNORECASE), - _SELECT_INTO_RELATION: re.compile(r"^\s+FROM\b", re.IGNORECASE), - _SECOND_RELATION: re.compile(r"^\s*$"), + _COPY_INTO_RELATION: re.compile(r"^\s++FROM\b", re.IGNORECASE), + _SELECT_INTO_RELATION: re.compile(r"^\s++FROM\b", re.IGNORECASE), + _SECOND_RELATION: re.compile(r"^\s*+$"), } _CREATE_SCHEMA_IF_NOT_EXISTS: re.Pattern[str] = re.compile( - r"^CREATE\s+SCHEMA\s+IF\s+NOT\s+EXISTS\s+[^\s;]+\s*$", re.IGNORECASE + r"^CREATE\s++SCHEMA\s++IF\s++NOT\s++EXISTS\s++[^\s;]++\s*+$", re.IGNORECASE ) _RELATION_NEUTRAL_OBJECT: re.Pattern[str] = re.compile( - r"^(?:CREATE|DROP|ALTER)\s+(?:OR\s+REPLACE\s+)?(?:(?:SECURE|TEMP|TEMPORARY|AGGREGATE)\s+)*" - r"(?:FUNCTION|PROCEDURE|SEQUENCE|STAGE|FILE\s+FORMAT|INDEX|UNIQUE\s+INDEX|WAREHOUSE|MACRO)\b", + r"^(?:CREATE|DROP|ALTER)\s++(?:OR\s++REPLACE\s++)?(?:(?:SECURE|TEMP|TEMPORARY|AGGREGATE)\s++)*+" + r"(?:FUNCTION|PROCEDURE|SEQUENCE|STAGE|FILE\s++FORMAT|INDEX|UNIQUE\s++INDEX|WAREHOUSE|MACRO)\b", re.IGNORECASE, ) -_RENAMING_ALTER: re.Pattern[str] = re.compile(r"^\s+(?:RENAME\s+TO|SWAP\s+WITH)\b", re.IGNORECASE) +_RENAMING_ALTER: re.Pattern[str] = re.compile( + r"^\s++(?:RENAME\s++TO|SWAP\s++WITH)\b", re.IGNORECASE +) _QUERY_STATEMENTS: frozenset[str] = frozenset({"SELECT", "WITH"}) _READ_OR_DATA_STATEMENTS: frozenset[str] = frozenset( { @@ -148,9 +129,19 @@ def statement_metadata_effect(sql: str) -> StatementMetadataEffect: """Return which cached relation names ``sql`` may have changed; ambiguity changes all.""" - effects: set[StatementMetadataEffect] = { - _code_effect(code=_code_text(sql=sql, lexer=lexer)) for lexer in _LEXERS - } + backslash_modes: tuple[bool, ...] = (False, True) if _BACKSLASH in sql else (False,) + bracket_modes: tuple[bool, ...] = (False, True) if _OPEN_BRACKET in sql else (False,) + effects: set[StatementMetadataEffect] = set() + backslash_escapes: bool + for backslash_escapes in backslash_modes: + effects.update( + _code_effect( + code=_code_text( + sql=sql, backslash_escapes=backslash_escapes, bracket_identifiers=brackets + ) + ) + for brackets in bracket_modes + ) if len(effects) != 1: return _ALL_RELATIONS return effects.pop() @@ -162,23 +153,109 @@ def relation_name_key(name: str) -> str: return _strip_identifier_quotes(name).casefold() -def qualified_relation_name_key(qualified: str) -> str: - """Return the invalidation key for the last part of a possibly qualified relation name.""" +def qualified_relation_name_key(qualified: str) -> str | None: + """Return the key for the last part of a qualified name, or None when it cannot be parsed.""" - return _unqualified_name_key(qualified) + if _LOOSE_QUALIFIED_NAME.match(qualified) is None: + return None + parts: list[str] = re.findall(_LOOSE_IDENTIFIER_PART, qualified) + return _part_name_key(parts[-1]) + + +def _code_text(*, sql: str, backslash_escapes: bool, bracket_identifiers: bool) -> str | None: + pieces: list[str] = [] + position: int = 0 + while True: + start: re.Match[str] | None = _LEXEME_START.search(sql, position) + if start is None: + pieces.append(sql[position:]) + return "".join(pieces) + pieces.append(sql[position : start.start()]) + lexeme: tuple[str, int] | None = _read_lexeme( + sql=sql, + start=start, + backslash_escapes=backslash_escapes, + bracket_identifiers=bracket_identifiers, + ) + if lexeme is None: + return None + pieces.append(lexeme[0]) + position = lexeme[1] + + +def _read_lexeme( + *, sql: str, start: re.Match[str], backslash_escapes: bool, bracket_identifiers: bool +) -> tuple[str, int] | None: + token: str = start.group(0) + begin: int = start.start() + if token == _LINE_COMMENT_START: + line_end: int = sql.find(_NEWLINE, begin) + return _COMMENT_REPLACEMENT, (len(sql) if line_end < 0 else line_end) + if token == _BLOCK_COMMENT_START: + return _closed( + text=_COMMENT_REPLACEMENT, end=sql.find(_BLOCK_COMMENT_END, begin + 2), width=2 + ) + if token == _SINGLE_QUOTE: + return _closed( + text=_EMPTY_LITERAL, + end=_literal_end(sql=sql, begin=begin, backslash_escapes=backslash_escapes), + width=1, + ) + if token == _DOLLAR: + return _dollar_lexeme(sql=sql, begin=begin) + if token == _OPEN_BRACKET and not bracket_identifiers: + return token, begin + 1 + closing: str = _CLOSE_BRACKET if token == _OPEN_BRACKET else token + end: int = _quoted_end(sql=sql, begin=begin, quote=closing) + if end < 0: + return None + return sql[begin : end + 1], end + 1 -def _code_text(*, sql: str, lexer: re.Pattern[str]) -> str | None: - if _UNTERMINATED_MARKER in sql: - return None - code: str = lexer.sub(_replace_lexeme, sql) - if _UNTERMINATED_MARKER in code: +def _closed(*, text: str, end: int, width: int) -> tuple[str, int] | None: + if end < 0: return None - return code + return text, end + width + + +def _dollar_lexeme(*, sql: str, begin: int) -> tuple[str, int] | None: + tag: re.Match[str] | None = _DOLLAR_TAG.match(sql, begin) + if tag is None: + return _DOLLAR, begin + 1 + return _closed( + text=_EMPTY_LITERAL, end=sql.find(tag.group(0), tag.end()), width=len(tag.group(0)) + ) + +def _literal_end(*, sql: str, begin: int, backslash_escapes: bool) -> int: + position: int = begin + 1 + while True: + stop: int = _literal_stop(sql=sql, position=position, backslash_escapes=backslash_escapes) + if stop < 0: + return stop + if sql[stop] == _BACKSLASH: + position = stop + 2 + continue + if sql.startswith(_EMPTY_LITERAL, stop): + position = stop + 2 + continue + return stop -def _replace_lexeme(match: re.Match[str]) -> str: - return _LEXEME_REPLACEMENTS.get(match.lastgroup or "", match.group(0)) + +def _literal_stop(*, sql: str, position: int, backslash_escapes: bool) -> int: + if not backslash_escapes: + return sql.find(_SINGLE_QUOTE, position) + stop: re.Match[str] | None = _BACKSLASH_LITERAL_STOP.search(sql, position) + return -1 if stop is None else stop.start() + + +def _quoted_end(*, sql: str, begin: int, quote: str) -> int: + position: int = begin + 1 + while True: + end: int = sql.find(quote, position) + if end < 0 or quote == _CLOSE_BRACKET or not sql.startswith(quote * 2, end): + return end + position = end + 2 def _code_effect(*, code: str | None) -> StatementMetadataEffect: @@ -246,10 +323,13 @@ def _unambiguous_names( def _unqualified_name_key(qualified: str) -> str: parts: list[str] = re.findall(_IDENTIFIER_PART, qualified) - last: str = parts[-1] - if last.startswith(BACKTICK): - return _strip_identifier_quotes(last).rsplit(".", 1)[-1].casefold() - return relation_name_key(last) + return _part_name_key(parts[-1]) + + +def _part_name_key(part: str) -> str: + if part.startswith(BACKTICK): + return _strip_identifier_quotes(part).rsplit(".", 1)[-1].casefold() + return relation_name_key(part) def _strip_identifier_quotes(part: str) -> str: diff --git a/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py b/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py index d1b56cfd4c..6f745f81f2 100644 --- a/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py +++ b/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py @@ -43,8 +43,12 @@ def lookup[ValueT]( def invalidate_relations(self, *, qualified_names: tuple[str, ...]) -> None: """Evict relations changed outside SQL statements, such as by a warehouse copy job.""" - names: frozenset[str] = frozenset(map(qualified_relation_name_key, qualified_names)) + keys: tuple[str | None, ...] = tuple(map(qualified_relation_name_key, qualified_names)) + names: frozenset[str] = frozenset(key for key in keys if key is not None) with self._lock: + if None in keys: + self._invalidate_all() + return self._invalidate_names(names=names) def observe_statement(self, *, sql: str) -> None: diff --git a/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/_test_types.py b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/_test_types.py index 93508cddb1..863d3c561b 100644 --- a/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/_test_types.py +++ b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/_test_types.py @@ -12,3 +12,21 @@ class StatementMetadataEffectTestCase: expected_relation_names: frozenset[str] expected_invalidates_all: bool expected_ends_transaction: bool = False + + +@dataclass(frozen=True) +class QualifiedNameKeyTestCase: + """One adapter-supplied qualified name and its invalidation key, if parseable.""" + + description: str + qualified: str + expected_key: str | None + + +@dataclass(frozen=True) +class AdversarialStatementTestCase: + """Hostile SQL that statement scanners must still handle in bounded time.""" + + description: str + sql: str + expected_max_seconds: float = 2.0 diff --git a/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_adversarial_input_runtime.py b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_adversarial_input_runtime.py new file mode 100644 index 0000000000..aa7e533133 --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_adversarial_input_runtime.py @@ -0,0 +1,99 @@ +"""Statement classification and the Snowflake session hook stay linear on adversarial SQL.""" + +from __future__ import annotations + +import time + +import pytest + +from sqlbuild.adapter.relations._helpers.ddl_effects import statement_metadata_effect +from sqlbuild.adapters.snowflake._helpers.session_context import may_switch_session_database +from tests.unit.src.sqlbuild.adapter.relations._helpers.ddl_effects._test_types import ( + AdversarialStatementTestCase, +) + +_REALISTIC_CTAS: str = ( + "CREATE OR REPLACE TABLE analytics.marts.orders AS SELECT " + + ", ".join(f'CAST(col_{index} AS VARCHAR) AS "c{index}" -- note\n' for index in range(25_000)) + + " FROM raw.orders WHERE status = 'it''s' AND body = $$x$$" +) + + +@pytest.mark.parametrize( + "test_case", + [ + AdversarialStatementTestCase( + description="10k leading spaces", + sql=" " * 10_000 + "SELECT 1", + ), + AdversarialStatementTestCase( + description="100k leading spaces", + sql=" " * 100_000 + "SELECT 1", + ), + AdversarialStatementTestCase( + description="100k mixed leading whitespace", + sql="\n\t " * 33_000 + "USE DATABASE analytics", + ), + AdversarialStatementTestCase( + description="repeated unterminated block comments", + sql="/*" * 20_000, + ), + AdversarialStatementTestCase( + description="spaced unterminated block comments", + sql="SELECT " + "/* " * 20_000, + ), + AdversarialStatementTestCase( + description="many dollar tag openers", + sql="SELECT " + "$a$ " * 10_000, + ), + AdversarialStatementTestCase( + description="many unclosed dollar tags", + sql="SELECT " + "$tag" * 10_000, + ), + AdversarialStatementTestCase( + description="drop with long whitespace and junk", + sql="DROP TABLE marts.orders" + " " * 20_000 + "x", + ), + AdversarialStatementTestCase( + description="create with long whitespace and junk", + sql="CREATE TABLE marts.orders" + " " * 20_000 + "x", + ), + AdversarialStatementTestCase( + description="alter with long whitespace and junk", + sql="ALTER TABLE marts.orders" + " " * 20_000 + "x", + ), + AdversarialStatementTestCase( + description="many single quotes", + sql="SELECT " + "'" * 20_001, + ), + AdversarialStatementTestCase( + description="many double quotes", + sql="SELECT " + '"' * 20_001, + ), + AdversarialStatementTestCase( + description="many line comments", + sql="-- c\n" * 20_000 + "SELECT 1", + ), + AdversarialStatementTestCase( + description="1 MB create table as select", + sql=_REALISTIC_CTAS, + ), + ], + ids=lambda case: case.description, +) +def test_given_adversarial_sql_when_scanning_then_completes_within_bound( + test_case: AdversarialStatementTestCase, +) -> None: + classifier_started: float = time.perf_counter() + _ = statement_metadata_effect(test_case.sql) + classifier_seconds: float = time.perf_counter() - classifier_started + hook_started: float = time.perf_counter() + _ = may_switch_session_database(test_case.sql) + hook_seconds: float = time.perf_counter() - hook_started + + assert classifier_seconds < test_case.expected_max_seconds + assert hook_seconds < test_case.expected_max_seconds + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) diff --git a/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py index 676cc41e61..ae3d75af7a 100644 --- a/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py +++ b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py @@ -4,9 +4,13 @@ import pytest -from sqlbuild.adapter.relations._helpers.ddl_effects import statement_metadata_effect +from sqlbuild.adapter.relations._helpers.ddl_effects import ( + qualified_relation_name_key, + statement_metadata_effect, +) from sqlbuild.adapter.relations.models import StatementMetadataEffect from tests.unit.src.sqlbuild.adapter.relations._helpers.ddl_effects._test_types import ( + QualifiedNameKeyTestCase, StatementMetadataEffectTestCase, ) @@ -106,6 +110,24 @@ expected_relation_names=frozenset({"orders_history"}), expected_invalidates_all=False, ), + StatementMetadataEffectTestCase( + description="databricks shallow clone names only the new relation", + sql="CREATE OR REPLACE TABLE main.marts.orders__prev SHALLOW CLONE main.marts.orders", + expected_relation_names=frozenset({"orders__prev"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="databricks deep clone names only the new relation", + sql="CREATE TABLE IF NOT EXISTS main.marts.orders__copy DEEP CLONE main.marts.orders", + expected_relation_names=frozenset({"orders__copy"}), + expected_invalidates_all=False, + ), + StatementMetadataEffectTestCase( + description="unquoted name starting with a digit keeps its digits", + sql="CREATE OR REPLACE TABLE ds.2024_orders AS SELECT 1 AS id", + expected_relation_names=frozenset({"2024_orders"}), + expected_invalidates_all=False, + ), StatementMetadataEffectTestCase( description="copy into a table may evolve its columns", sql="COPY INTO raw.orders FROM @landing/orders", @@ -238,5 +260,42 @@ def test_given_transaction_end_statement_when_classifying_then_reports_possible_ assert effect.ends_transaction is test_case.expected_ends_transaction +@pytest.mark.parametrize( + "test_case", + [ + QualifiedNameKeyTestCase( + description="unquoted part starting with a digit keeps its digits", + qualified="my-proj.ds.2024_orders", + expected_key="2024_orders", + ), + QualifiedNameKeyTestCase( + description="backticked path keys its last segment", + qualified="`my-proj.ds.Orders`", + expected_key="orders", + ), + QualifiedNameKeyTestCase( + description="quoted dotted part keeps its dot", + qualified='"ANALYTICS"."MARTS"."my.table"', + expected_key="my.table", + ), + QualifiedNameKeyTestCase( + description="trailing text is unparseable", + qualified="ds.orders extra", + expected_key=None, + ), + QualifiedNameKeyTestCase( + description="unquoted non-ascii part is unparseable", + qualified="ds.café", + expected_key=None, + ), + ], + ids=lambda case: case.description, +) +def test_given_qualified_name_when_keying_then_uses_last_part_or_reports_unparseable( + test_case: QualifiedNameKeyTestCase, +) -> None: + assert qualified_relation_name_key(test_case.qualified) == test_case.expected_key + + if __name__ == "__main__": pytest.main([__file__, "-vv"])