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/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py b/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py new file mode 100644 index 0000000000..f78d68623c --- /dev/null +++ b/src/sqlbuild/adapter/relations/_helpers/ddl_effects.py @@ -0,0 +1,343 @@ +"""Classify executed SQL by its effect on cached relation existence and columns.""" + +from __future__ import annotations + +import re + +from sqlbuild.adapter.relations.constants import ( + BACKTICK, + DOUBLE_QUOTE, + IDENTIFIER_QUOTE_PAIRS, + QUOTED_IDENTIFIER_MIN_LENGTH, + STATEMENT_SEPARATOR, +) +from sqlbuild.adapter.relations.models import StatementMetadataEffect + +_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-z0-9_][A-Za-z0-9_$]*+)' +_QUALIFIED_NAME: str = ( + 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++" +) +_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 +) +_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|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), + _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*+$", 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( + { + "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; ambiguity changes all.""" + + 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() + + +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 | None: + """Return the key for the last part of a qualified name, or None when it cannot be parsed.""" + + 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 _closed(*, text: str, end: int, width: int) -> tuple[str, int] | None: + if end < 0: + return None + 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 _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: + 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 + 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 _select_into_effect(*, statement: str) -> StatementMetadataEffect: + if _SELECT_INTO_KEYWORD.search(statement) is None: + return _NO_EFFECT + 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) + 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) + 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: + 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..6f745f81f2 --- /dev/null +++ b/src/sqlbuild/adapter/relations/classes/relation_metadata_cache.py @@ -0,0 +1,108 @@ +"""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 ( + qualified_relation_name_key, + 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 invalidate_relations(self, *, qualified_names: tuple[str, ...]) -> None: + """Evict relations changed outside SQL statements, such as by a warehouse copy job.""" + + 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: + """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..53d4ad02c5 100644 --- a/src/sqlbuild/adapter/relations/constants.py +++ b/src/sqlbuild/adapter/relations/constants.py @@ -4,3 +4,10 @@ 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" +BACKTICK: str = "`" 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/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/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/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/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/_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..7555d23333 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -118,16 +118,19 @@ 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 ( BASE_TABLE_METADATA_TYPE, + CURRENT_DATABASE_ATTRIBUTE, EXTERNAL_BROWSER_AUTHENTICATOR, MAX_STATEMENT_TIMEOUT_SECONDS, MFA_AUTHENTICATOR, OAUTH_AUTHORIZATION_CODE_AUTHENTICATOR, SECONDARY_ROLES_ALL, SECONDARY_ROLES_NONE, + SHOW_LIKE_ESCAPE_CHARACTER, STATUS_COLUMN_NAME, SUCCESS_STATUS_TOKENS, TRUE_METADATA_VALUE, @@ -149,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" @@ -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)] @@ -1985,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() @@ -1995,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( @@ -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/classes/snowflake_connection.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py index 19349018df..1ee72c1948 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_connection.py @@ -2,7 +2,9 @@ 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 class _SnowflakeConnection: @@ -21,4 +23,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 may_switch_session_database(sql): + _ = 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 01c94d449e..69692ed103 100644 --- a/src/sqlbuild/adapters/snowflake/constants.py +++ b/src/sqlbuild/adapters/snowflake/constants.py @@ -22,3 +22,5 @@ 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 = "\\" +CURRENT_DATABASE_ATTRIBUTE: str = "_sqlbuild_current_database" 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/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/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..863d3c561b --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/_test_types.py @@ -0,0 +1,32 @@ +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 + + +@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 new file mode 100644 index 0000000000..ae3d75af7a --- /dev/null +++ b/tests/unit/src/sqlbuild/adapter/relations/_helpers/ddl_effects/test_statement_metadata_effect.py @@ -0,0 +1,301 @@ +"""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 ( + 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, +) + +_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="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", + 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="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", + 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", + "CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$ SELECT 1; SELECT 2 $$", + "INSERT INTO marts.notes VALUES ('drop table u; --')", + ) + ) + ], + 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", " "), + ("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, +) +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 + + +@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"]) 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"]) 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"]) 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..3eb644e84e 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,48 @@ 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 + + +@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 b4fc0d1b2d..2e774c0b65 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( @@ -49,6 +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'^(?s:.*?)USE DATABASE "?(?P[^";]+)"?(?s:.*)$' +) _SHOW_SCHEMAS_PATTERN: re.Pattern[str] = re.compile( r"^SHOW SCHEMAS LIKE '(?P(?:[^']|'')*)' IN DATABASE \"(?P[^\"]*)\"$" ) @@ -59,6 +63,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 +113,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", @@ -239,6 +251,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.""" @@ -295,6 +320,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), ) @@ -351,6 +377,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} @@ -378,12 +412,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 +433,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 +499,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 +593,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 +1006,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_session_database_switch.py b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py new file mode 100644 index 0000000000..53ad14f0bb --- /dev/null +++ b/tests/unit/src/sqlbuild/adapters/snowflake/inspection/test_session_database_switch.py @@ -0,0 +1,76 @@ +"""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="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", + 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"]) 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", [ 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