Skip to content

Commit 93ae8be

Browse files
committed
Merge remote-tracking branch 'origin/main' into fix/snowflake-alter-session-quoting
# Conflicts: # tests/core/engine_adapter/test_snowflake.py
2 parents 8019d3e + 3bd5972 commit 93ae8be

6 files changed

Lines changed: 329 additions & 22 deletions

File tree

‎Makefile‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,10 @@ bigquery-test: guard-BIGQUERY_KEYFILE engine-bigquery-install
238238
pytest -n auto -m "bigquery" --reruns 3 --junitxml=test-results/junit-bigquery.xml
239239

240240
databricks-test: guard-DATABRICKS_CATALOG guard-DATABRICKS_SERVER_HOSTNAME guard-DATABRICKS_HTTP_PATH guard-DATABRICKS_CONNECT_VERSION engine-databricks-install
241-
$(PIP) install 'databricks-connect==${DATABRICKS_CONNECT_VERSION}'
241+
# Connect 15.3.x requires numpy<2 and pyarrow>=4 (no upper bound). Cap PyArrow
242+
# in the same resolve that downgrades NumPy, or PyArrow 25+ is left installed
243+
# and cannot be imported.
244+
$(PIP) install 'databricks-connect==${DATABRICKS_CONNECT_VERSION}' 'pyarrow<25'
242245
pytest -n auto -m "databricks" --reruns 3 --junitxml=test-results/junit-databricks.xml
243246

244247
redshift-test: guard-REDSHIFT_HOST guard-REDSHIFT_USER guard-REDSHIFT_PASSWORD guard-REDSHIFT_DATABASE engine-redshift-install

‎sqlmesh/core/engine_adapter/snowflake.py‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -757,14 +757,32 @@ def close(self) -> t.Any:
757757
def get_table_last_modified_ts(self, table_names: t.List[TableName]) -> t.List[int]:
758758
from sqlmesh.utils.date import to_timestamp
759759

760-
num_tables = len(table_names)
761-
762-
query = "SELECT LAST_ALTERED FROM INFORMATION_SCHEMA.TABLES WHERE"
763-
for i, table_name in enumerate(table_names):
760+
# INFORMATION_SCHEMA only contains the objects of the database it belongs to, so the tables
761+
# are grouped by database and each group is looked up in its own database's INFORMATION_SCHEMA.
762+
# The groups are keyed on the catalog identifier rather than its name, because quoting matters
763+
# in Snowflake: unquoted raw is the database RAW, while quoted "raw" is a different database
764+
catalogs_to_tables: t.Dict[t.Optional[exp.Identifier], t.List[exp.Table]] = {}
765+
for table_name in table_names:
764766
table = exp.to_table(table_name)
765-
query += f"""(TABLE_NAME = '{table.name}' AND TABLE_SCHEMA = '{table.db}' AND TABLE_CATALOG = '{table.catalog}')"""
766-
if i < num_tables - 1:
767-
query += " OR "
767+
catalogs_to_tables.setdefault(table.args.get("catalog"), []).append(table)
768+
769+
results = []
770+
for catalog, tables in catalogs_to_tables.items():
771+
query = (
772+
exp.select("LAST_ALTERED")
773+
.from_(exp.table_("TABLES", db="INFORMATION_SCHEMA", catalog=catalog))
774+
.where(
775+
exp.or_(
776+
*(
777+
exp.and_(
778+
exp.column("TABLE_NAME").eq(exp.Literal.string(table.name)),
779+
exp.column("TABLE_SCHEMA").eq(exp.Literal.string(table.db)),
780+
)
781+
for table in tables
782+
)
783+
)
784+
)
785+
)
786+
results.extend(self.fetchall(query))
768787

769-
result = self.fetchall(query)
770-
return [to_timestamp(row[0]) for row in result]
788+
return [to_timestamp(row[0]) for row in results]

‎sqlmesh/utils/concurrency.py‎

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -65,26 +65,40 @@ def run(self) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]:
6565
return self._node_errors, self._skipped_nodes
6666

