diff --git a/sqlmesh/core/engine_adapter/snowflake.py b/sqlmesh/core/engine_adapter/snowflake.py index 37ca6eee22..758e4847c1 100644 --- a/sqlmesh/core/engine_adapter/snowflake.py +++ b/sqlmesh/core/engine_adapter/snowflake.py @@ -757,14 +757,32 @@ def close(self) -> t.Any: def get_table_last_modified_ts(self, table_names: t.List[TableName]) -> t.List[int]: from sqlmesh.utils.date import to_timestamp - num_tables = len(table_names) - - query = "SELECT LAST_ALTERED FROM INFORMATION_SCHEMA.TABLES WHERE" - for i, table_name in enumerate(table_names): + # INFORMATION_SCHEMA only contains the objects of the database it belongs to, so the tables + # are grouped by database and each group is looked up in its own database's INFORMATION_SCHEMA. + # The groups are keyed on the catalog identifier rather than its name, because quoting matters + # in Snowflake: unquoted raw is the database RAW, while quoted "raw" is a different database + catalogs_to_tables: t.Dict[t.Optional[exp.Identifier], t.List[exp.Table]] = {} + for table_name in table_names: table = exp.to_table(table_name) - query += f"""(TABLE_NAME = '{table.name}' AND TABLE_SCHEMA = '{table.db}' AND TABLE_CATALOG = '{table.catalog}')""" - if i < num_tables - 1: - query += " OR " + catalogs_to_tables.setdefault(table.args.get("catalog"), []).append(table) + + results = [] + for catalog, tables in catalogs_to_tables.items(): + query = ( + exp.select("LAST_ALTERED") + .from_(exp.table_("TABLES", db="INFORMATION_SCHEMA", catalog=catalog)) + .where( + exp.or_( + *( + exp.and_( + exp.column("TABLE_NAME").eq(exp.Literal.string(table.name)), + exp.column("TABLE_SCHEMA").eq(exp.Literal.string(table.db)), + ) + for table in tables + ) + ) + ) + ) + results.extend(self.fetchall(query)) - result = self.fetchall(query) - return [to_timestamp(row[0]) for row in result] + return [to_timestamp(row[0]) for row in results] diff --git a/tests/core/engine_adapter/test_snowflake.py b/tests/core/engine_adapter/test_snowflake.py index 05f0133ec8..89f19fc588 100644 --- a/tests/core/engine_adapter/test_snowflake.py +++ b/tests/core/engine_adapter/test_snowflake.py @@ -1,4 +1,5 @@ import typing as t +from datetime import datetime, timezone import pandas as pd # noqa: TID253 import pytest @@ -14,6 +15,7 @@ from sqlmesh.core.model import load_sql_based_model from sqlmesh.core.model.definition import SqlModel from sqlmesh.core.node import IntervalUnit +from sqlmesh.utils.date import to_timestamp from sqlmesh.utils.errors import SQLMeshError from sqlmesh.utils import optional_import from tests.core.engine_adapter import to_sql_calls @@ -73,6 +75,38 @@ def test_get_data_objects_lowercases_columns( assert data_object.clustering_key == "ID" +def test_get_table_last_modified_ts_multiple_catalogs( + snowflake_mocked_engine_adapter: SnowflakeEngineAdapter, +) -> None: + adapter = snowflake_mocked_engine_adapter + adapter.cursor.fetchall.side_effect = [ + [ + (datetime(2025, 1, 1, tzinfo=timezone.utc),), + (datetime(2025, 1, 2, tzinfo=timezone.utc),), + ], + [(datetime(2025, 1, 3, tzinfo=timezone.utc),)], + ] + + result = adapter.get_table_last_modified_ts( + [ + '"ANALYTICS"."STAGING"."ORDERS"', + '"RAW"."SHOPIFY"."ORDERS"', + '"ANALYTICS"."STAGING"."CUSTOMERS"', + ] + ) + + # INFORMATION_SCHEMA only covers the database it belongs to, so each database is queried separately + assert to_sql_calls(adapter) == [ + """SELECT LAST_ALTERED FROM "ANALYTICS".INFORMATION_SCHEMA.TABLES WHERE (TABLE_NAME = 'ORDERS' AND TABLE_SCHEMA = 'STAGING') OR (TABLE_NAME = 'CUSTOMERS' AND TABLE_SCHEMA = 'STAGING')""", + """SELECT LAST_ALTERED FROM "RAW".INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'ORDERS' AND TABLE_SCHEMA = 'SHOPIFY'""", + ] + assert result == [ + to_timestamp("2025-01-01"), + to_timestamp("2025-01-02"), + to_timestamp("2025-01-03"), + ] + + @pytest.mark.parametrize( "current_warehouse, current_warehouse_exp, configured_warehouse, configured_warehouse_exp, should_change", [ @@ -171,6 +205,29 @@ def test_session( assert to_sql_calls(adapter) == expected_exception_calls +def test_get_table_last_modified_ts_quoted_and_unquoted_catalogs( + snowflake_mocked_engine_adapter: SnowflakeEngineAdapter, +) -> None: + # unquoted raw resolves to the database RAW, quoted "raw" is a different database, + # so the two must be looked up in two different information schemas + adapter = snowflake_mocked_engine_adapter + adapter.cursor.fetchall.side_effect = [ + [(datetime(2025, 1, 1, tzinfo=timezone.utc),)], + [(datetime(2025, 1, 2, tzinfo=timezone.utc),)], + ] + + result = adapter.get_table_last_modified_ts(['"raw".shopify.orders', "raw.shopify.customers"]) + + assert result == [ + to_timestamp(datetime(2025, 1, 1, tzinfo=timezone.utc)), + to_timestamp(datetime(2025, 1, 2, tzinfo=timezone.utc)), + ] + assert to_sql_calls(adapter) == [ + """SELECT LAST_ALTERED FROM "raw".INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'orders' AND TABLE_SCHEMA = 'shopify'""", + """SELECT LAST_ALTERED FROM raw.INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'customers' AND TABLE_SCHEMA = 'shopify'""", + ] + + def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture): adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)