feat: Add progress bar to CLI from feast apply

franciscojavierarceo · franciscojavierarceo · commit e44dcb034e67 · 2026-01-16T10:35:38.000-05:00
Signed-off-by: Francisco Javier Arceo &lt;farceo@redhat.com&gt;
diff --git a/sdk/python/feast/cli/cli.py b/sdk/python/feast/cli/cli.py
@@ -23,6 +23,7 @@
 from colorama import Fore, Style
 from dateutil import parser
 from pygments import formatters, highlight, lexers
+from tqdm import tqdm
 
 from feast import utils
 from feast.cli.data_sources import data_sources_cmd
@@ -270,9 +271,17 @@ def plan_command(
     is_flag=True,
     help="Don't validate feature views. Use with caution as this skips important checks.",
 )
+@click.option(
+    "--no-progress",
+    is_flag=True,
+    help="Disable progress bars during apply operation.",
+)
 @click.pass_context
 def apply_total_command(
-    ctx: click.Context, skip_source_validation: bool, skip_feature_view_validation: bool
+    ctx: click.Context,
+    skip_source_validation: bool,
+    skip_feature_view_validation: bool,
+    no_progress: bool,
 ):
     """
     Create or update a feature store deployment
@@ -282,9 +291,21 @@ def apply_total_command(
     cli_check_repo(repo, fs_yaml_file)
 
     repo_config = load_repo_config(repo, fs_yaml_file)
+
+    # Create tqdm_builder for progress tracking
+    tqdm_builder = None
+    if not no_progress:
+
+        def tqdm_builder(length):
+            return tqdm(total=length, ncols=100)
+
     try:
         apply_total(
-            repo_config, repo, skip_source_validation, skip_feature_view_validation
+            repo_config,
+            repo,
+            skip_source_validation,
+            skip_feature_view_validation,
+            tqdm_builder=tqdm_builder,
         )
     except FeastProviderLoginError as e:
         print(str(e))
diff --git a/sdk/python/feast/diff/apply_progress.py b/sdk/python/feast/diff/apply_progress.py
@@ -0,0 +1,90 @@
+"""
+Progress tracking infrastructure for feast apply operations.
+
+This module provides the ApplyProgressContext class that manages progress bars
+during apply operations, following the same tqdm_builder pattern used in
+materialization operations.
+"""
+
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+from tqdm import tqdm
+
+
+@dataclass
+class ApplyProgressContext:
+    """
+    Context object for tracking progress during feast apply operations.
+
+    This class manages progress bars for the main phases of apply:
+    1. Planning changes (computing diffs)
+    2. Updating infrastructure (table creation/deletion)
+    3. Updating registry (metadata updates)
+
+    Follows the same tqdm_builder pattern used throughout Feast for consistency.
+    """
+
+    tqdm_builder: Callable[[int], tqdm]
+    current_phase: str = ""
+    overall_progress: Optional[tqdm] = None
+    phase_progress: Optional[tqdm] = None
+
+    # Phase tracking
+    total_phases: int = 3
+    completed_phases: int = 0
+
+    # Infrastructure operation tracking
+    total_infra_operations: int = 0
+    completed_infra_operations: int = 0
+
+    def start_overall_progress(self):
+        """Initialize the overall progress bar for apply phases."""
+        if self.overall_progress is None:
+            self.overall_progress = self.tqdm_builder(self.total_phases)
+            self.overall_progress.set_description("Applying changes")
+
+    def start_phase(self, phase_name: str, operations_count: int = 0):
+        """
+        Start tracking a new phase.
+
+        Args:
+            phase_name: Human-readable name of the phase
+            operations_count: Number of operations in this phase (0 for unknown)
+        """
+        self.current_phase = phase_name
+        if operations_count > 0:
+            self.phase_progress = self.tqdm_builder(operations_count)
+            self.phase_progress.set_description(f"{phase_name}")
+
+    def update_phase_progress(self, description: Optional[str] = None):
+        """
+        Update progress within the current phase.
+
+        Args:
+            description: Optional description of current operation
+        """
+        if self.phase_progress:
+            self.phase_progress.update(1)
+            if description:
+                self.phase_progress.set_description(
+                    f"{self.current_phase}: {description}"
+                )
+
+    def complete_phase(self):
+        """Mark current phase as complete and advance overall progress."""
+        if self.phase_progress:
+            self.phase_progress.close()
+            self.phase_progress = None
+        if self.overall_progress:
+            self.overall_progress.update(1)
+        self.completed_phases += 1
+
+    def cleanup(self):
+        """Clean up all progress bars. Should be called in finally blocks."""
+        if self.phase_progress:
+            self.phase_progress.close()
+            self.phase_progress = None
+        if self.overall_progress:
+            self.overall_progress.close()
+            self.overall_progress = None
diff --git a/sdk/python/feast/diff/infra_diff.py b/sdk/python/feast/diff/infra_diff.py
@@ -1,5 +1,8 @@
 from dataclasses import dataclass
-from typing import Generic, Iterable, List, Optional, Tuple, TypeVar
+from typing import TYPE_CHECKING, Generic, Iterable, List, Optional, Tuple, TypeVar
+
+if TYPE_CHECKING:
+    from feast.diff.apply_progress import ApplyProgressContext
 
 from feast.diff.property_diff import PropertyDiff, TransitionType
 from feast.infra.infra_object import (
@@ -33,8 +36,9 @@ class InfraDiff:
     def __init__(self):
         self.infra_object_diffs = []
 
-    def update(self):
+    def update(self, progress_ctx: Optional["ApplyProgressContext"] = None):
         """Apply the infrastructure changes specified in this object."""
+
         for infra_object_diff in self.infra_object_diffs:
             if infra_object_diff.transition_type in [
                 TransitionType.DELETE,
@@ -43,6 +47,10 @@ def update(self):
                 infra_object = InfraObject.from_proto(
                     infra_object_diff.current_infra_object
                 )
+                if progress_ctx:
+                    progress_ctx.update_phase_progress(
+                        f"Tearing down {infra_object_diff.name}"
+                    )
                 infra_object.teardown()
             elif infra_object_diff.transition_type in [
                 TransitionType.CREATE,
@@ -51,8 +59,19 @@ def update(self):
                 infra_object = InfraObject.from_proto(
                     infra_object_diff.new_infra_object
                 )
+                if progress_ctx:
+                    progress_ctx.update_phase_progress(
+                        f"Creating/updating {infra_object_diff.name}"
+                    )
                 infra_object.update()
 
+            # Update progress after each operation (except unchanged operations)
+            if (
+                progress_ctx
+                and infra_object_diff.transition_type != TransitionType.UNCHANGED
+            ):
+                progress_ctx.update_phase_progress()
+
     def to_string(self):
         from colorama import Fore, Style
 
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -18,6 +18,7 @@
 from datetime import datetime, timedelta
 from pathlib import Path
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
@@ -31,6 +32,9 @@
     cast,
 )
 
+if TYPE_CHECKING:
+    from feast.diff.apply_progress import ApplyProgressContext
+
 import pandas as pd
 import pyarrow as pa
 from colorama import Fore, Style
@@ -726,6 +730,7 @@ def plan(
         self,
         desired_repo_contents: RepoContents,
         skip_feature_view_validation: bool = False,
+        progress_ctx: Optional["ApplyProgressContext"] = None,
     ) -> Tuple[RegistryDiff, InfraDiff, Infra]:
         """Dry-run registering objects to metadata store.
 
