Handle snowflake lack of CTAS tracking

Author: treysp

sqlmesh/core/engine_adapter/base.py

@@ -2413,6 +2413,11 @@ def _log_sql(
 
         logger.log(self._execute_log_level, "Executing SQL: %s", sql_to_log)
 
+    def _record_execution_stats(
+        self, sql: str, rowcount: t.Optional[int] = None, bytes_processed: t.Optional[int] = None
+    ) -> None:
+        QueryExecutionTracker.record_execution(sql, rowcount, bytes_processed)
+
     def _execute(self, sql: str, track_execution_stats: bool = False, **kwargs: t.Any) -> None:
         self.cursor.execute(sql, **kwargs)
 
@@ -2429,7 +2434,7 @@ def _execute(self, sql: str, track_execution_stats: bool = False, **kwargs: t.An
                 except (TypeError, ValueError):
                     pass
 
-            QueryExecutionTracker.record_execution(sql, rowcount, None)
+            self._record_execution_stats(sql, rowcount)
 
     @contextlib.contextmanager
     def temp_table(

sqlmesh/core/engine_adapter/snowflake.py

@@ -2,6 +2,7 @@
 
 import contextlib
 import logging
+import re
 import typing as t
 
 from sqlglot import exp
@@ -24,6 +25,7 @@
     set_catalog,
 )
 from sqlmesh.core.schema_diff import SchemaDiffer
+from sqlmesh.core.snapshot.execution_tracker import QueryExecutionTracker
 from sqlmesh.utils import optional_import, get_source_columns_to_types
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pandas import columns_to_types_from_dtypes
@@ -72,6 +74,7 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
     )
     MANAGED_TABLE_KIND = "DYNAMIC TABLE"
     SNOWPARK = "snowpark"
+    SUPPORTS_QUERY_EXECUTION_TRACKING = True
 
     @contextlib.contextmanager
     def session(self, properties: SessionProperties) -> t.Iterator[None]:
@@ -664,3 +667,33 @@ def close(self) -> t.Any:
             self._connection_pool.set_attribute(self.SNOWPARK, None)
 
         return super().close()
+
+    def _record_execution_stats(
+        self, sql: str, rowcount: t.Optional[int] = None, bytes_processed: t.Optional[int] = None
+    ) -> None:
+        """Snowflake does not report row counts for CTAS like other DML operations.
+
+        They neither report the sentinel value -1 nor do they report 0 rows. Instead, they return a single data row
+        containing the string "Table <table_name> successfully created." and a row count of 1.
+
+        We do not want to record the row count of 1 for CTAS operations, so we check for that data pattern and return
+        early if it is detected.
+
+        Regex explanation - Snowflake identifiers may be:
+        - An unquoted contiguous set of [a-zA-Z0-9_$] characters
+        - A double-quoted string that may contain spaces and nested double-quotes represented by `""`
+          - Example: " my ""table"" name "
+          - Pattern: "(?:[^"]|"")+"
+            - ?: is a non-capturing group
+            - [^"] matches any single character except a double-quote
+            - "" matches two sequential double-quotes
+        """
+        if rowcount == 1:
+            results = self.cursor.fetchall()
+            if results and len(results) == 1:
+                is_ctas = re.match(
+                    r'Table ([a-zA-Z0-9_$]+|"(?:[^"]|"")+") successfully created\.', results[0][0]
+                )
+                if is_ctas:
+                    return
+        QueryExecutionTracker.record_execution(sql, rowcount, bytes_processed)

sqlmesh/core/snapshot/execution_tracker.py

@@ -10,8 +10,8 @@
 @dataclass
 class QueryExecutionStats:
     snapshot_batch_id: str
-    total_rows_processed: int = 0
-    total_bytes_processed: int = 0
+    total_rows_processed: t.Optional[int] = None
+    total_bytes_processed: t.Optional[int] = None
     query_count: int = 0
     queries_executed: t.List[t.Tuple[str, t.Optional[int], t.Optional[int], float]] = field(
         default_factory=list
@@ -42,12 +42,18 @@ def add_execution(
         self, sql: str, row_count: t.Optional[int], bytes_processed: t.Optional[int]
     ) -> None:
         if row_count is not None and row_count >= 0:
-            self.stats.total_rows_processed += row_count
+            if self.stats.total_rows_processed is None:
+                self.stats.total_rows_processed = row_count
+            else:
+                self.stats.total_rows_processed += row_count
 
             # conditional on row_count because we should only count bytes corresponding to
             # DML actions whose rows were captured
             if bytes_processed is not None and bytes_processed >= 0:
-                self.stats.total_bytes_processed += bytes_processed
+                if self.stats.total_bytes_processed is None:
+                    self.stats.total_bytes_processed = bytes_processed
+                else:
+                    self.stats.total_bytes_processed += bytes_processed
 
         self.stats.query_count += 1
         # TODO: remove this

tests/core/engine_adapter/integration/test_integration.py

@@ -2413,7 +2413,12 @@ def capture_execution_stats(
 
     if ctx.engine_adapter.SUPPORTS_QUERY_EXECUTION_TRACKING:
         assert actual_execution_stats["incremental_model"].total_rows_processed == 7
-        assert actual_execution_stats["full_model"].total_rows_processed == 3
+        # snowflake doesn't track rows for CTAS
+        assert actual_execution_stats["full_model"].total_rows_processed == (
+            None if ctx.mark.startswith("snowflake") else 3
+        )
+        # seed rows aren't tracked
+        assert actual_execution_stats["seed_model"].total_rows_processed is None
 
         if ctx.mark.startswith("bigquery"):
             assert actual_execution_stats["incremental_model"].total_bytes_processed