Add BQ support and track bytes processed

Author: treysp

sqlmesh/core/console.py

@@ -31,6 +31,7 @@
 from sqlmesh.core.environment import EnvironmentNamingInfo, EnvironmentSummary
 from sqlmesh.core.linter.rule import RuleViolation
 from sqlmesh.core.model import Model
+from sqlmesh.core.execution_tracker import QueryExecutionStats
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotChangeCategory,
@@ -428,7 +429,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        rows_processed: t.Optional[int] = None,
+        execution_stats: t.Optional[QueryExecutionStats] = None,
         auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         """Updates the snapshot evaluation progress."""
@@ -577,7 +578,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        rows_processed: t.Optional[int] = None,
+        execution_stats: t.Optional[QueryExecutionStats] = None,
         auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         pass
@@ -1016,7 +1017,7 @@ def start_evaluation_progress(
             # determine column widths
             self.evaluation_column_widths["annotation"] = (
                 _calculate_annotation_str_len(
-                    batched_intervals, self.AUDIT_PADDING, len(" (XXXXXX rows processed)")
+                    batched_intervals, self.AUDIT_PADDING, len(" (123.4m rows, 123.4 KiB)")
                 )
                 + 3  # brackets and opening escape backslash
             )
@@ -1062,7 +1063,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        rows_processed: t.Optional[int] = None,
+        execution_stats: t.Optional[QueryExecutionStats] = None,
         auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         """Update the snapshot evaluation progress."""
@@ -1083,7 +1084,7 @@ def update_snapshot_evaluation_progress(
                 ).ljust(self.evaluation_column_widths["name"])
 
                 annotation = _create_evaluation_model_annotation(
-                    snapshot, _format_evaluation_model_interval(snapshot, interval), rows_processed
+                    snapshot, _format_evaluation_model_interval(snapshot, interval), execution_stats
                 )
                 audits_str = ""
                 if num_audits_passed:
@@ -3647,7 +3648,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        rows_processed: t.Optional[int] = None,
+        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]
@@ -3818,7 +3819,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        rows_processed: t.Optional[int] = None,
+        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}"
@@ -4153,11 +4154,27 @@ def _format_evaluation_model_interval(snapshot: Snapshot, interval: Interval) ->
 
 
 def _create_evaluation_model_annotation(
-    snapshot: Snapshot, interval_info: t.Optional[str], rows_processed: t.Optional[int]
+    snapshot: Snapshot,
+    interval_info: t.Optional[str],
+    execution_stats: t.Optional[QueryExecutionStats],
 ) -> str:
     annotation = None
-    num_rows_processed = str(rows_processed) if rows_processed else ""
-    rows_processed_str = f" ({num_rows_processed} rows)" if num_rows_processed else ""
+    execution_stats_str = ""
+    if execution_stats:
+        rows_processed = execution_stats.total_rows_processed
+        execution_stats_str += (
+            f"{_abbreviate_integer_count(rows_processed)} row{'s' if rows_processed > 1 else ''}"
+            if rows_processed
+            else ""
+        )
+
+        bytes_processed = execution_stats.total_bytes_processed
+        execution_stats_str += (
+            f"{', ' if execution_stats_str else ''}{_format_bytes(bytes_processed)}"
+            if bytes_processed
+            else ""
+        )
+    execution_stats_str = f" ({execution_stats_str})" if execution_stats_str else ""
 
     if snapshot.is_audit:
         annotation = "run standalone audit"
