Add snowflake test

Author: treysp

tests/core/engine_adapter/integration/test_integration.py

@@ -6,12 +6,13 @@
 import sys
 import typing as t
 import shutil
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, date
 from unittest.mock import patch
 import numpy as np  # noqa: TID253
 import pandas as pd  # noqa: TID253
 import pytest
 import pytz
+import time_machine
 from sqlglot import exp, parse_one
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.qualify_columns import quote_identifiers

tests/core/engine_adapter/integration/test_integration_snowflake.py

@@ -13,6 +13,7 @@
 from tests.core.engine_adapter.integration import TestContext
 from sqlmesh import model, ExecutionContext
 from sqlmesh.core.model import ModelKindName
+from sqlmesh.core.snapshot.execution_tracker import QueryExecutionTracker
 from datetime import datetime
 
 from tests.core.engine_adapter.integration import (
@@ -307,3 +308,21 @@ def fetch_database_names() -> t.Set[str]:
 
     engine_adapter.drop_catalog(sqlmesh_managed_catalog)  # works, catalog is SQLMesh-managed
     assert fetch_database_names() == {non_sqlmesh_managed_catalog}
+
+
+def test_rows_tracker(ctx: TestContext, engine_adapter: SnowflakeEngineAdapter):
+    sqlmesh = ctx.create_context()
+    tracker = QueryExecutionTracker()
+
+    with tracker.track_execution("a"):
+        # Snowflake doesn't report row counts for CTAS, so this should not be tracked
+        engine_adapter.execute(
+            "CREATE TABLE a (id int) AS SELECT 1 as id", track_rows_processed=True
+        )
+        engine_adapter.execute("INSERT INTO a VALUES (2), (3)", track_rows_processed=True)
+        engine_adapter.execute("INSERT INTO a VALUES (4)", track_rows_processed=True)
+
+    stats = tracker.get_execution_stats("a")
+    assert stats is not None
+    assert stats.query_count == 2
+    assert stats.total_rows_processed == 3