Add rows tracking

Author: treysp

sqlmesh/core/console.py

@@ -428,6 +428,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        rows_processed: t.Optional[int] = None,
     ) -> None:
         """Updates the snapshot evaluation progress."""
 
@@ -575,6 +576,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        rows_processed: t.Optional[int] = None,
     ) -> None:
         pass
 
@@ -1011,7 +1013,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(" (XXXXXX rows processed)")
+                )
                 + 3  # brackets and opening escape backslash
             )
             self.evaluation_column_widths["name"] = max(
@@ -1056,6 +1060,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        rows_processed: t.Optional[int] = None,
     ) -> None:
         """Update the snapshot evaluation progress."""
         if (
@@ -1075,7 +1080,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), rows_processed
                 )
                 audits_str = ""
                 if num_audits_passed:
@@ -3639,6 +3644,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        rows_processed: t.Optional[int] = None,
     ) -> None:
         view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
 
@@ -3808,6 +3814,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        rows_processed: t.Optional[int] = None,
     ) -> None:
         message = f"Evaluating {snapshot.name} | batch={batch_idx} | duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"
 
@@ -3988,7 +3995,8 @@ def show_table_diff_summary(self, table_diff: TableDiff) -> None:
         self._write(f"Join On: {keys}")
 
 
-_CONSOLE: Console = NoopConsole()
+# _CONSOLE: Console = NoopConsole()
+_CONSOLE: Console = TerminalConsole()
 
 
 def set_console(console: Console) -> None:
@@ -4135,33 +4143,49 @@ 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], rows_processed: t.Optional[int]
+) -> str:
+    annotation = None
+    num_rows_processed = str(rows_processed) if rows_processed else ""
+    rows_processed_str = f" ({num_rows_processed} rows processed)" if num_rows_processed 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:
+            # no "processed" for seeds
+            seed_num_rows_inserted = (
+                f" ({num_rows_processed} rows inserted)" if num_rows_processed else ""
+            )
+            annotation = f"insert seed file{seed_num_rows_inserted}"
+        if snapshot.model.kind.is_full:
+            annotation = f"full refresh{rows_processed_str}"
+        if snapshot.model.kind.is_incremental_by_unique_key:
+            annotation = f"insert/update rows{rows_processed_str}"
+        if snapshot.model.kind.is_incremental_by_partition:
+            annotation = f"insert partitions{rows_processed_str}"
+
+    if annotation:
+        return annotation
+
+    return f"{interval_info}{rows_processed_str}" if interval_info else ""
+
+
+def _calculate_interval_str_len(
+    snapshot: Snapshot, intervals: t.List[Interval], rows_processed: t.Optional[int] = 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), rows_processed
                 )
             ),
         )
@@ -4214,13 +4238,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,
+    rows_processed_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)
+            + rows_processed_len,
         )
     return annotation_str_len

sqlmesh/core/engine_adapter/base.py

@@ -40,6 +40,7 @@
 )
 from sqlmesh.core.model.kind import TimeColumn
 from sqlmesh.core.schema_diff import SchemaDiffer
+from sqlmesh.core.execution_tracker import record_execution as track_execution_record
 from sqlmesh.utils import CorrelationId, columns_to_types_all_known, random_id
 from sqlmesh.utils.connection_pool import ConnectionPool, create_connection_pool
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_time_column
@@ -759,6 +760,7 @@ def _create_table_from_source_queries(
         table_description: t.Optional[str] = None,
         column_descriptions: t.Optional[t.Dict[str, str]] = None,
         table_kind: t.Optional[str] = None,
+        track_row_count: bool = True,
         **kwargs: t.Any,
     ) -> None:
         table = exp.to_table(table_name)