@@ -4167,30 +4184,32 @@ def _create_evaluation_model_annotation(
         if snapshot.model.kind.is_view:
             annotation = "recreate view"
         if snapshot.model.kind.is_seed:
-            annotation = f"insert seed file{rows_processed_str}"
+            annotation = f"insert seed file{execution_stats_str}"
         if snapshot.model.kind.is_full:
-            annotation = f"full refresh{rows_processed_str}"
+            annotation = f"full refresh{execution_stats_str}"
         if snapshot.model.kind.is_incremental_by_unique_key:
-            annotation = f"insert/update rows{rows_processed_str}"
+            annotation = f"insert/update rows{execution_stats_str}"
         if snapshot.model.kind.is_incremental_by_partition:
-            annotation = f"insert partitions{rows_processed_str}"
+            annotation = f"insert partitions{execution_stats_str}"
 
     if annotation:
         return annotation
 
-    return f"{interval_info}{rows_processed_str}" if interval_info else ""
+    return f"{interval_info}{execution_stats_str}" if interval_info else ""
 
 
 def _calculate_interval_str_len(
-    snapshot: Snapshot, intervals: t.List[Interval], rows_processed: t.Optional[int] = None
+    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), rows_processed
+                    snapshot, _format_evaluation_model_interval(snapshot, interval), execution_stats
                 )
             ),
         )
@@ -4245,14 +4264,50 @@ 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,
-    rows_processed_len: 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)
-            + rows_processed_len,
+            + execution_stats_len,
         )
     return annotation_str_len
+
+
+# Convert number of bytes to a human-readable string
+# https://github.com/dbt-labs/dbt-adapters/blob/34fd178539dcb6f82e18e738adc03de7784c032f/dbt-bigquery/src/dbt/adapters/bigquery/connections.py#L165
+def _format_bytes(num_bytes: t.Optional[int]) -> str:
+    if num_bytes and num_bytes > 0:
+        if num_bytes < 1024:
+            return f"{num_bytes} Bytes"
+
+        num_bytes_float = float(num_bytes) / 1024.0
+        for unit in ["KiB", "MiB", "GiB", "TiB", "PiB"]:
+            if num_bytes_float < 1024.0:
+                return f"{num_bytes_float:3.1f} {unit}"
+            num_bytes_float /= 1024.0
+
+        num_bytes_float *= 1024.0  # undo last division in loop
+        return f"{num_bytes_float:3.1f} {unit}"
+    return ""
+
+
+# Abbreviate integer count. Example: 1,000,000,000 -> 1b
+# https://github.com/dbt-labs/dbt-adapters/blob/34fd178539dcb6f82e18e738adc03de7784c032f/dbt-bigquery/src/dbt/adapters/bigquery/connections.py#L178
+def _abbreviate_integer_count(count: t.Optional[int]) -> str:
+    if count and count > 0:
+        if count < 1000:
+            return str(count)
+
+        count_float = float(count) / 1000.0
+        for unit in ["k", "m", "b", "t"]:
+            if count_float < 1000.0:
+                return f"{count_float:3.1f}{unit}".strip()
+            count_float /= 1000.0
+
+        count_float *= 1000.0  # undo last division in loop
+        return f"{count_float:3.1f}{unit}".strip()
+    return ""

sqlmesh/core/engine_adapter/base.py

@@ -837,7 +837,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,
+        track_execution_stats: bool = True,
         **kwargs: t.Any,
     ) -> None:
         table = exp.to_table(table_name)
@@ -883,15 +883,15 @@ def _create_table_from_source_queries(
                             replace=replace,
                             table_description=table_description,
                             table_kind=table_kind,
-                            track_row_count=track_row_count,
+                            track_execution_stats=track_execution_stats,
                             **kwargs,
                         )
                     else:
                         self._insert_append_query(
                             table_name,
                             query,
                             target_columns_to_types or self.columns(table),
-                            track_row_count=track_row_count,
+                            track_execution_stats=track_execution_stats,
                         )
 
         # Register comments with commands if the engine supports comments and we weren't able to
