diff --git a/docs/cli.md b/docs/cli.md index 568a2b3..efcebfe 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -105,6 +105,7 @@ snowcap export --all | `--exclude ` | Exclude resource types (used with --all) | | `--out ` | Write exported config to a file | | `--format [json\|yml]` | Output format (default: yml) | +| `--use-account-usage / --no-use-account-usage` | Read from `SNOWFLAKE.ACCOUNT_USAGE` instead of `SHOW` (default: on). `SHOW ... IN ACCOUNT` is capped at 10,000 rows; `ACCOUNT_USAGE` lags live state by up to ~2 hours | **Examples:** diff --git a/snowcap/blueprint.py b/snowcap/blueprint.py index 33d83c2..6912d25 100644 --- a/snowcap/blueprint.py +++ b/snowcap/blueprint.py @@ -1666,20 +1666,22 @@ def fetch_remote_state(self, session, manifest: Manifest) -> State: if self._config.sync_resources: urns = [item for item in manifest.urns if item.resource_type not in self._config.sync_resources] for resource_type in self._config.sync_resources: - # Future grants are read in full whenever grants are synced, and neither - # the query nor the set of roles queried is narrowed to what the manifest - # declares. Narrowing would be sound for a plan that only creates, but - # syncing a resource type means removing what is not declared, and a future - # grant absent from config is precisely what has to be found. - # - # Skipping the query when the manifest declared no future grants kept the - # ones already in Snowflake out of remote state, so sync could not propose - # dropping them -- unseen rather than deliberately kept, with nothing in - # the plan to say so. Migrating a config from ALL plus FUTURE pairs to - # inherited grants removes the last future grant and hit exactly that: 26 - # orphaned future grants, zero drops, no warning. - list_kwargs: dict[str, Any] = {} + # Sync drops whatever remote state omits, so the listing must read the + # source the config asked for. list_resource drops unsupported kwargs. + list_kwargs: dict[str, Any] = {"use_account_usage": self._config.use_account_usage} if resource_type == ResourceType.GRANT: + # Future grants are read in full whenever grants are synced, and neither + # the query nor the set of roles queried is narrowed to what the manifest + # declares. Narrowing would be sound for a plan that only creates, but + # syncing a resource type means removing what is not declared, and a future + # grant absent from config is precisely what has to be found. + # + # Skipping the query when the manifest declared no future grants kept the + # ones already in Snowflake out of remote state, so sync could not propose + # dropping them -- unseen rather than deliberately kept, with nothing in + # the plan to say so. Migrating a config from ALL plus FUTURE pairs to + # inherited grants removes the last future grant and hit exactly that: 26 + # orphaned future grants, zero drops, no warning. list_kwargs["include_future_grants"] = True for fqn in data_provider.list_resource(session, resource_label_for_type(resource_type), **list_kwargs): if self._config.scope == BlueprintScope.DATABASE and fqn.database != self._config.database: diff --git a/snowcap/cli.py b/snowcap/cli.py index a771cff..51a58ee 100644 --- a/snowcap/cli.py +++ b/snowcap/cli.py @@ -358,7 +358,15 @@ def apply( ) @click.option("--out", type=str, help="Write exported config to a file", metavar="") @click.option("--format", type=click.Choice(["json", "yml"]), default="yml", help="Output format") -def export(resources, export_all, exclude_resources, out, format) -> None: +@click.option( + "--use-account-usage/--no-use-account-usage", + default=True, + help=( + "Read from SNOWFLAKE.ACCOUNT_USAGE instead of SHOW. Required past the 10,000-row " + "SHOW cap, but lags live state by up to ~2 hours" + ), +) +def export(resources, export_all, exclude_resources, out, format, use_account_usage) -> None: """ Generate a resource config for existing Snowflake resources @@ -388,9 +396,9 @@ def export(resources, export_all, exclude_resources, out, format) -> None: resource_config: dict[str, Any] = {} if resources: - resource_config = export_resources(include=resources) + resource_config = export_resources(include=resources, use_account_usage=use_account_usage) elif export_all: - resource_config = export_resources(exclude=exclude_resources) + resource_config = export_resources(exclude=exclude_resources, use_account_usage=use_account_usage) else: raise diff --git a/snowcap/client.py b/snowcap/client.py index 32a7d9b..e8415f9 100644 --- a/snowcap/client.py +++ b/snowcap/client.py @@ -35,10 +35,18 @@ # Track queries currently being executed to prevent duplicate execution _PENDING_QUERIES: dict[tuple[str, str], threading.Event] = {} +# Invalidation callbacks for caches derived from _EXECUTION_CACHE, registered by the modules +# that index cached rows -- this module cannot import them, they import it. +_CACHE_RESET_HOOKS: list[Callable[[], None]] = [] + + +def register_cache_reset_hook(hook: Callable[[], None]) -> None: + _CACHE_RESET_HOOKS.append(hook) + def reset_cache(): """ - Reset the SQL execution cache. + Reset the SQL execution cache and every cache derived from it. This clears cached query results so subsequent queries will re-execute. Note: This does NOT clear ACCOUNT_USAGE caches, which are designed to persist @@ -47,6 +55,8 @@ def reset_cache(): """ global _EXECUTION_CACHE _EXECUTION_CACHE = {} + for hook in _CACHE_RESET_HOOKS: + hook() def execute( diff --git a/snowcap/data_provider.py b/snowcap/data_provider.py index a80cac4..edf170b 100644 --- a/snowcap/data_provider.py +++ b/snowcap/data_provider.py @@ -3,6 +3,7 @@ import json import logging import sys +import threading from typing import Any, Optional, TypedDict, Union import pytz @@ -26,6 +27,7 @@ UNSUPPORTED_FEATURE, execute, execute_in_parallel, + register_cache_reset_hook, ) from .enums import ( INHERITED_GRANTS_FEATURE_FLAG, @@ -397,15 +399,44 @@ def _drop_inherited_grants(rows: list[dict[str, Any]], context: str) -> list[dic return kept -def _fetch_grant_to_role( +# (session id, role, role type, grant type) -> {(granted_on label, privilege, name): grant}. +_GRANT_LOOKUP_INDEX_CACHE: dict[tuple, dict[tuple, dict[str, Any]]] = {} +_GRANT_LOOKUP_INDEX_LOCK = threading.Lock() + + +def _reset_grant_lookup_index() -> None: + with _GRANT_LOOKUP_INDEX_LOCK: + _GRANT_LOOKUP_INDEX_CACHE.clear() + + +# Built from cacheable SHOW GRANTS rows, so it expires with the SQL execution cache. +register_cache_reset_hook(_reset_grant_lookup_index) + + +def _grant_name_key(name: str) -> str: + """ + Canonical key for a grant target name, matching ResourceName equality semantics. + + ResourceName is unsafe as a dict key: __hash__ is hash(str(self)), which keeps the + quotes, while __eq__ treats quoted "FOO" and unquoted FOO as equal. "Exact if quoted, + upper-cased if not" reproduces __eq__ across all four combinations. + """ + rendered = str(ResourceName(name)) + return rendered[1:-1] if rendered.startswith('"') else rendered + + +def _grant_lookup_index( session: SnowflakeConnection, grant_type: GrantType, role: ResourceName, - granted_on: str, - on_name: str, - privilege: str, - role_type: ResourceType = ResourceType.ROLE, -): + role_type: ResourceType, +) -> dict[tuple, dict[str, Any]]: + """Return (building it first if needed) the (granted_on, privilege, name) index for one role.""" + cache_key = (id(session), str(role), role_type, grant_type) + index = _GRANT_LOOKUP_INDEX_CACHE.get(cache_key) + if index is not None: + return index + if grant_type == GrantType.FUTURE: if role_type == ResourceType.DATABASE_ROLE: grants = _show_future_grants_to_database_role(session, str(role), cacheable=True) @@ -413,18 +444,39 @@ def _fetch_grant_to_role( grants = _show_future_grants_to_role(session, role, cacheable=True) else: grants = _show_grants_to_role(session, role, role_type=role_type, cacheable=True) - # Compare object types the way list_grants does. Snowflake sometimes reports a grant - # against a different name than the one its DDL uses -- CORTEX_AGENT_SERVER for what - # GRANT calls an MCP SERVER -- so a raw string comparison never matches the declared - # grant, and the plan proposes creating it on every run. - wanted_type = _granted_on_label(granted_on) - for grant in grants: - name = "ACCOUNT" if grant["granted_on"] == "ACCOUNT" else grant["name"] - # Use ResourceName for comparison to handle quoted identifiers correctly - name_matches = ResourceName(name) == ResourceName(on_name) if name != "ACCOUNT" else name == on_name - if _granted_on_label(grant["granted_on"]) == wanted_type and grant["privilege"] == privilege and name_matches: - return grant - return None + + with _GRANT_LOOKUP_INDEX_LOCK: + # Re-check: another thread may have built it while we waited for the lock. + index = _GRANT_LOOKUP_INDEX_CACHE.get(cache_key) + if index is None: + index = {} + for grant in grants: + is_account = grant["granted_on"] == "ACCOUNT" + name = "ACCOUNT" if is_account else grant["name"] + key = ( + # Snowflake reports type synonyms (CORTEX_AGENT_SERVER for MCP SERVER). + _granted_on_label(grant["granted_on"]), + grant["privilege"], + name if is_account else _grant_name_key(name), + ) + # First match wins, as the linear scan did, when a grant is duplicated. + index.setdefault(key, grant) + _GRANT_LOOKUP_INDEX_CACHE[cache_key] = index + return index + + +def _fetch_grant_to_role( + session: SnowflakeConnection, + grant_type: GrantType, + role: ResourceName, + granted_on: str, + on_name: str, + privilege: str, + role_type: ResourceType = ResourceType.ROLE, +): + index = _grant_lookup_index(session, grant_type, role, role_type) + name_key = on_name if granted_on == "ACCOUNT" else _grant_name_key(on_name) + return index.get((_granted_on_label(granted_on), privilege, name_key)) def _filter_result(result, **kwargs): @@ -834,13 +886,7 @@ def _show_all_grants_to_role( if role_type == ResourceType.ROLE: session_id = id(session) if session_id in _ACCOUNT_USAGE_GRANTS_CACHE: - # Filter cached grants by role name (case-insensitive) - role_upper = str(role).upper() - filtered_grants = [ - grant - for grant in _ACCOUNT_USAGE_GRANTS_CACHE[session_id] - if grant["grantee_name"].upper() == role_upper and grant["granted_to"] == "ROLE" - ] + filtered_grants = _grants_by_role_index(session_id).get(str(role).upper(), []) logger.debug(f"Using ACCOUNT_USAGE cache for grants to role {role} ({len(filtered_grants)} grants)") return filtered_grants @@ -1204,6 +1250,30 @@ def fetch_role_privileges( # Stores the normalized grant list from GRANTS_TO_USERS _ACCOUNT_USAGE_USER_GRANTS_CACHE: dict[int, list[dict[str, Any]]] = {} +# Role-name index over _ACCOUNT_USAGE_GRANTS_CACHE: session id, then upper-cased grantee. +_ACCOUNT_USAGE_GRANTS_BY_ROLE_CACHE: dict[int, dict[str, list[dict[str, Any]]]] = {} +_ACCOUNT_USAGE_GRANTS_INDEX_LOCK = threading.Lock() + + +def _grants_by_role_index(session_id: int) -> dict[str, list[dict[str, Any]]]: + """Return (building it first if needed) the role -> grants index for a session's grant cache.""" + index = _ACCOUNT_USAGE_GRANTS_BY_ROLE_CACHE.get(session_id) + if index is not None: + return index + with _ACCOUNT_USAGE_GRANTS_INDEX_LOCK: + # Re-check: another thread may have built it while we waited for the lock. + index = _ACCOUNT_USAGE_GRANTS_BY_ROLE_CACHE.get(session_id) + if index is None: + index = {} + for grant in _ACCOUNT_USAGE_GRANTS_CACHE.get(session_id, []): + if grant["granted_to"] != "ROLE": + continue + index.setdefault(grant["grantee_name"].upper(), []).append(grant) + _ACCOUNT_USAGE_GRANTS_BY_ROLE_CACHE[session_id] = index + logger.debug(f"Built ACCOUNT_USAGE grant index: {len(index)} roles") + return index + + # Tracks sessions whose GRANTS_TO_ROLES view has no IS_INHERITED column (keyed by session # id). Absent means "assume the column exists"; the first query proves it either way. _ACCOUNT_USAGE_INHERITED_COLUMN: dict[int, bool] = {} @@ -1212,6 +1282,9 @@ def fetch_role_privileges( # parameter could not be read. _INHERITED_GRANTS_ENABLED_CACHE: dict[int, Optional[bool]] = {} +# Sessions already warned that ACCOUNT_USAGE listing lags live state. +_ACCOUNT_USAGE_STALENESS_WARNED: set[int] = set() + def reset_account_usage_caches() -> None: """ @@ -1223,6 +1296,8 @@ def reset_account_usage_caches() -> None: global _ACCOUNT_USAGE_ACCESS_CACHE, _ACCOUNT_USAGE_FALLBACK_CACHE global _ACCOUNT_USAGE_GRANTS_CACHE, _ACCOUNT_USAGE_USER_GRANTS_CACHE global _ACCOUNT_USAGE_INHERITED_COLUMN, _INHERITED_GRANTS_ENABLED_CACHE + global _ACCOUNT_USAGE_GRANTS_BY_ROLE_CACHE, _GRANT_LOOKUP_INDEX_CACHE + global _ACCOUNT_USAGE_STALENESS_WARNED _ACCOUNT_USAGE_ACCESS_CACHE.clear() _ACCOUNT_USAGE_FALLBACK_CACHE.clear() @@ -1230,6 +1305,10 @@ def reset_account_usage_caches() -> None: _ACCOUNT_USAGE_USER_GRANTS_CACHE.clear() _ACCOUNT_USAGE_INHERITED_COLUMN.clear() _INHERITED_GRANTS_ENABLED_CACHE.clear() + _ACCOUNT_USAGE_STALENESS_WARNED.clear() + # Derived indexes must not outlive the rows they were built from. + _ACCOUNT_USAGE_GRANTS_BY_ROLE_CACHE.clear() + _reset_grant_lookup_index() def _mark_account_usage_fallback(session: SnowflakeConnection) -> None: @@ -4858,7 +4937,77 @@ def list_shares(session: SnowflakeConnection) -> list[FQN]: return shares -def list_stages(session: SnowflakeConnection) -> list[FQN]: +def _warn_account_usage_staleness(session: SnowflakeConnection) -> None: + """Say once per session that a listing read from ACCOUNT_USAGE may be behind live state.""" + session_id = id(session) + if session_id in _ACCOUNT_USAGE_STALENESS_WARNED: + return + _ACCOUNT_USAGE_STALENESS_WARNED.add(session_id) + logger.warning( + "Listing objects from ACCOUNT_USAGE, which lags live state by up to ~2 hours: " + "recently created objects may be missing. Pass --no-use-account-usage to list with " + "real-time SHOW instead." + ) + + +def _list_schema_scoped_from_account_usage( + session: SnowflakeConnection, sql: str, use_account_usage: bool = False +) -> Optional[list[FQN]]: + """ + List schema-scoped objects from ACCOUNT_USAGE, which SHOW ... IN ACCOUNT cannot past + its 10,000-row cap (error 090153). + + `sql` must alias its columns to the names SHOW uses ("database_name", "schema_name", + "name") and exclude dropped objects. Returns None when ACCOUNT_USAGE is off, + unavailable, or the query fails, so the caller falls back to SHOW. + """ + if not _should_use_account_usage(session, use_account_usage): + return None + try: + rows = execute(session, sql, cacheable=True) + except Exception as e: # any failure should fall back to SHOW + logger.warning(f"ACCOUNT_USAGE listing failed, falling back to SHOW: {e}") + _mark_account_usage_fallback(session) + return None + + _warn_account_usage_staleness(session) + # Compare in ResourceName space: str() renders '"MyDb"' where ACCOUNT_USAGE returns + # MyDb, so comparing rendered strings drops every quoted or mixed-case database. + user_databases = _list_databases(session) + results = [] + for row in rows: + database = resource_name_from_snowflake_metadata(row["database_name"]) + if row["database_name"] in SYSTEM_DATABASES or database not in user_databases: + continue + if row["schema_name"] == "INFORMATION_SCHEMA": + continue + results.append( + FQN( + database=database, + schema=resource_name_from_snowflake_metadata(row["schema_name"]), + name=resource_name_from_snowflake_metadata(row["name"]), + ) + ) + return results + + +def list_stages(session: SnowflakeConnection, use_account_usage: bool = False) -> list[FQN]: + # Named stages only, matching the SHOW path's type filter below. + from_account_usage = _list_schema_scoped_from_account_usage( + session, + """ + SELECT stage_catalog AS "database_name", + stage_schema AS "schema_name", + stage_name AS "name" + FROM SNOWFLAKE.ACCOUNT_USAGE.STAGES + WHERE deleted IS NULL + AND stage_type IN ('Internal Named', 'External Named') + """, + use_account_usage=use_account_usage, + ) + if from_account_usage is not None: + return from_account_usage + show_result = execute(session, "SHOW STAGES IN ACCOUNT", cacheable=True) stages = [] for row in show_result: @@ -4884,7 +5033,27 @@ def list_streams(session: SnowflakeConnection) -> list[FQN]: return list_schema_scoped_resource(session, "STREAMS") -def list_tables(session: SnowflakeConnection) -> list[FQN]: +def list_tables(session: SnowflakeConnection, use_account_usage: bool = False) -> list[FQN]: + # The SHOW path's exclusions map onto table_type plus the is_* flags, which older + # accounts may lack -- hence COALESCE. + from_account_usage = _list_schema_scoped_from_account_usage( + session, + """ + SELECT table_catalog AS "database_name", + table_schema AS "schema_name", + table_name AS "name" + FROM SNOWFLAKE.ACCOUNT_USAGE.TABLES + WHERE deleted IS NULL + AND table_type = 'BASE TABLE' + AND COALESCE(is_iceberg, 'NO') = 'NO' + AND COALESCE(is_dynamic, 'NO') = 'NO' + AND COALESCE(is_hybrid, 'NO') = 'NO' + """, + use_account_usage=use_account_usage, + ) + if from_account_usage is not None: + return from_account_usage + show_result = execute(session, "SHOW TABLES IN ACCOUNT", cacheable=True) user_databases = _list_databases(session) tables = [] @@ -5140,7 +5309,23 @@ def error_handler(err: Exception, sql: str): return key_pairs -def list_views(session: SnowflakeConnection) -> list[FQN]: +def list_views(session: SnowflakeConnection, use_account_usage: bool = False) -> list[FQN]: + # .TABLES, not .VIEWS: table_type excludes materialized views, as the SHOW path does. + from_account_usage = _list_schema_scoped_from_account_usage( + session, + """ + SELECT table_catalog AS "database_name", + table_schema AS "schema_name", + table_name AS "name" + FROM SNOWFLAKE.ACCOUNT_USAGE.TABLES + WHERE deleted IS NULL + AND table_type = 'VIEW' + """, + use_account_usage=use_account_usage, + ) + if from_account_usage is not None: + return from_account_usage + show_result = execute(session, "SHOW VIEWS IN ACCOUNT", cacheable=True) views = [] for row in show_result: diff --git a/snowcap/operations/export.py b/snowcap/operations/export.py index e931e32..112025c 100644 --- a/snowcap/operations/export.py +++ b/snowcap/operations/export.py @@ -28,13 +28,14 @@ def export_resources( include: Optional[list[ResourceType]] = None, exclude: Optional[list[ResourceType]] = None, threads: int = DEFAULT_EXPORT_THREADS, + use_account_usage: bool = True, ) -> dict[str, list]: if session is None: session = connect() # Pre-populate ACCOUNT_USAGE caches before parallel fetching to avoid # multiple threads hitting the slow ACCOUNT_USAGE access check simultaneously - if threads > 1: + if threads > 1 and use_account_usage: try: populate_account_usage_caches(session) except Exception as e: @@ -53,7 +54,7 @@ def export_resources( ) continue try: - config.update(export_resource(session, resource_type, threads=threads)) + config.update(export_resource(session, resource_type, threads=threads, use_account_usage=use_account_usage)) # No list method for resource except AttributeError: logger.warning(f"Skipping {resource_type} because it has no list method") @@ -84,9 +85,16 @@ def _fetch_resource_safe(session, urn: URN): return None -def export_resource(session, resource_type: ResourceType, threads: int = DEFAULT_EXPORT_THREADS) -> dict[str, list]: +def export_resource( + session, + resource_type: ResourceType, + threads: int = DEFAULT_EXPORT_THREADS, + use_account_usage: bool = True, +) -> dict[str, list]: resource_label = resource_label_for_type(resource_type) - resource_names = list_resource(session, resource_label) + # list_resource forwards only the kwargs a given list_* function declares, so this is a + # no-op for listers that do not take it. + resource_names = list_resource(session, resource_label, use_account_usage=use_account_usage) if len(resource_names) == 0: return {} diff --git a/tests/test_blueprint.py b/tests/test_blueprint.py index bbed0c2..40aa49f 100644 --- a/tests/test_blueprint.py +++ b/tests/test_blueprint.py @@ -2796,6 +2796,59 @@ def test_the_query_is_not_narrowed_to_roles_named_in_config(self): assert "future_grant_database_roles" not in kwargs +class TestSyncListingHonoursUseAccountUsage: + """Sync mode drops whatever remote state omits, so a listing read from ACCOUNT_USAGE -- + which lags live state by up to ~2 hours -- proposes dropping objects that exist.""" + + @pytest.fixture(autouse=True) + def _ctx(self, session_ctx): + type(self).SESSION_CTX = session_ctx + from snowcap.data_provider import reset_account_usage_caches + + reset_account_usage_caches() + yield + reset_account_usage_caches() + + def _listing_sql(self, use_account_usage): + """SQL issued while sync mode lists tables. The later failure is mock plumbing for + reference resolution, not the behaviour under test.""" + from snowcap.blueprint_config import BlueprintConfig + + bp = Blueprint(resources=[res.Role(name="SOME_ROLE")]) + bp._config = BlueprintConfig( + sync_resources={ResourceType.TABLE}, + use_account_usage=use_account_usage, + ) + + with ( + patch("snowcap.blueprint.data_provider.fetch_session") as mock_session, + patch("snowcap.blueprint.data_provider.use_secondary_roles"), + patch("snowcap.blueprint.data_provider.populate_account_usage_caches"), + patch("snowcap.data_provider._has_account_usage_access", return_value=True), + patch("snowcap.data_provider._list_databases", return_value=[]), + patch("snowcap.data_provider.execute", return_value=[]) as mock_execute, + ): + mock_session.return_value = self.SESSION_CTX + manifest = bp.generate_manifest(self.SESSION_CTX) + try: + bp.fetch_remote_state(MagicMock(), manifest) + except Exception: + pass + return [call.args[1] for call in mock_execute.call_args_list if len(call.args) > 1] + + def test_tables_are_listed_with_show_when_account_usage_is_off(self): + sql = self._listing_sql(use_account_usage=False) + + assert any("SHOW TABLES IN ACCOUNT" in statement for statement in sql) + assert not any("ACCOUNT_USAGE.TABLES" in statement for statement in sql) + + def test_tables_are_listed_from_account_usage_when_it_is_on(self): + sql = self._listing_sql(use_account_usage=True) + + assert any("ACCOUNT_USAGE.TABLES" in statement for statement in sql) + assert not any("SHOW TABLES IN ACCOUNT" in statement for statement in sql) + + class TestSummarizePlanValue: """The plan table must not dump a multiline SQL body (alert THEN, task body).""" diff --git a/tests/test_data_provider.py b/tests/test_data_provider.py index de8a82b..34c4d8d 100644 --- a/tests/test_data_provider.py +++ b/tests/test_data_provider.py @@ -8,8 +8,16 @@ import pytest from unittest.mock import MagicMock, patch, PropertyMock +from snowcap.client import reset_cache from snowcap.data_provider import ( # Helper functions + _ACCOUNT_USAGE_GRANTS_CACHE, + _fetch_grant_to_role, + _grant_name_key, + _grants_by_role_index, + _list_schema_scoped_from_account_usage, + _show_all_grants_to_role, + reset_account_usage_caches, _quote_snowflake_identifier, _get_owner_identifier, _desc_result_to_dict, @@ -46,6 +54,9 @@ list_resource, list_account_scoped_resource, list_schema_scoped_resource, + list_stages, + list_tables, + list_views, # Session functions fetch_account_locator, fetch_region, @@ -58,6 +69,7 @@ import datetime import json +import logging import pytz @@ -2892,3 +2904,467 @@ def test_show_future_grants_container_inference_is_quote_aware(self, mock_execut assert grants[0]["granted_on"] == "DATABASE" assert grants[1]["granted_on"] == "SCHEMA" + + +def _reset_grant_indexes(): + """Both grant indexes are process-wide, so a test must not inherit another's.""" + from snowcap.data_provider import reset_account_usage_caches + + reset_account_usage_caches() + reset_cache() + + +class TestGrantNameKey: + """Tests for _grant_name_key, the dict key standing in for ResourceName equality.""" + + def test_resource_name_is_unsafe_as_a_dict_key(self): + # The reason the helper exists at all: __eq__ calls these the same name but + # __hash__ disagrees, so a plain {ResourceName: grant} dict misses the lookup. + assert ResourceName('"FOO"') == ResourceName("FOO") + assert hash(ResourceName('"FOO"')) != hash(ResourceName("FOO")) + assert _grant_name_key('"FOO"') == _grant_name_key("FOO") + + @pytest.mark.parametrize( + "left,right", + [ + ("FOO", "FOO"), + ("FOO", "foo"), + ('"FOO"', "FOO"), + ('"FOO"', "foo"), + ('"Foo"', '"Foo"'), + # A name that has to be quoted to be legal is the same name either way. + ("my-name", '"my-name"'), + ], + ) + def test_key_is_shared_by_names_resource_name_calls_equal(self, left, right): + assert ResourceName(left) == ResourceName(right) + assert _grant_name_key(left) == _grant_name_key(right) + + @pytest.mark.parametrize( + "left,right", + [ + ('"Foo"', "Foo"), + ('"foo"', "FOO"), + ('"Foo"', '"foo"'), + ("FOO", "BAR"), + ], + ) + def test_key_differs_for_names_resource_name_calls_different(self, left, right): + assert ResourceName(left) != ResourceName(right) + assert _grant_name_key(left) != _grant_name_key(right) + + +class TestGrantLookupIndex: + """Tests for _grant_lookup_index and _fetch_grant_to_role, which reads through it.""" + + @pytest.fixture(autouse=True) + def _clean_indexes(self): + _reset_grant_indexes() + yield + _reset_grant_indexes() + + def _fetch(self, session, granted_on, on_name, privilege="SELECT", role="SOME_ROLE"): + return _fetch_grant_to_role( + session, + GrantType.OBJECT, + ResourceName(role), + granted_on, + on_name, + privilege, + ) + + @patch("snowcap.data_provider._show_grants_to_role") + def test_finds_a_grant_by_type_privilege_and_name(self, mock_show_grants): + mock_show_grants.return_value = [ + _grant_to_role_row(privilege="SELECT", granted_on="TABLE", name="MY_DB.MY_SCHEMA.MY_TABLE"), + ] + + result = self._fetch(MagicMock(), "TABLE", "MY_DB.MY_SCHEMA.MY_TABLE") + + assert result is not None + assert result["privilege"] == "SELECT" + + @pytest.mark.parametrize( + "reported_name,requested_name", + [ + ("MY_TABLE", '"MY_TABLE"'), + ('"MY_TABLE"', "MY_TABLE"), + ("MY_TABLE", "my_table"), + ('"MixedCase"', '"MixedCase"'), + ], + ) + @patch("snowcap.data_provider._show_grants_to_role") + def test_quoting_does_not_change_the_match(self, mock_show_grants, reported_name, requested_name): + # The scan this replaced compared with ResourceName, so quoted and unquoted + # spellings of one name matched. The index has to keep that true. + mock_show_grants.return_value = [_grant_to_role_row(granted_on="TABLE", name=reported_name)] + + assert self._fetch(MagicMock(), "TABLE", requested_name, privilege="USAGE") is not None + + @patch("snowcap.data_provider._show_grants_to_role") + def test_a_genuinely_different_name_does_not_match(self, mock_show_grants): + mock_show_grants.return_value = [_grant_to_role_row(granted_on="TABLE", name='"MixedCase"')] + + assert self._fetch(MagicMock(), "TABLE", "MIXEDCASE", privilege="USAGE") is None + + @patch("snowcap.data_provider._show_grants_to_role") + def test_wrong_privilege_does_not_match(self, mock_show_grants): + mock_show_grants.return_value = [_grant_to_role_row(privilege="USAGE", granted_on="TABLE", name="MY_TABLE")] + + assert self._fetch(MagicMock(), "TABLE", "MY_TABLE", privilege="SELECT") is None + + @patch("snowcap.data_provider._show_grants_to_role") + def test_account_grants_match_on_the_account_keyword(self, mock_show_grants): + mock_show_grants.return_value = [ + _grant_to_role_row(privilege="AUDIT", granted_on="ACCOUNT", name="SOME_ACCOUNT_LOCATOR"), + ] + + assert self._fetch(MagicMock(), "ACCOUNT", "ACCOUNT", privilege="AUDIT") is not None + + @patch("snowcap.data_provider._show_grants_to_role") + def test_object_type_synonyms_still_match(self, mock_show_grants): + # Snowflake reports the type with an underscore where the DDL spells it with a + # space. Both sides go through _granted_on_label, so the index has to as well. + mock_show_grants.return_value = [ + _grant_to_role_row(privilege="READ", granted_on="GIT_REPOSITORY", name="MY_DB.PUBLIC.MY_REPO"), + ] + + result = self._fetch(MagicMock(), "GIT REPOSITORY", "MY_DB.PUBLIC.MY_REPO", privilege="READ") + + assert result is not None + + @patch("snowcap.data_provider._show_grants_to_role") + def test_duplicate_grants_resolve_to_the_first_one(self, mock_show_grants): + # Two roles can grant the same privilege on the same object. The linear scan + # returned the first row; the index has to agree. + mock_show_grants.return_value = [ + _grant_to_role_row(granted_on="TABLE", name="MY_TABLE", granted_by="ROLE_A"), + _grant_to_role_row(granted_on="TABLE", name="MY_TABLE", granted_by="ROLE_B"), + ] + + result = self._fetch(MagicMock(), "TABLE", "MY_TABLE", privilege="USAGE") + + assert result is not None + assert result["granted_by"] == "ROLE_A" + + @patch("snowcap.data_provider._show_grants_to_role") + def test_a_role_is_read_once_however_many_grants_are_looked_up(self, mock_show_grants): + # The point of the index: 90k grant lookups must not re-read a role 90k times. + mock_show_grants.return_value = [_grant_to_role_row(granted_on="TABLE", name=f"MY_TABLE_{i}") for i in range(5)] + session = MagicMock() + + for i in range(5): + assert self._fetch(session, "TABLE", f"MY_TABLE_{i}", privilege="USAGE") is not None + + assert mock_show_grants.call_count == 1 + + @patch("snowcap.data_provider._show_grants_to_role") + def test_each_role_gets_its_own_index(self, mock_show_grants): + mock_show_grants.side_effect = [ + [_grant_to_role_row(granted_on="TABLE", name="TABLE_A")], + [_grant_to_role_row(granted_on="TABLE", name="TABLE_B")], + ] + session = MagicMock() + + assert self._fetch(session, "TABLE", "TABLE_A", privilege="USAGE", role="ROLE_A") is not None + assert self._fetch(session, "TABLE", "TABLE_A", privilege="USAGE", role="ROLE_B") is None + assert self._fetch(session, "TABLE", "TABLE_B", privilege="USAGE", role="ROLE_B") is not None + + @patch("snowcap.data_provider._show_future_grants_to_role") + @patch("snowcap.data_provider._show_grants_to_role") + def test_future_grants_are_indexed_from_the_future_grant_source(self, mock_show_grants, mock_show_future): + mock_show_future.return_value = [_grant_to_role_row(privilege="SELECT", granted_on="SCHEMA", name="MY_DB.SCH")] + + result = _fetch_grant_to_role( + MagicMock(), GrantType.FUTURE, ResourceName("SOME_ROLE"), "SCHEMA", "MY_DB.SCH", "SELECT" + ) + + assert result is not None + mock_show_grants.assert_not_called() + + @patch("snowcap.data_provider._show_grants_to_role") + def test_resetting_the_caches_rebuilds_the_index(self, mock_show_grants): + # The index is derived state. Leaving it behind would hide grants that an apply + # created after it was built. + mock_show_grants.side_effect = [ + [], + [_grant_to_role_row(granted_on="TABLE", name="MY_TABLE")], + ] + session = MagicMock() + + assert self._fetch(session, "TABLE", "MY_TABLE", privilege="USAGE") is None + _reset_grant_indexes() + + assert self._fetch(session, "TABLE", "MY_TABLE", privilege="USAGE") is not None + + @patch("snowcap.data_provider._show_grants_to_role") + def test_resetting_the_sql_cache_rebuilds_the_index(self, mock_show_grants): + # Blueprint.plan() calls reset_cache() and nothing else, so a second plan on the + # same session was answered from the first plan's grants. + mock_show_grants.side_effect = [ + [_grant_to_role_row(granted_on="TABLE", name="MY_TABLE")], + [], + ] + session = MagicMock() + + assert self._fetch(session, "TABLE", "MY_TABLE", privilege="USAGE") is not None + reset_cache() + + assert self._fetch(session, "TABLE", "MY_TABLE", privilege="USAGE") is None + assert mock_show_grants.call_count == 2 + + @patch("snowcap.data_provider._show_grants_to_role") + def test_the_index_survives_within_one_plan(self, mock_show_grants): + # The counterpart: only a reset invalidates the index, or the lookup is pointless. + mock_show_grants.return_value = [_grant_to_role_row(granted_on="TABLE", name="MY_TABLE")] + session = MagicMock() + + self._fetch(session, "TABLE", "MY_TABLE", privilege="USAGE") + self._fetch(session, "TABLE", "MY_TABLE", privilege="USAGE") + + assert mock_show_grants.call_count == 1 + + +class TestGrantsByRoleIndex: + """Tests for _grants_by_role_index, the role index over the ACCOUNT_USAGE grant cache.""" + + @pytest.fixture(autouse=True) + def _clean_indexes(self): + _reset_grant_indexes() + yield + _reset_grant_indexes() + + def _populate(self, session, grants): + _ACCOUNT_USAGE_GRANTS_CACHE[id(session)] = grants + + def test_groups_grants_by_grantee(self): + session = MagicMock() + self._populate( + session, + [ + _grant_to_role_row(name="DB_A", grantee_name="ROLE_A"), + _grant_to_role_row(name="DB_B", grantee_name="ROLE_B"), + _grant_to_role_row(name="DB_C", grantee_name="ROLE_A"), + ], + ) + + index = _grants_by_role_index(id(session)) + + assert sorted(index) == ["ROLE_A", "ROLE_B"] + assert [grant["name"] for grant in index["ROLE_A"]] == ["DB_A", "DB_C"] + + def test_grantee_names_are_matched_case_insensitively(self): + # ACCOUNT_USAGE reports the grantee as stored; the caller looks it up by the + # upper-cased role name, which is what the scan this replaced compared. + session = MagicMock() + self._populate(session, [_grant_to_role_row(grantee_name="role_a")]) + + assert _grants_by_role_index(id(session))["ROLE_A"] + + def test_grants_to_anything_other_than_a_role_are_left_out(self): + session = MagicMock() + self._populate( + session, + [ + _grant_to_role_row(grantee_name="ROLE_A", granted_to="ROLE"), + _grant_to_role_row(grantee_name="ROLE_A", granted_to="DATABASE_ROLE"), + ], + ) + + assert len(_grants_by_role_index(id(session))["ROLE_A"]) == 1 + + @patch("snowcap.data_provider.execute") + def test_show_grants_reads_the_index_instead_of_the_database(self, mock_execute): + session = MagicMock() + self._populate( + session, + [ + _grant_to_role_row(name="DB_A", grantee_name="ROLE_A"), + _grant_to_role_row(name="DB_B", grantee_name="ROLE_B"), + ], + ) + + grants = _show_all_grants_to_role(session, ResourceName("ROLE_A")) + + assert [grant["name"] for grant in grants] == ["DB_A"] + mock_execute.assert_not_called() + + @patch("snowcap.data_provider.execute") + def test_a_role_with_no_grants_reads_as_empty(self, mock_execute): + session = MagicMock() + self._populate(session, [_grant_to_role_row(grantee_name="ROLE_A")]) + + assert _show_all_grants_to_role(session, ResourceName("ROLE_B")) == [] + mock_execute.assert_not_called() + + def test_the_index_is_built_once_per_session(self): + session = MagicMock() + self._populate(session, [_grant_to_role_row(grantee_name="ROLE_A")]) + + first = _grants_by_role_index(id(session)) + second = _grants_by_role_index(id(session)) + + assert first is second + + def test_resetting_the_caches_drops_the_index(self): + session = MagicMock() + self._populate(session, [_grant_to_role_row(grantee_name="ROLE_A")]) + _grants_by_role_index(id(session)) + + _reset_grant_indexes() + + assert _grants_by_role_index(id(session)) == {} + + +def _account_usage_row(database_name="MY_DB", schema_name="PUBLIC", name="MY_TABLE"): + """A row shaped the way the ACCOUNT_USAGE listing queries alias their columns.""" + return {"database_name": database_name, "schema_name": schema_name, "name": name} + + +class TestListSchemaScopedFromAccountUsage: + """Tests for _list_schema_scoped_from_account_usage, the SHOW-cap workaround.""" + + @pytest.fixture(autouse=True) + def _clean_caches(self): + reset_account_usage_caches() + yield + reset_account_usage_caches() + + @patch("snowcap.data_provider._has_account_usage_access", return_value=True) + @patch("snowcap.data_provider._list_databases") + @patch("snowcap.data_provider.execute") + def test_a_quoted_database_is_kept(self, mock_execute, mock_list_databases, _mock_access): + # ACCOUNT_USAGE returns the bare name (MyDb) while str(ResourceName) renders it with + # quotes ('"MyDb"'), so comparing rendered strings drops every database that is not + # plain upper case. The comparison has to happen in ResourceName space. + mock_list_databases.return_value = [ResourceName('"MyDb"')] + mock_execute.return_value = [_account_usage_row(database_name="MyDb")] + + result = _list_schema_scoped_from_account_usage(MagicMock(), "SELECT 1", use_account_usage=True) + + assert [str(fqn.database) for fqn in result] == ['"MyDb"'] + + @patch("snowcap.data_provider._has_account_usage_access", return_value=True) + @patch("snowcap.data_provider._list_databases") + @patch("snowcap.data_provider.execute") + def test_an_uppercase_database_is_kept(self, mock_execute, mock_list_databases, _mock_access): + mock_list_databases.return_value = [ResourceName("MY_DB")] + mock_execute.return_value = [_account_usage_row()] + + result = _list_schema_scoped_from_account_usage(MagicMock(), "SELECT 1", use_account_usage=True) + + assert [str(fqn) for fqn in result] == ["MY_DB.PUBLIC.MY_TABLE"] + + @patch("snowcap.data_provider._has_account_usage_access", return_value=True) + @patch("snowcap.data_provider._list_databases") + @patch("snowcap.data_provider.execute") + def test_a_database_snowcap_does_not_manage_is_dropped(self, mock_execute, mock_list_databases, _mock_access): + # Shares and imported databases are not in _list_databases, and exporting objects + # out of them would produce config that cannot be applied. + mock_list_databases.return_value = [ResourceName("MY_DB")] + mock_execute.return_value = [_account_usage_row(database_name="SOME_SHARE")] + + assert _list_schema_scoped_from_account_usage(MagicMock(), "SELECT 1", use_account_usage=True) == [] + + @patch("snowcap.data_provider._has_account_usage_access", return_value=True) + @patch("snowcap.data_provider._list_databases") + @patch("snowcap.data_provider.execute") + def test_system_databases_and_information_schema_are_dropped(self, mock_execute, mock_list_databases, _mock_access): + mock_list_databases.return_value = [ResourceName("MY_DB")] + mock_execute.return_value = [ + _account_usage_row(database_name="SNOWFLAKE"), + _account_usage_row(schema_name="INFORMATION_SCHEMA"), + _account_usage_row(), + ] + + result = _list_schema_scoped_from_account_usage(MagicMock(), "SELECT 1", use_account_usage=True) + + assert [str(fqn) for fqn in result] == ["MY_DB.PUBLIC.MY_TABLE"] + + @patch("snowcap.data_provider._has_account_usage_access", return_value=True) + @patch("snowcap.data_provider.execute") + def test_opting_out_returns_none_without_querying(self, mock_execute, _mock_access): + # --no-use-account-usage has to reach all the way down, not just skip the warning. + assert _list_schema_scoped_from_account_usage(MagicMock(), "SELECT 1", use_account_usage=False) is None + mock_execute.assert_not_called() + + @patch("snowcap.data_provider._has_account_usage_access", return_value=False) + @patch("snowcap.data_provider.execute") + def test_no_account_usage_access_returns_none_so_the_caller_can_show(self, mock_execute, _mock_access): + assert _list_schema_scoped_from_account_usage(MagicMock(), "SELECT 1", use_account_usage=True) is None + mock_execute.assert_not_called() + + @patch("snowcap.data_provider._has_account_usage_access", return_value=True) + @patch("snowcap.data_provider.execute", side_effect=Exception("boom")) + def test_a_failed_query_returns_none_and_stops_retrying(self, mock_execute, _mock_access): + session = MagicMock() + + assert _list_schema_scoped_from_account_usage(session, "SELECT 1", use_account_usage=True) is None + # The session is marked as fallen back, so the next listing does not pay for the + # same failure again. + assert _list_schema_scoped_from_account_usage(session, "SELECT 2", use_account_usage=True) is None + assert mock_execute.call_count == 1 + + @patch("snowcap.data_provider._has_account_usage_access", return_value=True) + @patch("snowcap.data_provider._list_databases") + @patch("snowcap.data_provider.execute") + def test_staleness_is_warned_about_once_per_session(self, mock_execute, mock_list_databases, _mock_access, caplog): + # ACCOUNT_USAGE lags live state by up to ~2 hours. That is worth saying once a run, + # not once per resource type. + mock_list_databases.return_value = [ResourceName("MY_DB")] + mock_execute.return_value = [_account_usage_row()] + session = MagicMock() + + with caplog.at_level(logging.WARNING, logger="snowcap"): + _list_schema_scoped_from_account_usage(session, "SELECT 1", use_account_usage=True) + _list_schema_scoped_from_account_usage(session, "SELECT 2", use_account_usage=True) + + warnings = [record for record in caplog.records if "--no-use-account-usage" in record.message] + assert len(warnings) == 1 + + +class TestSchemaScopedListersUseAccountUsage: + """list_tables / list_views / list_stages read ACCOUNT_USAGE, then fall back to SHOW.""" + + @pytest.fixture(autouse=True) + def _clean_caches(self): + reset_account_usage_caches() + yield + reset_account_usage_caches() + + @pytest.mark.parametrize("lister", [list_tables, list_views, list_stages]) + @patch("snowcap.data_provider._list_schema_scoped_from_account_usage") + @patch("snowcap.data_provider.execute") + def test_account_usage_answers_without_a_show(self, mock_execute, mock_from_account_usage, lister): + # SHOW ... IN ACCOUNT is capped at 10,000 rows, so on a large account it is not an + # option at all -- it fails outright and takes the export with it. + listed = [FQN(database=ResourceName("MY_DB"), schema=ResourceName("PUBLIC"), name=ResourceName("MY_OBJECT"))] + mock_from_account_usage.return_value = listed + + assert lister(MagicMock(), use_account_usage=True) == listed + mock_execute.assert_not_called() + + @pytest.mark.parametrize("lister", [list_tables, list_views, list_stages]) + @patch("snowcap.data_provider._has_account_usage_access", return_value=True) + @patch("snowcap.data_provider._list_databases", return_value=[]) + @patch("snowcap.data_provider.execute", return_value=[]) + def test_the_default_is_show(self, mock_execute, _mock_databases, _mock_access, lister): + # A caller that says nothing keeps real-time SHOW; the lagging view is opt-in. + lister(MagicMock()) + + assert "SHOW" in mock_execute.call_args_list[0].args[1] + + @pytest.mark.parametrize("lister", [list_tables, list_views, list_stages]) + @patch("snowcap.data_provider._list_schema_scoped_from_account_usage", return_value=None) + @patch("snowcap.data_provider.execute", return_value=[]) + def test_show_still_runs_when_account_usage_is_unavailable(self, mock_execute, _mock_from_account_usage, lister): + assert lister(MagicMock()) == [] + assert "SHOW" in mock_execute.call_args[0][1] + + @pytest.mark.parametrize("lister", [list_tables, list_views, list_stages]) + @patch("snowcap.data_provider._list_schema_scoped_from_account_usage", return_value=None) + @patch("snowcap.data_provider.execute", return_value=[]) + def test_the_opt_out_reaches_the_listing_helper(self, _mock_execute, mock_from_account_usage, lister): + lister(MagicMock(), use_account_usage=False) + + assert mock_from_account_usage.call_args.kwargs["use_account_usage"] is False