@@ -802,11 +804,15 @@ def _create_table_from_source_queries(
                             replace=replace,
                             table_description=table_description,
                             table_kind=table_kind,
+                            track_row_count=track_row_count,
                             **kwargs,
                         )
                     else:
                         self._insert_append_query(
-                            table_name, query, columns_to_types or self.columns(table)
+                            table_name,
+                            query,
+                            columns_to_types or self.columns(table),
+                            track_row_count=track_row_count,
                         )
 
         # Register comments with commands if the engine supports comments and we weren't able to
@@ -830,6 +836,7 @@ def _create_table(
         table_description: t.Optional[str] = None,
         column_descriptions: t.Optional[t.Dict[str, str]] = None,
         table_kind: t.Optional[str] = None,
+        track_row_count: bool = True,
         **kwargs: t.Any,
     ) -> None:
         self.execute(
@@ -846,7 +853,8 @@ def _create_table(
                 ),
                 table_kind=table_kind,
                 **kwargs,
-            )
+            ),
+            track_row_count=track_row_count,
         )
 
     def _build_create_table_exp(
@@ -1308,34 +1316,44 @@ def insert_append(
         table_name: TableName,
         query_or_df: QueryOrDF,
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
+        track_row_count: bool = True,
     ) -> None:
         source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
             query_or_df, columns_to_types, target_table=table_name
         )
-        self._insert_append_source_queries(table_name, source_queries, columns_to_types)
+        self._insert_append_source_queries(
+            table_name, source_queries, columns_to_types, track_row_count
+        )
 
     def _insert_append_source_queries(
         self,
         table_name: TableName,
         source_queries: t.List[SourceQuery],
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
+        track_row_count: bool = True,
     ) -> None:
         with self.transaction(condition=len(source_queries) > 0):
             columns_to_types = columns_to_types or self.columns(table_name)
             for source_query in source_queries:
                 with source_query as query:
-                    self._insert_append_query(table_name, query, columns_to_types)
+                    self._insert_append_query(
+                        table_name, query, columns_to_types, track_row_count=track_row_count
+                    )
 
     def _insert_append_query(
         self,
         table_name: TableName,
         query: Query,
         columns_to_types: t.Dict[str, exp.DataType],
         order_projections: bool = True,
+        track_row_count: bool = True,
     ) -> None:
         if order_projections:
             query = self._order_projections_and_filter(query, columns_to_types)
