feat: track rows processed during model evaluation (#5162)

Author: treysp

.circleci/continue_config.yml

@@ -239,7 +239,7 @@ jobs:
       - checkout
       - run:
           name: Install OS-level dependencies
-          command: ./.circleci/install-prerequisites.sh "<< parameters.engine >>"      
+          command: ./.circleci/install-prerequisites.sh "<< parameters.engine >>"
       - run:
           name: Generate database name
           command: |
@@ -307,7 +307,7 @@ workflows:
                 - redshift
                 - bigquery
                 - clickhouse-cloud
-                - athena                
+                - athena
                 - fabric
                 - gcp-postgres
           filters:

docs/integrations/engines/snowflake.md

@@ -250,6 +250,14 @@ And confirm that our schemas and objects exist in the Snowflake catalog:
 
 Congratulations - your SQLMesh project is up and running on Snowflake!
 
+### Where are the row counts?
+
+SQLMesh reports the number of rows processed by each model in its `plan` and `run` terminal output.
+
+However, due to limitations in the Snowflake Python connector, row counts cannot be determined for `CREATE TABLE AS` statements. Therefore, SQLMesh does not report row counts for certain model kinds, such as `FULL` models.
+
+Learn more about the connector limitation [on Github](https://github.com/snowflakedb/snowflake-connector-python/issues/645).
+
 ## Local/Built-in Scheduler
 **Engine Adapter Type**: `snowflake`
 

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "croniter",
     "duckdb>=0.10.0,!=0.10.3",
     "dateparser<=1.2.1",
+    "humanize",
     "hyperscript>=0.1.0",
     "importlib-metadata; python_version<'3.12'",
     "ipywidgets",

sqlmesh/core/console.py

@@ -7,6 +7,7 @@
 import uuid
 import logging
 import textwrap
+from humanize import metric, naturalsize
 from itertools import zip_longest
 from pathlib import Path
 from hyperscript import h
@@ -39,6 +40,7 @@
     SnapshotInfoLike,
 )
 from sqlmesh.core.snapshot.definition import Interval, Intervals, SnapshotTableInfo
+from sqlmesh.core.snapshot.execution_tracker import QueryExecutionStats
 from sqlmesh.core.test import ModelTest
 from sqlmesh.utils import rich as srich
 from sqlmesh.utils import Verbosity
@@ -439,6 +441,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        execution_stats: t.Optional[QueryExecutionStats] = None,
         auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         """Updates the snapshot evaluation progress."""
@@ -587,6 +590,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        execution_stats: t.Optional[QueryExecutionStats] = None,
         auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         pass
@@ -1032,7 +1036,9 @@ def start_evaluation_progress(
 
             # determine column widths
             self.evaluation_column_widths["annotation"] = (
-                _calculate_annotation_str_len(batched_intervals, self.AUDIT_PADDING)
+                _calculate_annotation_str_len(
+                    batched_intervals, self.AUDIT_PADDING, len(" (123.4m rows, 123.4 KiB)")
+                )
                 + 3  # brackets and opening escape backslash
             )
             self.evaluation_column_widths["name"] = max(
@@ -1077,6 +1083,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        execution_stats: t.Optional[QueryExecutionStats] = None,
         auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         """Update the snapshot evaluation progress."""
@@ -1097,7 +1104,7 @@ def update_snapshot_evaluation_progress(
                 ).ljust(self.evaluation_column_widths["name"])
 
                 annotation = _create_evaluation_model_annotation(
-                    snapshot, _format_evaluation_model_interval(snapshot, interval)
+                    snapshot, _format_evaluation_model_interval(snapshot, interval), execution_stats
                 )
                 audits_str = ""
                 if num_audits_passed:
@@ -3668,6 +3675,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        execution_stats: t.Optional[QueryExecutionStats] = None,
         auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
@@ -3838,6 +3846,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        execution_stats: t.Optional[QueryExecutionStats] = None,
         auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         message = f"Evaluated {snapshot.name} | batch={batch_idx} | duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"
@@ -4169,33 +4178,62 @@ def _format_evaluation_model_interval(snapshot: Snapshot, interval: Interval) ->
     return ""
 
 
-def _create_evaluation_model_annotation(snapshot: Snapshot, interval_info: t.Optional[str]) -> str:
+def _create_evaluation_model_annotation(
+    snapshot: Snapshot,
+    interval_info: t.Optional[str],
+    execution_stats: t.Optional[QueryExecutionStats],
+) -> str:
+    annotation = None
+    execution_stats_str = ""
+    if execution_stats:
+        rows_processed = execution_stats.total_rows_processed
+        if rows_processed:
+            # 1.00 and 1.0 to 1
+            rows_processed_str = metric(rows_processed).replace(".00", "").replace(".0", "")
+            execution_stats_str += f"{rows_processed_str} row{'s' if rows_processed > 1 else ''}"
+
+        bytes_processed = execution_stats.total_bytes_processed
+        execution_stats_str += (
+            f"{', ' if execution_stats_str else ''}{naturalsize(bytes_processed, binary=True)}"
+            if bytes_processed
+            else ""
+        )
+    execution_stats_str = f" ({execution_stats_str})" if execution_stats_str else ""
+
     if snapshot.is_audit:
-        return "run standalone audit"
-    if snapshot.is_model and snapshot.model.kind.is_external:
-        return "run external audits"
-    if snapshot.model.kind.is_seed:
-        return "insert seed file"
-    if snapshot.model.kind.is_full:
-        return "full refresh"
-    if snapshot.model.kind.is_view:
-        return "recreate view"
-    if snapshot.model.kind.is_incremental_by_unique_key:
-        return "insert/update rows"
-    if snapshot.model.kind.is_incremental_by_partition:
-        return "insert partitions"
-
-    return interval_info if interval_info else ""
-
-
-def _calculate_interval_str_len(snapshot: Snapshot, intervals: t.List[Interval]) -> int:
+        annotation = "run standalone audit"
+    if snapshot.is_model:
+        if snapshot.model.kind.is_external:
+            annotation = "run external audits"
+        if snapshot.model.kind.is_view:
+            annotation = "recreate view"
+        if snapshot.model.kind.is_seed:
+            annotation = f"insert seed file{execution_stats_str}"
+        if snapshot.model.kind.is_full:
+            annotation = f"full refresh{execution_stats_str}"
+        if snapshot.model.kind.is_incremental_by_unique_key:
+            annotation = f"insert/update rows{execution_stats_str}"
+        if snapshot.model.kind.is_incremental_by_partition:
+            annotation = f"insert partitions{execution_stats_str}"
+
+    if annotation:
+        return annotation
+
+    return f"{interval_info}{execution_stats_str}" if interval_info else ""
+
+
+def _calculate_interval_str_len(
+    snapshot: Snapshot,
+    intervals: t.List[Interval],
+    execution_stats: t.Optional[QueryExecutionStats] = None,
+) -> int:
     interval_str_len = 0
     for interval in intervals:
         interval_str_len = max(
             interval_str_len,
             len(
                 _create_evaluation_model_annotation(
-                    snapshot, _format_evaluation_model_interval(snapshot, interval)
+                    snapshot, _format_evaluation_model_interval(snapshot, interval), execution_stats
                 )
             ),
         )
@@ -4248,13 +4286,16 @@ def _calculate_audit_str_len(snapshot: Snapshot, audit_padding: int = 0) -> int:
 
 
 def _calculate_annotation_str_len(
-    batched_intervals: t.Dict[Snapshot, t.List[Interval]], audit_padding: int = 0
+    batched_intervals: t.Dict[Snapshot, t.List[Interval]],
+    audit_padding: int = 0,
+    execution_stats_len: int = 0,
 ) -> int:
     annotation_str_len = 0
     for snapshot, intervals in batched_intervals.items():
         annotation_str_len = max(
             annotation_str_len,
             _calculate_interval_str_len(snapshot, intervals)
-            + _calculate_audit_str_len(snapshot, audit_padding),
+            + _calculate_audit_str_len(snapshot, audit_padding)
+            + execution_stats_len,
         )
     return annotation_str_len

sqlmesh/core/engine_adapter/athena.py

@@ -45,6 +45,7 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin, RowDiffMixin):
     # >>> self._execute('/* test */ DESCRIBE foo')
     #     pyathena.error.OperationalError: FAILED: ParseException line 1:0 cannot recognize input near '/' '*' 'test'
     ATTACH_CORRELATION_ID = False
+    SUPPORTS_QUERY_EXECUTION_TRACKING = True
     SUPPORTED_DROP_CASCADE_OBJECT_KINDS = ["DATABASE", "SCHEMA"]
 
     def __init__(