@@ -915,7 +915,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,
+        track_execution_stats: bool = True,
         **kwargs: t.Any,
     ) -> None:
         self.execute(
@@ -933,7 +933,7 @@ def _create_table(
                 table_kind=table_kind,
                 **kwargs,
             ),
-            track_row_count=track_row_count,
+            track_execution_stats=track_execution_stats,
         )
 
     def _build_create_table_exp(
@@ -1408,7 +1408,7 @@ def insert_append(
         table_name: TableName,
         query_or_df: QueryOrDF,
         target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
-        track_row_count: bool = True,
+        track_execution_stats: bool = True,
         source_columns: t.Optional[t.List[str]] = None,
     ) -> None:
         source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types(
@@ -1418,22 +1418,25 @@ def insert_append(
             source_columns=source_columns,
         )
         self._insert_append_source_queries(
-            table_name, source_queries, target_columns_to_types, track_row_count
+            table_name, source_queries, target_columns_to_types, track_execution_stats
         )
 
     def _insert_append_source_queries(
         self,
         table_name: TableName,
         source_queries: t.List[SourceQuery],
         target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
-        track_row_count: bool = True,
+        track_execution_stats: bool = True,
     ) -> None:
         with self.transaction(condition=len(source_queries) > 0):
             target_columns_to_types = target_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, target_columns_to_types, track_row_count=track_row_count
+                        table_name,
+                        query,
+                        target_columns_to_types,
+                        track_execution_stats=track_execution_stats,
                     )
 
     def _insert_append_query(
@@ -1442,13 +1445,13 @@ def _insert_append_query(
         query: Query,
         target_columns_to_types: t.Dict[str, exp.DataType],
         order_projections: bool = True,
-        track_row_count: bool = True,
+        track_execution_stats: bool = True,
     ) -> None:
         if order_projections:
             query = self._order_projections_and_filter(query, target_columns_to_types)
         self.execute(
             exp.insert(query, table_name, columns=list(target_columns_to_types)),
-            track_row_count=track_row_count,
+            track_execution_stats=track_execution_stats,
         )
 
     def insert_overwrite_by_partition(
@@ -1591,7 +1594,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, track_row_count=True)
+                        self.execute(insert_exp, track_execution_stats=True)
 
     def update_table(
         self,
@@ -1612,7 +1615,9 @@ 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), track_row_count=True)
+        self.execute(
+            exp.Merge(this=this, using=using, on=on, whens=whens), track_execution_stats=True
+        )
 
     def scd_type_2_by_time(
         self,
@@ -2361,7 +2366,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,
+        track_execution_stats: bool = False,
         **kwargs: t.Any,
     ) -> None:
         """Execute a sql query."""
@@ -2383,7 +2388,7 @@ def execute(
                     expression=e if isinstance(e, exp.Expression) else None,
                     quote_identifiers=quote_identifiers,
                 )
-                self._execute(sql, track_row_count, **kwargs)
+                self._execute(sql, track_execution_stats, **kwargs)
 
     def _attach_correlation_id(self, sql: str) -> str:
         if self.ATTACH_CORRELATION_ID and self.correlation_id:
@@ -2408,12 +2413,12 @@ def _log_sql(
 
         logger.log(self._execute_log_level, "Executing SQL: %s", sql_to_log)
 
-    def _execute(self, sql: str, track_row_count: bool = False, **kwargs: t.Any) -> None:
+    def _execute(self, sql: str, track_execution_stats: bool = False, **kwargs: t.Any) -> None:
         self.cursor.execute(sql, **kwargs)
 
         if (
             self.SUPPORTS_QUERY_EXECUTION_TRACKING
-            and track_row_count
+            and track_execution_stats
             and QueryExecutionTracker.is_tracking()
         ):
             rowcount_raw = getattr(self.cursor, "rowcount", None)
@@ -2424,7 +2429,7 @@ def _execute(self, sql: str, track_row_count: bool = False, **kwargs: t.Any) ->
                 except (TypeError, ValueError):
                     pass
 
-            QueryExecutionTracker.record_execution(sql, rowcount)
+            QueryExecutionTracker.record_execution(sql, rowcount, None)
 
     @contextlib.contextmanager
     def temp_table(
@@ -2470,7 +2475,7 @@ def temp_table(
                 exists=True,
                 table_description=None,
                 column_descriptions=None,
-                track_row_count=False,
+                track_execution_stats=False,
                 **kwargs,
             )
 
@@ -2722,7 +2727,7 @@ def _replace_by_key(
                     insert_statement.set("where", delete_filter)
                     insert_statement.set("this", exp.to_table(target_table))
 
-                self.execute(insert_statement, track_row_count=True)
+                self.execute(insert_statement, track_execution_stats=True)
             finally:
                 self.drop_table(temp_table)