-        self.execute(exp.insert(query, table_name, columns=list(columns_to_types)))
+        self.execute(
+            exp.insert(query, table_name, columns=list(columns_to_types)),
+            track_row_count=track_row_count,
+        )
 
     def insert_overwrite_by_partition(
         self,
@@ -1459,7 +1477,7 @@ def _insert_overwrite_by_condition(
                         )
                         if insert_overwrite_strategy.is_replace_where:
                             insert_exp.set("where", where or exp.true())
-                        self.execute(insert_exp)
+                        self.execute(insert_exp, track_row_count=True)
 
     def update_table(
         self,
@@ -1480,7 +1498,7 @@ def _merge(
         using = exp.alias_(
             exp.Subquery(this=query), alias=MERGE_SOURCE_ALIAS, copy=False, table=True
         )
-        self.execute(exp.Merge(this=this, using=using, on=on, whens=whens))
+        self.execute(exp.Merge(this=this, using=using, on=on, whens=whens), track_row_count=True)
 
     def scd_type_2_by_time(
         self,
@@ -2214,6 +2232,7 @@ def execute(
         expressions: t.Union[str, exp.Expression, t.Sequence[exp.Expression]],
         ignore_unsupported_errors: bool = False,
         quote_identifiers: bool = True,
+        track_row_count: bool = False,
         **kwargs: t.Any,
     ) -> None:
         """Execute a sql query."""
@@ -2235,7 +2254,7 @@ def execute(
                     expression=e if isinstance(e, exp.Expression) else None,
                     quote_identifiers=quote_identifiers,
                 )
-                self._execute(sql, **kwargs)
+                self._execute(sql, track_row_count, **kwargs)
 
     def _attach_correlation_id(self, sql: str) -> str:
         if self.ATTACH_CORRELATION_ID and self.correlation_id:
@@ -2260,9 +2279,20 @@ def _log_sql(
 
         logger.log(self._execute_log_level, "Executing SQL: %s", sql_to_log)
 
-    def _execute(self, sql: str, **kwargs: t.Any) -> None:
+    def _execute(self, sql: str, track_row_count: bool = False, **kwargs: t.Any) -> None:
         self.cursor.execute(sql, **kwargs)
 
+        if track_row_count:
+            rowcount_raw = getattr(self.cursor, "rowcount", None)
+            rowcount = None
+            if rowcount_raw is not None:
+                try:
+                    rowcount = int(rowcount_raw)
+                except (TypeError, ValueError):
+                    pass
+
+            track_execution_record(sql, rowcount)
+
     @contextlib.contextmanager
     def temp_table(
         self,
@@ -2303,6 +2333,7 @@ def temp_table(
                 exists=True,
                 table_description=None,
                 column_descriptions=None,
+                track_row_count=False,
                 **kwargs,
             )
 
@@ -2547,7 +2578,7 @@ def _replace_by_key(
                     insert_statement.set("where", delete_filter)
                     insert_statement.set("this", exp.to_table(target_table))
 
-                self.execute(insert_statement)
+                self.execute(insert_statement, track_row_count=True)
             finally:
                 self.drop_table(temp_table)
 

sqlmesh/core/engine_adapter/bigquery.py

@@ -20,6 +20,7 @@
     SourceQuery,
     set_catalog,
 )
+from sqlmesh.core.execution_tracker import record_execution as track_execution_record
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.schema_diff import SchemaDiffer
 from sqlmesh.utils import optional_import
@@ -1036,6 +1037,7 @@ def _db_call(self, func: t.Callable[..., t.Any], *args: t.Any, **kwargs: t.Any)
     def _execute(
         self,
         sql: str,
+        track_row_count: bool = False,
         **kwargs: t.Any,
     ) -> None:
         """Execute a sql query."""
@@ -1081,6 +1083,9 @@ def _execute(
         self.cursor._set_rowcount(query_results)
         self.cursor._set_description(query_results.schema)
 
+        if track_row_count:
+            track_execution_record(sql, query_results.total_rows)
+
     def _get_data_objects(
         self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
     ) -> t.List[DataObject]:

sqlmesh/core/engine_adapter/clickhouse.py

@@ -282,7 +282,7 @@ def _insert_overwrite_by_condition(
                     )
 
                 try:
-                    self.execute(existing_records_insert_exp)
+                    self.execute(existing_records_insert_exp, track_row_count=True)
                 finally:
                     if table_partition_exp:
                         self.drop_table(partitions_temp_table_name)
@@ -469,6 +469,7 @@ def _create_table(
         table_description: t.Optional[str] = None,
         column_descriptions: t.Optional[t.Dict[str, str]] = None,
         table_kind: t.Optional[str] = None,
+        track_row_count: bool = True,
         **kwargs: t.Any,
     ) -> None:
         """Creates a table in the database.
@@ -505,6 +506,7 @@ def _create_table(
             column_descriptions,
             table_kind,
             empty_ctas=(self.engine_run_mode.is_cloud and expression is not None),
+            track_row_count=track_row_count,
             **kwargs,
         )
 

sqlmesh/core/engine_adapter/duckdb.py

@@ -153,6 +153,7 @@ def _create_table(
         table_description: t.Optional[str] = None,
         column_descriptions: t.Optional[t.Dict[str, str]] = None,
         table_kind: t.Optional[str] = None,
+        track_row_count: bool = True,
         **kwargs: t.Any,
     ) -> None:
         catalog = self.get_current_catalog()
@@ -176,6 +177,7 @@ def _create_table(
             table_description,
             column_descriptions,
             table_kind,
+            track_row_count=track_row_count,
             **kwargs,
         )