6767
def _process_node(self, node: H, executor: Executor) -> None:
68+
with self._unprocessed_nodes_lock:
69+
if self._finished_future.done():
70+
return
71+
6872
try:
6973
self.fn(node)
7074

7175
with self._unprocessed_nodes_lock:
76+
if self._finished_future.done():
77+
return
7278
self._unprocessed_nodes_num -= 1
7379
self._submit_next_nodes(executor, node)
7480
except Exception as ex:
7581
error = NodeExecutionFailedError(node)
7682
error.__cause__ = ex
7783

78-
if self.raise_on_error:
79-
self._finished_future.set_exception(error)
80-
return
81-
8284
with self._unprocessed_nodes_lock:
85+
if self._finished_future.done():
86+
return
87+
if self.raise_on_error:
88+
self._finished_future.set_exception(error)
89+
return
8390
self._unprocessed_nodes_num -= 1
8491
self._node_errors.append(error)
8592
self._skip_next_nodes(node)
93+
except BaseException as ex:
94+
with self._unprocessed_nodes_lock:
95+
if not self._finished_future.done():
96+
self._finished_future.set_exception(ex)
8697

8798
def _submit_next_nodes(self, executor: Executor, processed_node: t.Optional[H] = None) -> None:
99+
if self._finished_future.done():
100+
return
101+
88102
if not self._unprocessed_nodes_num:
89103
self._finished_future.set_result(None)
90104
return
@@ -101,6 +115,9 @@ def _submit_next_nodes(self, executor: Executor, processed_node: t.Optional[H] =
101115
executor.submit(self._process_node, submitted_node, executor)
102116

103117
def _skip_next_nodes(self, parent: H) -> None:
118+
if self._finished_future.done():
119+
return
120+
104121
if not self._unprocessed_nodes_num:
105122
self._finished_future.set_result(None)
106123
return
@@ -267,9 +284,11 @@ def concurrent_apply_to_values(
267284

268285
def _process_value(value: A, index: int) -> None:
269286
try:
270-
futures[index].set_result(fn(value))
271-
except Exception as ex:
287+
result = fn(value)
288+
except BaseException as ex:
272289
futures[index].set_exception(ex)
290+
else:
291+
futures[index].set_result(result)
273292

274293
with ThreadPoolExecutor(max_workers=tasks_num) as pool:
275294
for index, value in enumerate(values):

‎tests/core/engine_adapter/test_snowflake.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import typing as t
2+
from datetime import datetime, timezone
23

34
import pandas as pd # noqa: TID253
45
import pytest
@@ -14,6 +15,7 @@
1415
from sqlmesh.core.model import load_sql_based_model
1516
from sqlmesh.core.model.definition import SqlModel
1617
from sqlmesh.core.node import IntervalUnit
18+
from sqlmesh.utils.date import to_timestamp
1719
from sqlmesh.utils.errors import SQLMeshError
1820
from sqlmesh.utils import optional_import
1921
from tests.core.engine_adapter import to_sql_calls
@@ -73,6 +75,38 @@ def test_get_data_objects_lowercases_columns(
7375
assert data_object.clustering_key == "ID"
7476

7577

78+
def test_get_table_last_modified_ts_multiple_catalogs(
79+
snowflake_mocked_engine_adapter: SnowflakeEngineAdapter,
80+
) -> None:
81+
adapter = snowflake_mocked_engine_adapter
82+
adapter.cursor.fetchall.side_effect = [
83+
[
84+
(datetime(2025, 1, 1, tzinfo=timezone.utc),),
85+
(datetime(2025, 1, 2, tzinfo=timezone.utc),),
86+
],
87+
[(datetime(2025, 1, 3, tzinfo=timezone.utc),)],
88+
]
89+
90+
result = adapter.get_table_last_modified_ts(
91+
[
92+
'"ANALYTICS"."STAGING"."ORDERS"',
93+
'"RAW"."SHOPIFY"."ORDERS"',
94+
'"ANALYTICS"."STAGING"."CUSTOMERS"',
95+
]
96+
)
97+
98+
# INFORMATION_SCHEMA only covers the database it belongs to, so each database is queried separately
99+
assert to_sql_calls(adapter) == [
100+
"""SELECT LAST_ALTERED FROM "ANALYTICS".INFORMATION_SCHEMA.TABLES WHERE (TABLE_NAME = 'ORDERS' AND TABLE_SCHEMA = 'STAGING') OR (TABLE_NAME = 'CUSTOMERS' AND TABLE_SCHEMA = 'STAGING')""",
101+
"""SELECT LAST_ALTERED FROM "RAW".INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'ORDERS' AND TABLE_SCHEMA = 'SHOPIFY'""",
102+
]
103+
assert result == [
104+
to_timestamp("2025-01-01"),
105+
to_timestamp("2025-01-02"),
106+
to_timestamp("2025-01-03"),
107+
]
108+
109+
76110
@pytest.mark.parametrize(
77111
"current_warehouse, current_warehouse_exp, configured_warehouse, configured_warehouse_exp, should_change",
78112
[
@@ -183,6 +217,29 @@ def test_alter_session_parameters_are_not_quoted(make_mocked_engine_adapter: t.C
183217
]
184218

185219

220+
def test_get_table_last_modified_ts_quoted_and_unquoted_catalogs(
221+
snowflake_mocked_engine_adapter: SnowflakeEngineAdapter,
222+
) -> None:
223+
# unquoted raw resolves to the database RAW, quoted "raw" is a different database,
224+
# so the two must be looked up in two different information schemas
225+
adapter = snowflake_mocked_engine_adapter
226+
adapter.cursor.fetchall.side_effect = [
227+
[(datetime(2025, 1, 1, tzinfo=timezone.utc),)],
228+
[(datetime(2025, 1, 2, tzinfo=timezone.utc),)],
229+
]
230+
231+
result = adapter.get_table_last_modified_ts(['"raw".shopify.orders', "raw.shopify.customers"])
232+
233+
assert result == [
234+
to_timestamp(datetime(2025, 1, 1, tzinfo=timezone.utc)),
235+
to_timestamp(datetime(2025, 1, 2, tzinfo=timezone.utc)),
236+
]
237+
assert to_sql_calls(adapter) == [
238+
"""SELECT LAST_ALTERED FROM "raw".INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'orders' AND TABLE_SCHEMA = 'shopify'""",
239+
"""SELECT LAST_ALTERED FROM raw.INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'customers' AND TABLE_SCHEMA = 'shopify'""",
240+
]
241+
242+
186243
def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
187244
adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)
188245

‎tests/core/test_context.py‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2609,19 +2609,14 @@ def test_requirements(copy_to_temp_path: t.Callable):
26092609
"dev", no_prompts=True, skip_tests=True, skip_backfill=True, auto_apply=True
26102610
).environment
26112611
requirements = {"ipywidgets", "numpy", "pandas", "test_package"}
2612-
if IS_WINDOWS:
2613-
requirements.add("pendulum")
26142612
assert environment.requirements["pandas"] == "2.2.2"
26152613
assert set(environment.requirements) == requirements
26162614

26172615
context._requirements = {"numpy": "2.1.2", "pandas": "2.2.1"}
26182616
context._excluded_requirements = {"ipywidgets", "ruamel.yaml", "ruamel.yaml.clib"}
26192617
diff = context.plan_builder("dev", skip_tests=True, skip_backfill=True).build().context_diff
26202618
assert set(diff.previous_requirements) == requirements
2621-
reqs = {"numpy", "pandas"}
2622-
if IS_WINDOWS:
2623-
reqs.add("pendulum")
2624-
assert set(diff.requirements) == reqs
2619+
assert set(diff.requirements) == {"numpy", "pandas"}
26252620

26262621

26272622
def test_deactivate_automatic_requirement_inference(copy_to_temp_path: t.Callable):

0 commit comments

Comments
 (0)