Feat: add progress bar for snapshot creation (#1157)

* Feat: add progress bar for snapshot creation

* Rephrase progress bar status message

* Cleanup

* Cleanup

Author: georgesittas

sqlmesh/core/console.py

@@ -44,38 +44,50 @@ class Console(abc.ABC):
     with them when their input is needed."""
 
     @abc.abstractmethod
-    def start_snapshot_progress(
+    def start_evaluation_progress(
         self, snapshot: Snapshot, total_batches: int, environment: str
     ) -> None:
-        """Indicates that a new load progress has begun."""
+        """Indicates that a new snapshot evaluation progress has begun."""
 
     @abc.abstractmethod
-    def update_snapshot_progress(self, snapshot_name: str, num_batches: int) -> None:
-        """Update snapshot progress."""
+    def update_evaluation_progress(self, snapshot_name: str, num_batches: int) -> None:
+        """Update the snapshot evaluation progress."""
 
     @abc.abstractmethod
-    def stop_snapshot_progress(self, success: bool = True) -> None:
-        """Stop the load progress."""
+    def stop_evaluation_progress(self, success: bool = True) -> None:
+        """Stop the snapshot evaluation progress."""
+
+    @abc.abstractmethod
+    def start_creation_progress(self, environment: str, total_tasks: int) -> None:
+        """Indicates that a new snapshot creation progress has begun."""
+
+    @abc.abstractmethod
+    def update_creation_progress(self, num_tasks: int) -> None:
+        """Update the snapshot creation progress."""
+
+    @abc.abstractmethod
+    def stop_creation_progress(self, success: bool = True) -> None:
+        """Stop the snapshot creation progress."""
 
     @abc.abstractmethod
     def start_promotion_progress(self, environment: str, total_tasks: int) -> None:
-        """Indicates that a new promotion progress has begun."""
+        """Indicates that a new snapshot promotion progress has begun."""
 
     @abc.abstractmethod
     def update_promotion_progress(self, num_tasks: int) -> None:
-        """Update promotion progress."""
+        """Update the snapshot promotion progress."""
 
     @abc.abstractmethod
     def stop_promotion_progress(self, success: bool = True) -> None:
-        """Stop the promotion progress."""
+        """Stop the snapshot promotion progress."""
 
     @abc.abstractmethod
     def start_migration_progress(self, total_tasks: int) -> None:
         """Indicates that a new migration progress has begun."""
 
     @abc.abstractmethod
     def update_migration_progress(self, num_tasks: int) -> None:
-        """Update migration progress."""
+        """Update the migration progress."""
 
     @abc.abstractmethod
     def stop_migration_progress(self, success: bool = True) -> None:
@@ -85,7 +97,7 @@ def stop_migration_progress(self, success: bool = True) -> None:
     def show_model_difference_summary(
         self, context_diff: ContextDiff, detailed: bool = False
     ) -> None:
-        """Displays a summary of differences for the given models"""
+        """Displays a summary of differences for the given models."""
 
     @abc.abstractmethod
     def plan(self, plan: Plan, auto_apply: bool) -> None:
@@ -151,6 +163,8 @@ def __init__(self, console: t.Optional[RichConsole] = None, **kwargs: t.Any) ->
         self.console: RichConsole = console or srich.console
         self.evaluation_progress: t.Optional[Progress] = None
         self.evaluation_tasks: t.Dict[str, t.Tuple[TaskID, int]] = {}
+        self.creation_progress: t.Optional[Progress] = None
+        self.creation_task: t.Optional[TaskID] = None
         self.promotion_progress: t.Optional[Progress] = None
         self.promotion_task: t.Optional[TaskID] = None
         self.migration_progress: t.Optional[Progress] = None
@@ -166,10 +180,10 @@ def _prompt(self, message: str, **kwargs: t.Any) -> t.Any:
     def _confirm(self, message: str, **kwargs: t.Any) -> bool:
         return Confirm.ask(message, console=self.console, **kwargs)
 
-    def start_snapshot_progress(
+    def start_evaluation_progress(
         self, snapshot: Snapshot, total_batches: int, environment: str
     ) -> None:
-        """Indicates that a new load progress has begun."""
+        """Indicates that a new snapshot evaluation progress has begun."""
         if not self.evaluation_progress:
             self.evaluation_progress = Progress(
                 TextColumn("[bold blue]{task.fields[view_name]}", justify="right"),
@@ -194,23 +208,59 @@ def start_snapshot_progress(
             total_batches,
         )
 
-    def update_snapshot_progress(self, snapshot_name: str, num_batches: int) -> None:
-        """Update snapshot progress."""
+    def update_evaluation_progress(self, snapshot_name: str, num_batches: int) -> None:
+        """Update the snapshot evaluation progress."""
         if self.evaluation_progress and self.evaluation_tasks:
             task_id = self.evaluation_tasks[snapshot_name][0]
             self.evaluation_progress.update(task_id, refresh=True, advance=num_batches)
 
-    def stop_snapshot_progress(self, success: bool = True) -> None:
-        """Stop the load progress"""
+    def stop_evaluation_progress(self, success: bool = True) -> None:
+        """Stop the snapshot evaluation progress."""
         self.evaluation_tasks = {}
         if self.evaluation_progress:
             self.evaluation_progress.stop()
             self.evaluation_progress = None
             if success:
                 self.log_success("All model batches have been executed successfully")
 
+    def start_creation_progress(self, environment: str, total_tasks: int) -> None:
+        """Indicates that a new creation progress has begun."""
+        if self.creation_progress is None:
+            self.creation_progress = Progress(
+                TextColumn(
+                    f"[bold blue]Creating new model versions for '{environment}'", justify="right"
+                ),
+                BarColumn(bar_width=40),
+                "[progress.percentage]{task.percentage:>3.1f}%",
+                "•",
+                srich.BatchColumn(),
+                "•",
+                TimeElapsedColumn(),
+                console=self.console,
+            )
+
+            self.creation_progress.start()
+            self.creation_task = self.creation_progress.add_task(
+                f"Creating new model versions for {environment}...",
+                total=total_tasks,
+            )
+
+    def update_creation_progress(self, num_tasks: int) -> None:
+        """Update the snapshot creation progress."""
+        if self.creation_progress is not None and self.creation_task is not None:
+            self.creation_progress.update(self.creation_task, refresh=True, advance=num_tasks)
+
+    def stop_creation_progress(self, success: bool = True) -> None:
+        """Stop the snapshot creation progress."""
+        self.creation_task = None
+        if self.creation_progress is not None:
+            self.creation_progress.stop()
+            self.creation_progress = None
+            if success:
+                self.log_success("All model versions have been created successfully")
+
     def start_promotion_progress(self, environment: str, total_tasks: int) -> None:
-        """Indicates that a new promotion progress has begun."""
+        """Indicates that a new snapshot promotion progress has begun."""
         if self.promotion_progress is None:
             self.promotion_progress = Progress(
                 TextColumn(f"[bold blue]Virtually Updating '{environment}'", justify="right"),
@@ -228,12 +278,12 @@ def start_promotion_progress(self, environment: str, total_tasks: int) -> None:
             )
 
     def update_promotion_progress(self, num_tasks: int) -> None:
-        """Update promotion progress."""
+        """Update the snapshot promotion progress."""
         if self.promotion_progress is not None and self.promotion_task is not None:
             self.promotion_progress.update(self.promotion_task, refresh=True, advance=num_tasks)
 
     def stop_promotion_progress(self, success: bool = True) -> None:
-        """Stop the promotion progress"""
+        """Stop the snapshot promotion progress."""
         self.promotion_task = None
         if self.promotion_progress is not None:
             self.promotion_progress.stop()
@@ -262,7 +312,7 @@ def start_migration_progress(self, total_tasks: int) -> None:
             )
 
     def update_migration_progress(self, num_tasks: int) -> None:
-        """Update migration progress."""
+        """Update the migration progress."""
         if self.migration_progress is not None and self.migration_task is not None:
             self.migration_progress.update(self.migration_task, refresh=True, advance=num_tasks)
 
@@ -361,7 +411,7 @@ def _show_options_after_categorization(self, plan: Plan, auto_apply: bool) -> No
             self._prompt_promote(plan)
 
     def _prompt_categorize(self, plan: Plan, auto_apply: bool) -> None:
-        """Get the user's change category for the directly modified models"""
+        """Get the user's change category for the directly modified models."""
         self.show_model_difference_summary(plan.context_diff)
 
         self._show_categorized_snapshots(plan)
@@ -398,7 +448,7 @@ def _show_categorized_snapshots(self, plan: Plan) -> None:
             self._print(tree)
 
     def _show_missing_dates(self, plan: Plan) -> None:
-        """Displays the models with missing dates"""
+        """Displays the models with missing dates."""
         if not plan.missing_intervals:
             return
         backfill = Tree("[bold]Models needing backfill (missing dates):")
@@ -608,13 +658,14 @@ def _snapshot_change_choices(
 
 
 def add_to_layout_widget(target_widget: LayoutWidget, *widgets: widgets.Widget) -> LayoutWidget:
-    """Helper function to add a widget to a layout widget
+    """Helper function to add a widget to a layout widget.
+
     Args:
-        target_widget: The layout widget to add the other widget(s) to
-        *widgets: The widgets to add to the layout widget
+        target_widget: The layout widget to add the other widget(s) to.
+        *widgets: The widgets to add to the layout widget.
 
     Returns:
-        The layout widget with the children added
+        The layout widget with the children added.
     """
     target_widget.children += tuple(widgets)
     return target_widget
@@ -914,7 +965,7 @@ def show_model_difference_summary(
         """Shows a summary of the differences.
 
         Args:
-            context_diff: The context diff to use to print the summary
+            context_diff: The context diff to use to print the summary.
             detailed: Show the actual SQL differences if True.
         """
         if context_diff.is_new_environment:
@@ -972,7 +1023,7 @@ def show_model_difference_summary(
                 self._print("\n")
 
     def _show_missing_dates(self, plan: Plan) -> None:
-        """Displays the models with missing dates"""
+        """Displays the models with missing dates."""
         if not plan.missing_intervals:
             return
         self._print("**Models needing backfill (missing dates):**\n\n")
@@ -1048,17 +1099,17 @@ def _confirm(self, message: str, **kwargs: t.Any) -> bool:
         self._print(message)
         return super()._confirm("", **kwargs)
 
-    def start_snapshot_progress(
+    def start_evaluation_progress(
         self, snapshot: Snapshot, total_batches: int, environment: str
     ) -> None:
-        """Indicates that a new load progress has begun."""
+        """Indicates that a new snapshot evaluation progress has begun."""
         if not self.evaluation_batch_progress.get(snapshot.name):
             view_name = snapshot.qualified_view_name.for_environment(environment)
             self.evaluation_batch_progress[snapshot.name] = (view_name, 0, total_batches)
             print(f"Starting '{view_name}', Total batches: {total_batches}")
 
-    def update_snapshot_progress(self, snapshot_name: str, num_batches: int) -> None:
-        """Update snapshot progress."""
+    def update_evaluation_progress(self, snapshot_name: str, num_batches: int) -> None:
+        """Update the snapshot evaluation progress."""
         view_name, loaded_batches, total_batches = self.evaluation_batch_progress[snapshot_name]
         loaded_batches += 1
         self.evaluation_batch_progress[snapshot_name] = (view_name, loaded_batches, total_batches)
@@ -1072,26 +1123,26 @@ def update_snapshot_progress(self, snapshot_name: str, num_batches: int) -> None
             total = len(self.evaluation_batch_progress)
             print(f"Completed Loading {total_finished_loading}/{total} Models")
 
-    def stop_snapshot_progress(self, success: bool = True) -> None:
-        """Stop the load progress"""
+    def stop_evaluation_progress(self, success: bool = True) -> None:
+        """Stop the snapshot evaluation progress."""
         self.evaluation_batch_progress = {}
         print(f"Loading {'succeeded' if success else 'failed'}")
 
     def start_promotion_progress(self, environment: str, total_tasks: int) -> None:
-        """Indicates that a new promotion progress has begun."""
+        """Indicates that a new snapshot promotion progress has begun."""
         self.promotion_status = (0, total_tasks)
         print(f"Virtually Updating '{environment}'")
 
     def update_promotion_progress(self, num_tasks: int) -> None:
-        """Update promotion progress."""
+        """Update the snapshot promotion progress."""
         num_promotions, total_promotions = self.promotion_status
         num_promotions += num_tasks
         self.promotion_status = (num_promotions, total_promotions)
         if num_promotions % 5 == 0:
             print(f"Virtually Updated {num_promotions}/{total_promotions}")
 
     def stop_promotion_progress(self, success: bool = True) -> None:
-        """Stop the promotion progress"""
+        """Stop the snapshot promotion progress."""
         print(f"Virtual Update {'succeeded' if success else 'failed'}")
 
 

sqlmesh/core/context.py

@@ -250,7 +250,9 @@ def __init__(
         self._test_engine_adapter = test_connection_config.create_engine_adapter()
 
         self.snapshot_evaluator = SnapshotEvaluator(
-            self.engine_adapter, ddl_concurrent_tasks=self.concurrent_tasks
+            self.engine_adapter,
+            ddl_concurrent_tasks=self.concurrent_tasks,
+            console=self.console,
         )
 
         self._provided_state_sync: t.Optional[StateSync] = state_sync

sqlmesh/core/plan/evaluator.py

@@ -116,9 +116,24 @@ def _push(self, plan: Plan) -> None:
         Args:
             plan: The plan to source snapshots from.
         """
-        self.snapshot_evaluator.create(
-            plan.new_snapshots, {s.snapshot_id: s for s in plan.snapshots}
-        )
+        snapshot_id_to_snapshot = {s.snapshot_id: s for s in plan.snapshots}
+
+        self.console.start_creation_progress(plan.environment.name, len(snapshot_id_to_snapshot))
+
+        def on_complete(snapshot: SnapshotInfoLike) -> None:
+            self.console.update_creation_progress(1)
+
+        completed = False
+        try:
+            self.snapshot_evaluator.create(
+                plan.new_snapshots,
+                snapshot_id_to_snapshot,
+                on_complete=on_complete,
+            )
+            completed = True
+        finally:
+            self.console.stop_creation_progress(success=completed)
+
         self.state_sync.push_snapshots(plan.new_snapshots)
 
     def _promote(self, plan: Plan) -> None:

sqlmesh/core/scheduler.py

@@ -63,7 +63,7 @@ def __init__(
         self.snapshot_per_version = _resolve_one_snapshot_per_version(self.snapshots.values())
         self.snapshot_evaluator = snapshot_evaluator
         self.max_workers = max_workers
-        self.console: Console = console or get_console()
+        self.console = console or get_console()
         self.notification_target_manager = (
             notification_target_manager or NotificationTargetManager()
         )
@@ -167,7 +167,7 @@ def evaluate(
                 )
             raise e
         self.state_sync.add_interval(snapshot, start, end, is_dev=is_dev)
-        self.console.update_snapshot_progress(snapshot.name, 1)
+        self.console.update_evaluation_progress(snapshot.name, 1)
 
     def run(
         self,
@@ -203,7 +203,7 @@ def run(
                 continue
             visited.add(snapshot)
             intervals = batches[snapshot]
-            self.console.start_snapshot_progress(snapshot, len(intervals), environment)
+            self.console.start_evaluation_progress(snapshot, len(intervals), environment)
 
         def evaluate_node(node: SchedulingUnit) -> None:
             assert latest
@@ -221,7 +221,7 @@ def evaluate_node(node: SchedulingUnit) -> None:
         finally:
             self.state_sync.recycle()
 
-        self.console.stop_snapshot_progress(success=not errors)
+        self.console.stop_evaluation_progress(success=not errors)
 
         for error in errors:
             sid = error.node[0]