@@ -793,6 +798,9 @@ def plan(
             self._registry, self.project, desired_repo_contents
         )
 
+        if progress_ctx:
+            progress_ctx.update_phase_progress("Computing infrastructure diff")
+
         # Compute the desired difference between the current infra, as stored in the registry,
         # and the desired infra.
         self._registry.refresh(project=self.project)
@@ -807,22 +815,45 @@ def plan(
         return registry_diff, infra_diff, new_infra
 
     def _apply_diffs(
-        self, registry_diff: RegistryDiff, infra_diff: InfraDiff, new_infra: Infra
+        self,
+        registry_diff: RegistryDiff,
+        infra_diff: InfraDiff,
+        new_infra: Infra,
+        progress_ctx: Optional["ApplyProgressContext"] = None,
     ):
         """Applies the given diffs to the metadata store and infrastructure.
 
         Args:
             registry_diff: The diff between the current registry and the desired registry.
             infra_diff: The diff between the current infra and the desired infra.
             new_infra: The desired infra.
+            progress_ctx: Optional progress context for tracking apply progress.
         """
-        infra_diff.update()
+        # Infrastructure phase
+        if progress_ctx:
+            infra_ops_count = len(infra_diff.infra_object_diffs)
+            progress_ctx.start_phase("Updating infrastructure", infra_ops_count)
+
+        infra_diff.update(progress_ctx=progress_ctx)
+
+        if progress_ctx:
+            progress_ctx.complete_phase()
+            progress_ctx.start_phase("Updating registry", 2)
+
+        # Registry phase
         apply_diff_to_registry(
             self._registry, registry_diff, self.project, commit=False
         )
 
+        if progress_ctx:
+            progress_ctx.update_phase_progress("Committing registry changes")
+
         self._registry.update_infra(new_infra, self.project, commit=True)
 
+        if progress_ctx:
+            progress_ctx.update_phase_progress("Registry update complete")
+            progress_ctx.complete_phase()
+
     def apply(
         self,
         objects: Union[
diff --git a/sdk/python/feast/repo_operations.py b/sdk/python/feast/repo_operations.py
@@ -10,10 +10,11 @@
 from importlib.abc import Loader
 from importlib.machinery import ModuleSpec
 from pathlib import Path
-from typing import List, Optional, Set, Union
+from typing import Callable, List, Optional, Set, Union
 
 import click
 from click.exceptions import BadParameter
+from tqdm import tqdm
 
 from feast import PushSource
 from feast.batch_feature_view import BatchFeatureView
@@ -342,6 +343,7 @@ def apply_total_with_repo_instance(
     repo: RepoContents,
     skip_source_validation: bool,
     skip_feature_view_validation: bool = False,
+    tqdm_builder: Optional[Callable[[int], tqdm]] = None,
 ):
     if not skip_source_validation:
         provider = store._get_provider()
@@ -358,22 +360,54 @@ def apply_total_with_repo_instance(
         views_to_delete,
     ) = extract_objects_for_apply_delete(project_name, registry, repo)
 
-    if store._should_use_plan():
-        registry_diff, infra_diff, new_infra = store.plan(
-            repo, skip_feature_view_validation=skip_feature_view_validation
-        )
-        click.echo(registry_diff.to_string())
+    # Create progress context if tqdm_builder is provided
+    progress_ctx = None
+    if tqdm_builder:
+        from feast.diff.apply_progress import ApplyProgressContext
 
-        store._apply_diffs(registry_diff, infra_diff, new_infra)
-        click.echo(infra_diff.to_string())
-    else:
-        store.apply(
-            all_to_apply,
-            objects_to_delete=all_to_delete,
-            partial=False,
-            skip_feature_view_validation=skip_feature_view_validation,
-        )
-        log_infra_changes(views_to_keep, views_to_delete)
+        progress_ctx = ApplyProgressContext(tqdm_builder=tqdm_builder)
+        progress_ctx.start_overall_progress()
+
+    try:
+        if store._should_use_plan():
+            # Planning phase
+            if progress_ctx:
+                progress_ctx.start_phase("Planning changes", 1)
+
+            registry_diff, infra_diff, new_infra = store.plan(
+                repo,
+                skip_feature_view_validation=skip_feature_view_validation,
+                progress_ctx=progress_ctx,
+            )
+            click.echo(registry_diff.to_string())
+
+            if progress_ctx:
+                progress_ctx.complete_phase()
+
+            # Apply phase
+            store._apply_diffs(
+                registry_diff, infra_diff, new_infra, progress_ctx=progress_ctx
+            )
+            click.echo(infra_diff.to_string())
+        else:
+            # Legacy apply path - simple single phase tracking
+            if progress_ctx:
+                progress_ctx.start_phase("Applying changes", 1)
+
+            store.apply(
+                all_to_apply,
+                objects_to_delete=all_to_delete,
+                partial=False,
+                skip_feature_view_validation=skip_feature_view_validation,
+            )
+            log_infra_changes(views_to_keep, views_to_delete)
+
+            if progress_ctx:
+                progress_ctx.update_phase_progress()
+                progress_ctx.complete_phase()
+    finally:
+        if progress_ctx:
+            progress_ctx.cleanup()
 
 
 def log_infra_changes(
@@ -416,6 +450,7 @@ def apply_total(
     repo_path: Path,
     skip_source_validation: bool,
     skip_feature_view_validation: bool = False,
+    tqdm_builder: Optional[Callable[[int], tqdm]] = None,
 ):
     os.chdir(repo_path)
     repo = _get_repo_contents(repo_path, repo_config.project, repo_config)
@@ -437,6 +472,7 @@ def apply_total(
             repo,
             skip_source_validation,
             skip_feature_view_validation,
+            tqdm_builder=tqdm_builder,
         )
 
 
