sqlmesh.core.console

   1from __future__ import annotations
   2
   3import abc
   4import datetime
   5import typing as t
   6import unittest
   7import uuid
   8
   9from hyperscript import h
  10from rich.console import Console as RichConsole
  11from rich.live import Live
  12from rich.progress import (
  13    BarColumn,
  14    Progress,
  15    SpinnerColumn,
  16    TaskID,
  17    TextColumn,
  18    TimeElapsedColumn,
  19)
  20from rich.prompt import Confirm, Prompt
  21from rich.status import Status
  22from rich.syntax import Syntax
  23from rich.table import Table
  24from rich.tree import Tree
  25
  26from sqlmesh.core.environment import EnvironmentNamingInfo
  27from sqlmesh.core.snapshot import (
  28    Snapshot,
  29    SnapshotChangeCategory,
  30    SnapshotId,
  31    SnapshotInfoLike,
  32    start_date,
  33)
  34from sqlmesh.core.test import ModelTest
  35from sqlmesh.utils import rich as srich
  36from sqlmesh.utils.date import time_like_to_str, to_date, yesterday_ds
  37
  38if t.TYPE_CHECKING:
  39    import ipywidgets as widgets
  40
  41    from sqlmesh.core.context_diff import ContextDiff
  42    from sqlmesh.core.plan import Plan, PlanBuilder
  43    from sqlmesh.core.table_diff import RowDiff, SchemaDiff
  44
  45    LayoutWidget = t.TypeVar("LayoutWidget", bound=t.Union[widgets.VBox, widgets.HBox])
  46
  47
  48SNAPSHOT_CHANGE_CATEGORY_STR = {
  49    None: "Unknown",
  50    SnapshotChangeCategory.BREAKING: "Breaking",
  51    SnapshotChangeCategory.NON_BREAKING: "Non-breaking",
  52    SnapshotChangeCategory.FORWARD_ONLY: "Forward-only",
  53    SnapshotChangeCategory.INDIRECT_BREAKING: "Indirect Breaking",
  54    SnapshotChangeCategory.INDIRECT_NON_BREAKING: "Indirect Non-breaking",
  55}
  56
  57
  58class Console(abc.ABC):
  59    """Abstract base class for defining classes used for displaying information to the user and also interact
  60    with them when their input is needed."""
  61
  62    @abc.abstractmethod
  63    def start_plan_evaluation(self, plan: Plan) -> None:
  64        """Indicates that a new evaluation has begun."""
  65
  66    @abc.abstractmethod
  67    def stop_plan_evaluation(self) -> None:
  68        """Indicates that the evaluation has ended."""
  69
  70    @abc.abstractmethod
  71    def start_evaluation_progress(
  72        self,
  73        batches: t.Dict[Snapshot, int],
  74        environment_naming_info: EnvironmentNamingInfo,
  75        default_catalog: t.Optional[str],
  76    ) -> None:
  77        """Indicates that a new snapshot evaluation progress has begun."""
  78
  79    @abc.abstractmethod
  80    def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
  81        """Starts the snapshot evaluation progress."""
  82
  83    @abc.abstractmethod
  84    def update_snapshot_evaluation_progress(
  85        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
  86    ) -> None:
  87        """Updates the snapshot evaluation progress."""
  88
  89    @abc.abstractmethod
  90    def stop_evaluation_progress(self, success: bool = True) -> None:
  91        """Stops the snapshot evaluation progress."""
  92
  93    @abc.abstractmethod
  94    def start_creation_progress(
  95        self,
  96        total_tasks: int,
  97        environment_naming_info: EnvironmentNamingInfo,
  98        default_catalog: t.Optional[str],
  99    ) -> None:
 100        """Indicates that a new snapshot creation progress has begun."""
 101
 102    @abc.abstractmethod
 103    def update_creation_progress(self, snapshot: SnapshotInfoLike) -> None:
 104        """Update the snapshot creation progress."""
 105
 106    @abc.abstractmethod
 107    def stop_creation_progress(self, success: bool = True) -> None:
 108        """Stop the snapshot creation progress."""
 109
 110    @abc.abstractmethod
 111    def update_cleanup_progress(self, object_name: str) -> None:
 112        """Update the snapshot cleanup progress."""
 113
 114    @abc.abstractmethod
 115    def start_promotion_progress(
 116        self,
 117        total_tasks: int,
 118        environment_naming_info: EnvironmentNamingInfo,
 119        default_catalog: t.Optional[str],
 120    ) -> None:
 121        """Indicates that a new snapshot promotion progress has begun."""
 122
 123    @abc.abstractmethod
 124    def update_promotion_progress(self, snapshot: SnapshotInfoLike, promoted: bool) -> None:
 125        """Update the snapshot promotion progress."""
 126
 127    @abc.abstractmethod
 128    def stop_promotion_progress(self, success: bool = True) -> None:
 129        """Stop the snapshot promotion progress."""
 130
 131    @abc.abstractmethod
 132    def start_migration_progress(self, total_tasks: int) -> None:
 133        """Indicates that a new migration progress has begun."""
 134
 135    @abc.abstractmethod
 136    def update_migration_progress(self, num_tasks: int) -> None:
 137        """Update the migration progress."""
 138
 139    @abc.abstractmethod
 140    def stop_migration_progress(self, success: bool = True) -> None:
 141        """Stop the migration progress."""
 142
 143    @abc.abstractmethod
 144    def show_model_difference_summary(
 145        self,
 146        context_diff: ContextDiff,
 147        environment_naming_info: EnvironmentNamingInfo,
 148        default_catalog: t.Optional[str],
 149        no_diff: bool = True,
 150        ignored_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None,
 151    ) -> None:
 152        """Displays a summary of differences for the given models."""
 153
 154    @abc.abstractmethod
 155    def plan(
 156        self,
 157        plan_builder: PlanBuilder,
 158        auto_apply: bool,
 159        default_catalog: t.Optional[str],
 160        no_diff: bool = False,
 161        no_prompts: bool = False,
 162    ) -> None:
 163        """The main plan flow.
 164
 165        The console should present the user with choices on how to backfill and version the snapshots
 166        of a plan.
 167
 168        Args:
 169            plan: The plan to make choices for.
 170            auto_apply: Whether to automatically apply the plan after all choices have been made.
 171            no_diff: Hide text differences for changed models.
 172            no_prompts: Whether to disable interactive prompts for the backfill time range. Please note that
 173                if this flag is set to true and there are uncategorized changes the plan creation will
 174                fail. Default: False
 175        """
 176
 177    @abc.abstractmethod
 178    def log_test_results(
 179        self, result: unittest.result.TestResult, output: str, target_dialect: str
 180    ) -> None:
 181        """Display the test result and output.
 182
 183        Args:
 184            result: The unittest test result that contains metrics like num success, fails, ect.
 185            output: The generated output from the unittest.
 186            target_dialect: The dialect that tests were run against. Assumes all tests run against the same dialect.
 187        """
 188
 189    @abc.abstractmethod
 190    def show_sql(self, sql: str) -> None:
 191        """Display to the user SQL."""
 192
 193    @abc.abstractmethod
 194    def log_status_update(self, message: str) -> None:
 195        """Display general status update to the user."""
 196
 197    @abc.abstractmethod
 198    def log_error(self, message: str) -> None:
 199        """Display error info to the user."""
 200
 201    @abc.abstractmethod
 202    def log_success(self, message: str) -> None:
 203        """Display a general successful message to the user."""
 204
 205    @abc.abstractmethod
 206    def loading_start(self, message: t.Optional[str] = None) -> uuid.UUID:
 207        """Starts loading and returns a unique ID that can be used to stop the loading. Optionally can display a message."""
 208
 209    @abc.abstractmethod
 210    def loading_stop(self, id: uuid.UUID) -> None:
 211        """Stop loading for the given id."""
 212
 213    @abc.abstractmethod
 214    def show_schema_diff(self, schema_diff: SchemaDiff) -> None:
 215        """Show table schema diff."""
 216
 217    @abc.abstractmethod
 218    def show_row_diff(self, row_diff: RowDiff, show_sample: bool = True) -> None:
 219        """Show table summary diff."""
 220
 221
 222class TerminalConsole(Console):
 223    """A rich based implementation of the console."""
 224
 225    def __init__(
 226        self, console: t.Optional[RichConsole] = None, verbose: bool = False, **kwargs: t.Any
 227    ) -> None:
 228        self.console: RichConsole = console or srich.console
 229
 230        self.evaluation_progress_live: t.Optional[Live] = None
 231        self.evaluation_total_progress: t.Optional[Progress] = None
 232        self.evaluation_total_task: t.Optional[TaskID] = None
 233        self.evaluation_model_progress: t.Optional[Progress] = None
 234        self.evaluation_model_tasks: t.Dict[str, TaskID] = {}
 235        self.evaluation_model_batches: t.Dict[Snapshot, int] = {}
 236
 237        # Put in temporary values that are replaced when evaluating
 238        self.environment_naming_info = EnvironmentNamingInfo()
 239        self.default_catalog: t.Optional[str] = None
 240
 241        self.creation_progress: t.Optional[Progress] = None
 242        self.creation_task: t.Optional[TaskID] = None
 243
 244        self.promotion_progress: t.Optional[Progress] = None
 245        self.promotion_task: t.Optional[TaskID] = None
 246
 247        self.migration_progress: t.Optional[Progress] = None
 248        self.migration_task: t.Optional[TaskID] = None
 249
 250        self.loading_status: t.Dict[uuid.UUID, Status] = {}
 251
 252        self.verbose = verbose
 253
 254    def _print(self, value: t.Any, **kwargs: t.Any) -> None:
 255        self.console.print(value, **kwargs)
 256
 257    def _prompt(self, message: str, **kwargs: t.Any) -> t.Any:
 258        return Prompt.ask(message, console=self.console, **kwargs)
 259
 260    def _confirm(self, message: str, **kwargs: t.Any) -> bool:
 261        return Confirm.ask(message, console=self.console, **kwargs)
 262
 263    def start_plan_evaluation(self, plan: Plan) -> None:
 264        pass
 265
 266    def stop_plan_evaluation(self) -> None:
 267        pass
 268
 269    def start_evaluation_progress(
 270        self,
 271        batches: t.Dict[Snapshot, int],
 272        environment_naming_info: EnvironmentNamingInfo,
 273        default_catalog: t.Optional[str],
 274    ) -> None:
 275        """Indicates that a new snapshot evaluation progress has begun."""
 276        if not self.evaluation_progress_live:
 277            self.evaluation_total_progress = Progress(
 278                TextColumn("[bold blue]Evaluating models", justify="right"),
 279                BarColumn(bar_width=40),
 280                "[progress.percentage]{task.percentage:>3.1f}%",
 281                "•",
 282                srich.BatchColumn(),
 283                "•",
 284                TimeElapsedColumn(),
 285                console=self.console,
 286            )
 287
 288            self.evaluation_model_progress = Progress(
 289                TextColumn("{task.fields[view_name]}", justify="right"),
 290                SpinnerColumn(spinner_name="simpleDots"),
 291                console=self.console,
 292            )
 293
 294            progress_table = Table.grid()
 295            progress_table.add_row(self.evaluation_total_progress)
 296            progress_table.add_row(self.evaluation_model_progress)
 297
 298            self.evaluation_progress_live = Live(progress_table, refresh_per_second=10)
 299            self.evaluation_progress_live.start()
 300
 301            self.evaluation_total_task = self.evaluation_total_progress.add_task(
 302                "Evaluating models...", total=sum(batches.values())
 303            )
 304
 305            self.evaluation_model_batches = batches
 306            self.environment_naming_info = environment_naming_info
 307            self.default_catalog = default_catalog
 308
 309    def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
 310        if self.evaluation_model_progress and snapshot.name not in self.evaluation_model_tasks:
 311            display_name = snapshot.display_name(self.environment_naming_info, self.default_catalog)
 312            self.evaluation_model_tasks[snapshot.name] = self.evaluation_model_progress.add_task(
 313                f"Evaluating {display_name}...",
 314                view_name=display_name,
 315                total=self.evaluation_model_batches[snapshot],
 316            )
 317
 318    def update_snapshot_evaluation_progress(
 319        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
 320    ) -> None:
 321        """Update the snapshot evaluation progress."""
 322        if (
 323            self.evaluation_total_progress
 324            and self.evaluation_model_progress
 325            and self.evaluation_progress_live
 326        ):
 327            total_batches = self.evaluation_model_batches[snapshot]
 328
 329            if duration_ms:
 330                self.evaluation_progress_live.console.print(
 331                    f"[{batch_idx + 1}/{total_batches}] {snapshot.display_name(self.environment_naming_info, self.default_catalog)} [green]evaluated[/green] in {(duration_ms / 1000.0):.2f}s"
 332                )
 333
 334            self.evaluation_total_progress.update(
 335                self.evaluation_total_task or TaskID(0), refresh=True, advance=1
 336            )
 337
 338            model_task_id = self.evaluation_model_tasks[snapshot.name]
 339            self.evaluation_model_progress.update(model_task_id, refresh=True, advance=1)
 340            if self.evaluation_model_progress._tasks[model_task_id].completed >= total_batches:
 341                self.evaluation_model_progress.remove_task(model_task_id)
 342
 343    def stop_evaluation_progress(self, success: bool = True) -> None:
 344        """Stop the snapshot evaluation progress."""
 345        if self.evaluation_progress_live:
 346            self.evaluation_progress_live.stop()
 347            if success:
 348                self.log_success("All model batches have been executed successfully")
 349
 350        self.evaluation_progress_live = None
 351        self.evaluation_total_progress = None
 352        self.evaluation_total_task = None
 353        self.evaluation_model_progress = None
 354        self.evaluation_model_tasks = {}
 355        self.evaluation_model_batches = {}
 356        self.environment_naming_info = EnvironmentNamingInfo()
 357        self.default_catalog = None
 358
 359    def start_creation_progress(
 360        self,
 361        total_tasks: int,
 362        environment_naming_info: EnvironmentNamingInfo,
 363        default_catalog: t.Optional[str],
 364    ) -> None:
 365        """Indicates that a new creation progress has begun."""
 366        if self.creation_progress is None:
 367            self.creation_progress = Progress(
 368                TextColumn("[bold blue]Creating physical tables", justify="right"),
 369                BarColumn(bar_width=40),
 370                "[progress.percentage]{task.percentage:>3.1f}%",
 371                "•",
 372                srich.BatchColumn(),
 373                "•",
 374                TimeElapsedColumn(),
 375                console=self.console,
 376            )
 377
 378            self.creation_progress.start()
 379            self.creation_task = self.creation_progress.add_task(
 380                "Creating physical tables...",
 381                total=total_tasks,
 382            )
 383
 384            self.environment_naming_info = environment_naming_info
 385            self.default_catalog = default_catalog
 386
 387    def update_creation_progress(self, snapshot: SnapshotInfoLike) -> None:
 388        """Update the snapshot creation progress."""
 389        if self.creation_progress is not None and self.creation_task is not None:
 390            if self.verbose:
 391                self.creation_progress.live.console.print(
 392                    f"{snapshot.display_name(self.environment_naming_info, self.default_catalog)} [green]created[/green]"
 393                )
 394            self.creation_progress.update(self.creation_task, refresh=True, advance=1)
 395
 396    def stop_creation_progress(self, success: bool = True) -> None:
 397        """Stop the snapshot creation progress."""
 398        self.creation_task = None
 399        if self.creation_progress is not None:
 400            self.creation_progress.stop()
 401            self.creation_progress = None
 402            if success:
 403                self.log_success("All model versions have been created successfully")
 404
 405        self.environment_naming_info = EnvironmentNamingInfo()
 406        self.default_catalog = None
 407
 408    def update_cleanup_progress(self, object_name: str) -> None:
 409        """Update the snapshot cleanup progress."""
 410        self._print(f"Deleted object {object_name}")
 411
 412    def start_promotion_progress(
 413        self,
 414        total_tasks: int,
 415        environment_naming_info: EnvironmentNamingInfo,
 416        default_catalog: t.Optional[str],
 417    ) -> None:
 418        """Indicates that a new snapshot promotion progress has begun."""
 419        if self.promotion_progress is None:
 420            self.promotion_progress = Progress(
 421                TextColumn(
 422                    f"[bold blue]Virtually Updating '{environment_naming_info.name}'",
 423                    justify="right",
 424                ),
 425                BarColumn(bar_width=40),
 426                "[progress.percentage]{task.percentage:>3.1f}%",
 427                "•",
 428                TimeElapsedColumn(),
 429                console=self.console,
 430            )
 431
 432            self.promotion_progress.start()
 433            self.promotion_task = self.promotion_progress.add_task(
 434                f"Virtually Updating {environment_naming_info.name}...",
 435                total=total_tasks,
 436            )
 437
 438            self.environment_naming_info = environment_naming_info
 439            self.default_catalog = default_catalog
 440
 441    def update_promotion_progress(self, snapshot: SnapshotInfoLike, promoted: bool) -> None:
 442        """Update the snapshot promotion progress."""
 443        if self.promotion_progress is not None and self.promotion_task is not None:
 444            if self.verbose:
 445                action_str = "[green]promoted[/green]" if promoted else "[yellow]demoted[/yellow]"
 446                self.promotion_progress.live.console.print(
 447                    f"{snapshot.display_name(self.environment_naming_info, self.default_catalog)} {action_str}"
 448                )
 449            self.promotion_progress.update(self.promotion_task, refresh=True, advance=1)
 450
 451    def stop_promotion_progress(self, success: bool = True) -> None:
 452        """Stop the snapshot promotion progress."""
 453        self.promotion_task = None
 454        if self.promotion_progress is not None:
 455            self.promotion_progress.stop()
 456            self.promotion_progress = None
 457            if success:
 458                self.log_success("The target environment has been updated successfully")
 459
 460        self.environment_naming_info = EnvironmentNamingInfo()
 461        self.default_catalog = None
 462
 463    def start_migration_progress(self, total_tasks: int) -> None:
 464        """Indicates that a new migration progress has begun."""
 465        if self.migration_progress is None:
 466            self.migration_progress = Progress(
 467                TextColumn("[bold blue]Migrating snapshots", justify="right"),
 468                BarColumn(bar_width=40),
 469                "[progress.percentage]{task.percentage:>3.1f}%",
 470                "•",
 471                srich.BatchColumn(),
 472                "•",
 473                TimeElapsedColumn(),
 474                console=self.console,
 475            )
 476
 477            self.migration_progress.start()
 478            self.migration_task = self.migration_progress.add_task(
 479                "Migrating snapshots...",
 480                total=total_tasks,
 481            )
 482
 483    def update_migration_progress(self, num_tasks: int) -> None:
 484        """Update the migration progress."""
 485        if self.migration_progress is not None and self.migration_task is not None:
 486            self.migration_progress.update(self.migration_task, refresh=True, advance=num_tasks)
 487
 488    def stop_migration_progress(self, success: bool = True) -> None:
 489        """Stop the migration progress."""
 490        self.migration_task = None
 491        if self.migration_progress is not None:
 492            self.migration_progress.stop()
 493            self.migration_progress = None
 494            if success:
 495                self.log_success("The migration has been completed successfully")
 496
 497    def show_model_difference_summary(
 498        self,
 499        context_diff: ContextDiff,
 500        environment_naming_info: EnvironmentNamingInfo,
 501        default_catalog: t.Optional[str],
 502        no_diff: bool = True,
 503        ignored_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None,
 504    ) -> None:
 505        """Shows a summary of the differences.
 506
 507        Args:
 508            context_diff: The context diff to use to print the summary
 509            environment_naming_info: The environment naming info to reference when printing model names
 510            default_catalog: The default catalog to reference when deciding to remove catalog from display names
 511            no_diff: Hide the actual SQL differences.
 512            ignored_snapshot_ids: A set of snapshot ids that are ignored
 513        """
 514        ignored_snapshot_ids = ignored_snapshot_ids or set()
 515        if context_diff.is_new_environment:
 516            self._print(
 517                Tree(
 518                    f"[bold]New environment `{context_diff.environment}` will be created from `{context_diff.create_from}`"
 519                )
 520            )
 521            if not context_diff.has_snapshot_changes:
 522                return
 523
 524        if not context_diff.has_changes:
 525            self._print(Tree(f"[bold]No differences when compared to `{context_diff.environment}`"))
 526            return
 527
 528        self._print(Tree(f"[bold]Summary of differences against `{context_diff.environment}`:"))
 529        self._show_summary_tree_for(
 530            context_diff,
 531            "Models",
 532            lambda x: x.is_model,
 533            environment_naming_info,
 534            default_catalog,
 535            no_diff=no_diff,
 536            ignored_snapshot_ids=ignored_snapshot_ids,
 537        )
 538        self._show_summary_tree_for(
 539            context_diff,
 540            "Standalone Audits",
 541            lambda x: x.is_audit,
 542            environment_naming_info,
 543            default_catalog,
 544            no_diff=no_diff,
 545            ignored_snapshot_ids=ignored_snapshot_ids,
 546        )
 547
 548    def plan(
 549        self,
 550        plan_builder: PlanBuilder,
 551        auto_apply: bool,
 552        default_catalog: t.Optional[str],
 553        no_diff: bool = False,
 554        no_prompts: bool = False,
 555    ) -> None:
 556        """The main plan flow.
 557
 558        The console should present the user with choices on how to backfill and version the snapshots
 559        of a plan.
 560
 561        Args:
 562            plan: The plan to make choices for.
 563            auto_apply: Whether to automatically apply the plan after all choices have been made.
 564            default_catalog: The default catalog to reference when deciding to remove catalog from display names
 565            no_diff: Hide text differences for changed models.
 566            no_prompts: Whether to disable interactive prompts for the backfill time range. Please note that
 567                if this flag is set to true and there are uncategorized changes the plan creation will
 568                fail. Default: False
 569        """
 570        self._prompt_categorize(
 571            plan_builder,
 572            auto_apply,
 573            no_diff=no_diff,
 574            no_prompts=no_prompts,
 575            default_catalog=default_catalog,
 576        )
 577
 578        if not no_prompts:
 579            self._show_options_after_categorization(
 580                plan_builder, auto_apply, default_catalog=default_catalog
 581            )
 582
 583        if auto_apply:
 584            plan_builder.apply()
 585
 586    def _get_ignored_tree(
 587        self,
 588        ignored_snapshot_ids: t.Set[SnapshotId],
 589        snapshots: t.Dict[SnapshotId, Snapshot],
 590        environment_naming_info: EnvironmentNamingInfo,
 591        default_catalog: t.Optional[str],
 592    ) -> Tree:
 593        ignored = Tree("[bold][ignored]Ignored Models (Expected Plan Start):")
 594        for s_id in ignored_snapshot_ids:
 595            snapshot = snapshots[s_id]
 596            ignored.add(
 597                f"[ignored]{snapshot.display_name(environment_naming_info, default_catalog)} ({snapshot.get_latest(start_date(snapshot, snapshots.values()))})"
 598            )
 599        return ignored
 600
 601    def _show_summary_tree_for(
 602        self,
 603        context_diff: ContextDiff,
 604        header: str,
 605        snapshot_selector: t.Callable[[SnapshotInfoLike], bool],
 606        environment_naming_info: EnvironmentNamingInfo,
 607        default_catalog: t.Optional[str],
 608        no_diff: bool = True,
 609        ignored_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None,
 610    ) -> None:
 611        ignored_snapshot_ids = ignored_snapshot_ids or set()
 612        selected_snapshots = {
 613            s_id: snapshot
 614            for s_id, snapshot in context_diff.snapshots.items()
 615            if snapshot_selector(snapshot)
 616        }
 617        selected_ignored_snapshot_ids = {
 618            s_id for s_id in selected_snapshots if s_id in ignored_snapshot_ids
 619        }
 620        added_snapshot_ids = {
 621            s_id for s_id in context_diff.added if snapshot_selector(context_diff.snapshots[s_id])
 622        } - selected_ignored_snapshot_ids
 623        removed_snapshot_ids = {
 624            s_id
 625            for s_id, snapshot in context_diff.removed_snapshots.items()
 626            if snapshot_selector(snapshot)
 627        } - selected_ignored_snapshot_ids
 628        modified_snapshot_ids = {
 629            current_snapshot.snapshot_id
 630            for _, (current_snapshot, _) in context_diff.modified_snapshots.items()
 631            if snapshot_selector(current_snapshot)
 632        } - selected_ignored_snapshot_ids
 633
 634        tree_sets = (
 635            added_snapshot_ids,
 636            removed_snapshot_ids,
 637            modified_snapshot_ids,
 638            selected_ignored_snapshot_ids,
 639        )
 640        if all(not s_ids for s_ids in tree_sets):
 641            return
 642
 643        tree = Tree(f"[bold]{header}:")
 644        if added_snapshot_ids:
 645            added_tree = Tree("[bold][added]Added:")
 646            for s_id in added_snapshot_ids:
 647                snapshot = context_diff.snapshots[s_id]
 648                added_tree.add(
 649                    f"[added]{snapshot.display_name(environment_naming_info, default_catalog)}"
 650                )
 651            tree.add(added_tree)
 652        if removed_snapshot_ids:
 653            removed_tree = Tree("[bold][removed]Removed:")
 654            for s_id in removed_snapshot_ids:
 655                snapshot_table_info = context_diff.removed_snapshots[s_id]
 656                removed_tree.add(
 657                    f"[removed]{snapshot_table_info.display_name(environment_naming_info, default_catalog)}"
 658                )
 659            tree.add(removed_tree)
 660        if modified_snapshot_ids:
 661            direct = Tree("[bold][direct]Directly Modified:")
 662            indirect = Tree("[bold][indirect]Indirectly Modified:")
 663            metadata = Tree("[bold][metadata]Metadata Updated:")
 664            for s_id in modified_snapshot_ids:
 665                name = s_id.name
 666                display_name = context_diff.snapshots[s_id].display_name(
 667                    environment_naming_info, default_catalog
 668                )
 669                if context_diff.directly_modified(name):
 670                    direct.add(
 671                        f"[direct]{display_name}"
 672                        if no_diff
 673                        else Syntax(f"{display_name}\n{context_diff.text_diff(name)}", "sql")
 674                    )
 675                elif context_diff.indirectly_modified(name):
 676                    indirect.add(f"[indirect]{display_name}")
 677                elif context_diff.metadata_updated(name):
 678                    metadata.add(
 679                        f"[metadata]{display_name}"
 680                        if no_diff
 681                        else Syntax(f"{display_name}", "sql", word_wrap=True)
 682                    )
 683            if direct.children:
 684                tree.add(direct)
 685            if indirect.children:
 686                tree.add(indirect)
 687            if metadata.children:
 688                tree.add(metadata)
 689        if selected_ignored_snapshot_ids:
 690            tree.add(
 691                self._get_ignored_tree(
 692                    selected_ignored_snapshot_ids,
 693                    selected_snapshots,
 694                    environment_naming_info,
 695                    default_catalog,
 696                )
 697            )
 698        self._print(tree)
 699
 700    def _show_options_after_categorization(
 701        self, plan_builder: PlanBuilder, auto_apply: bool, default_catalog: t.Optional[str]
 702    ) -> None:
 703        plan = plan_builder.build()
 704        if plan.forward_only and plan.new_snapshots:
 705            self._prompt_effective_from(plan_builder, auto_apply, default_catalog)
 706
 707        if plan.requires_backfill:
 708            self._show_missing_dates(plan_builder.build(), default_catalog)
 709            self._prompt_backfill(plan_builder, auto_apply, default_catalog)
 710        elif plan.has_changes and not auto_apply:
 711            self._prompt_promote(plan_builder)
 712        elif plan.has_unmodified_unpromoted and not auto_apply:
 713            self.log_status_update("\n[bold]Virtually updating unmodified models\n")
 714            self._prompt_promote(plan_builder)
 715
 716    def _prompt_categorize(
 717        self,
 718        plan_builder: PlanBuilder,
 719        auto_apply: bool,
 720        no_diff: bool,
 721        no_prompts: bool,
 722        default_catalog: t.Optional[str],
 723    ) -> None:
 724        """Get the user's change category for the directly modified models."""
 725        plan = plan_builder.build()
 726
 727        self.show_model_difference_summary(
 728            plan.context_diff,
 729            plan.environment_naming_info,
 730            default_catalog=default_catalog,
 731            ignored_snapshot_ids=plan.ignored,
 732        )
 733
 734        if not no_diff:
 735            self._show_categorized_snapshots(plan, default_catalog)
 736
 737        for snapshot in plan.uncategorized:
 738            if not no_diff:
 739                self.show_sql(plan.context_diff.text_diff(snapshot.name))
 740            tree = Tree(
 741                f"[bold][direct]Directly Modified: {snapshot.display_name(plan.environment_naming_info, default_catalog)}"
 742            )
 743            indirect_tree = None
 744
 745            for child_sid in sorted(plan.indirectly_modified.get(snapshot.snapshot_id, set())):
 746                child_snapshot = plan.context_diff.snapshots[child_sid]
 747                if not indirect_tree:
 748                    indirect_tree = Tree("[indirect]Indirectly Modified Children:")
 749                    tree.add(indirect_tree)
 750                indirect_tree.add(
 751                    f"[indirect]{child_snapshot.display_name(plan.environment_naming_info, default_catalog)}"
 752                )
 753            self._print(tree)
 754            if not no_prompts:
 755                self._get_snapshot_change_category(
 756                    snapshot, plan_builder, auto_apply, default_catalog
 757                )
 758
 759    def _show_categorized_snapshots(self, plan: Plan, default_catalog: t.Optional[str]) -> None:
 760        context_diff = plan.context_diff
 761
 762        for snapshot in plan.categorized:
 763            if not context_diff.directly_modified(snapshot.name):
 764                continue
 765
 766            category_str = SNAPSHOT_CHANGE_CATEGORY_STR[snapshot.change_category]
 767            tree = Tree(
 768                f"[bold][direct]Directly Modified: {snapshot.display_name(plan.environment_naming_info, default_catalog)} ({category_str})"
 769            )
 770            indirect_tree = None
 771            for child_sid in sorted(plan.indirectly_modified.get(snapshot.snapshot_id, set())):
 772                child_snapshot = context_diff.snapshots[child_sid]
 773                if not indirect_tree:
 774                    indirect_tree = Tree("[indirect]Indirectly Modified Children:")
 775                    tree.add(indirect_tree)
 776                child_category_str = SNAPSHOT_CHANGE_CATEGORY_STR[child_snapshot.change_category]
 777                indirect_tree.add(
 778                    f"[indirect]{child_snapshot.display_name(plan.environment_naming_info, default_catalog)} ({child_category_str})"
 779                )
 780            self._print(Syntax(context_diff.text_diff(snapshot.name), "sql", word_wrap=True))
 781            self._print(tree)
 782
 783    def _show_missing_dates(self, plan: Plan, default_catalog: t.Optional[str]) -> None:
 784        """Displays the models with missing dates."""
 785        missing_intervals = plan.missing_intervals
 786        if not missing_intervals:
 787            return
 788        backfill = Tree("[bold]Models needing backfill (missing dates):")
 789        for missing in missing_intervals:
 790            snapshot = plan.context_diff.snapshots[missing.snapshot_id]
 791            if not snapshot.is_model:
 792                continue
 793
 794            preview_modifier = ""
 795            if not plan.deployability_index.is_deployable(snapshot):
 796                preview_modifier = " ([orange1]preview[/orange1])"
 797
 798            backfill.add(
 799                f"{snapshot.display_name(plan.environment_naming_info, default_catalog)}: {missing.format_intervals(snapshot.node.interval_unit)}{preview_modifier}"
 800            )
 801        self._print(backfill)
 802
 803    def _prompt_effective_from(
 804        self, plan_builder: PlanBuilder, auto_apply: bool, default_catalog: t.Optional[str]
 805    ) -> None:
 806        if not plan_builder.build().effective_from:
 807            effective_from = self._prompt(
 808                "Enter the effective date (eg. '1 year', '2020-01-01') to apply forward-only changes retroactively or blank to only apply them going forward once changes are deployed to prod"
 809            )
 810            if effective_from:
 811                plan_builder.set_effective_from(effective_from)
 812
 813    def _prompt_backfill(
 814        self, plan_builder: PlanBuilder, auto_apply: bool, default_catalog: t.Optional[str]
 815    ) -> None:
 816        plan = plan_builder.build()
 817        is_forward_only_dev = plan.is_dev and plan.forward_only
 818        backfill_or_preview = "preview" if is_forward_only_dev else "backfill"
 819
 820        if plan_builder.is_start_and_end_allowed:
 821            if not plan_builder.override_start:
 822                if is_forward_only_dev:
 823                    if plan.effective_from:
 824                        blank_meaning = f"to preview starting from the effective date ('{time_like_to_str(plan.effective_from)}')"
 825                        default_start = plan.effective_from
 826                    else:
 827                        blank_meaning = "to preview starting from yesterday"
 828                        default_start = yesterday_ds()
 829                else:
 830                    if plan.provided_start:
 831                        blank_meaning = f"starting from '{time_like_to_str(plan.provided_start)}'"
 832                    else:
 833                        blank_meaning = "from the beginning of history"
 834                    default_start = None
 835
 836                start = self._prompt(
 837                    f"Enter the {backfill_or_preview} start date (eg. '1 year', '2020-01-01') or blank to backfill {blank_meaning}",
 838                )
 839                if start:
 840                    plan_builder.set_start(start)
 841                elif default_start:
 842                    plan_builder.set_start(default_start)
 843
 844            if not plan_builder.override_end:
 845                if plan.provided_end:
 846                    blank_meaning = f"'{time_like_to_str(plan.provided_end)}'"
 847                else:
 848                    blank_meaning = "now"
 849                end = self._prompt(
 850                    f"Enter the {backfill_or_preview} end date (eg. '1 month ago', '2020-01-01') or blank to {backfill_or_preview} up until {blank_meaning}",
 851                )
 852                if end:
 853                    plan_builder.set_end(end)
 854
 855            plan = plan_builder.build()
 856
 857        if plan.ignored:
 858            self._print(
 859                self._get_ignored_tree(
 860                    plan.ignored,
 861                    plan.context_diff.snapshots,
 862                    plan.environment_naming_info,
 863                    default_catalog,
 864                )
 865            )
 866        if not auto_apply and self._confirm(f"Apply - {backfill_or_preview.capitalize()} Tables"):
 867            plan_builder.apply()
 868
 869    def _prompt_promote(self, plan_builder: PlanBuilder) -> None:
 870        if self._confirm(
 871            "Apply - Virtual Update",
 872        ):
 873            plan_builder.apply()
 874
 875    def log_test_results(
 876        self, result: unittest.result.TestResult, output: str, target_dialect: str
 877    ) -> None:
 878        divider_length = 70
 879        if result.wasSuccessful():
 880            self._print("=" * divider_length)
 881            self._print(
 882                f"Successfully Ran {str(result.testsRun)} tests against {target_dialect}",
 883                style="green",
 884            )
 885            self._print("-" * divider_length)
 886        else:
 887            self._print("-" * divider_length)
 888            self._print("Test Failure Summary")
 889            self._print("=" * divider_length)
 890            self._print(
 891                f"Num Successful Tests: {result.testsRun - len(result.failures) - len(result.errors)}"
 892            )
 893            for test, _ in result.failures + result.errors:
 894                if isinstance(test, ModelTest):
 895                    self._print(f"Failure Test: {test.model.name} {test.test_name}")
 896            self._print("=" * divider_length)
 897            self._print(output)
 898
 899    def show_sql(self, sql: str) -> None:
 900        self._print(Syntax(sql, "sql", word_wrap=True), crop=False)
 901
 902    def log_status_update(self, message: str) -> None:
 903        self._print(message)
 904
 905    def log_error(self, message: str) -> None:
 906        self._print(f"[red]{message}[/red]")
 907
 908    def log_success(self, message: str) -> None:
 909        self._print(f"\n[green]{message}[/green]\n")
 910
 911    def loading_start(self, message: t.Optional[str] = None) -> uuid.UUID:
 912        id = uuid.uuid4()
 913        self.loading_status[id] = Status(message or "", console=self.console, spinner="line")
 914        self.loading_status[id].start()
 915        return id
 916
 917    def loading_stop(self, id: uuid.UUID) -> None:
 918        self.loading_status[id].stop()
 919        del self.loading_status[id]
 920
 921    def show_schema_diff(self, schema_diff: SchemaDiff) -> None:
 922        source_name = schema_diff.source
 923        if schema_diff.source_alias:
 924            source_name = schema_diff.source_alias.upper()
 925        target_name = schema_diff.target
 926        if schema_diff.target_alias:
 927            target_name = schema_diff.target_alias.upper()
 928
 929        first_line = f"\n[b]Schema Diff Between '[yellow]{source_name}[/yellow]' and '[green]{target_name}[/green]'"
 930        if schema_diff.model_name:
 931            first_line = (
 932                first_line + f" environments for model '[blue]{schema_diff.model_name}[/blue]'"
 933            )
 934
 935        tree = Tree(first_line + ":")
 936
 937        if any([schema_diff.added, schema_diff.removed, schema_diff.modified]):
 938            if schema_diff.added:
 939                added = Tree("[green]Added Columns:")
 940                for c, t in schema_diff.added:
 941                    added.add(f"[green]{c} ({t})")
 942                tree.add(added)
 943
 944            if schema_diff.removed:
 945                removed = Tree("[red]Removed Columns:")
 946                for c, t in schema_diff.removed:
 947                    removed.add(f"[red]{c} ({t})")
 948                tree.add(removed)
 949
 950            if schema_diff.modified:
 951                modified = Tree("[magenta]Modified Columns:")
 952                for c, (ft, tt) in schema_diff.modified.items():
 953                    modified.add(f"[magenta]{c} ({ft} -> {tt})")
 954                tree.add(modified)
 955        else:
 956            tree.add("[b]Schemas match")
 957
 958        self.console.print(tree)
 959
 960    def show_row_diff(self, row_diff: RowDiff, show_sample: bool = True) -> None:
 961        source_name = row_diff.source
 962        if row_diff.source_alias:
 963            source_name = row_diff.source_alias.upper()
 964        target_name = row_diff.target
 965        if row_diff.target_alias:
 966            target_name = row_diff.target_alias.upper()
 967
 968        tree = Tree("[b]Row Counts:[/b]")
 969        tree.add(f" [b][blue]COMMON[/blue]:[/b] {row_diff.join_count} rows")
 970        tree.add(f" [b][yellow]{source_name} ONLY[/yellow]:[/b] {row_diff.s_only_count} rows")
 971        tree.add(f" [b][green]{target_name} ONLY[/green]:[/b] {row_diff.t_only_count} rows")
 972        self.console.print("\n", tree)
 973
 974        self.console.print("\n[b][blue]COMMON ROWS[/blue] column comparison stats:[/b]")
 975        if row_diff.column_stats.shape[0] > 0:
 976            self.console.print(row_diff.column_stats.to_string(index=True), end="\n\n")
 977        else:
 978            self.console.print("  No columns with same name and data type in both tables")
 979
 980        if show_sample:
 981            self.console.print("\n[b][blue]COMMON ROWS[/blue] sample data differences:[/b]")
 982            if row_diff.joined_sample.shape[0] > 0:
 983                self.console.print(row_diff.joined_sample.to_string(index=False), end="\n\n")
 984            else:
 985                self.console.print("  All joined rows match")
 986
 987            if row_diff.s_sample.shape[0] > 0:
 988                self.console.print(f"\n[b][yellow]{source_name} ONLY[/yellow] sample rows:[/b]")
 989                self.console.print(row_diff.s_sample.to_string(index=False), end="\n\n")
 990
 991            if row_diff.t_sample.shape[0] > 0:
 992                self.console.print(f"\n[b][green]{target_name} ONLY[/green] sample rows:[/b]")
 993                self.console.print(row_diff.t_sample.to_string(index=False), end="\n\n")
 994
 995    def _get_snapshot_change_category(
 996        self,
 997        snapshot: Snapshot,
 998        plan_builder: PlanBuilder,
 999        auto_apply: bool,
1000        default_catalog: t.Optional[str],
1001    ) -> None:
1002        choices = self._snapshot_change_choices(
1003            snapshot, plan_builder.environment_naming_info, default_catalog
1004        )
1005        response = self._prompt(
1006            "\n".join([f"[{i+1}] {choice}" for i, choice in enumerate(choices.values())]),
1007            show_choices=False,
1008            choices=[f"{i+1}" for i in range(len(choices))],
1009        )
1010        choice = list(choices)[int(response) - 1]
1011        plan_builder.set_choice(snapshot, choice)
1012
1013    def _snapshot_change_choices(
1014        self,
1015        snapshot: Snapshot,
1016        environment_naming_info: EnvironmentNamingInfo,
1017        default_catalog: t.Optional[str],
1018        use_rich_formatting: bool = True,
1019    ) -> t.Dict[SnapshotChangeCategory, str]:
1020        direct = snapshot.display_name(environment_naming_info, default_catalog)
1021        if use_rich_formatting:
1022            direct = f"[direct]{direct}[/direct]"
1023        indirect = "indirectly modified children"
1024        if use_rich_formatting:
1025            indirect = f"[indirect]{indirect}[/indirect]"
1026        if snapshot.is_view:
1027            choices = {
1028                SnapshotChangeCategory.BREAKING: f"Update {direct} and backfill {indirect}",
1029                SnapshotChangeCategory.NON_BREAKING: f"Update {direct} but don't backfill {indirect}",
1030            }
1031        elif snapshot.is_symbolic:
1032            choices = {
1033                SnapshotChangeCategory.BREAKING: f"Backfill {indirect}",
1034                SnapshotChangeCategory.NON_BREAKING: f"Don't backfill {indirect}",
1035            }
1036        else:
1037            choices = {
1038                SnapshotChangeCategory.BREAKING: f"Backfill {direct} and {indirect}",
1039                SnapshotChangeCategory.NON_BREAKING: f"Backfill {direct} but not {indirect}",
1040            }
1041        labeled_choices = {
1042            k: f"[{SNAPSHOT_CHANGE_CATEGORY_STR[k]}] {v}" for k, v in choices.items()
1043        }
1044        return labeled_choices
1045
1046
1047def add_to_layout_widget(target_widget: LayoutWidget, *widgets: widgets.Widget) -> LayoutWidget:
1048    """Helper function to add a widget to a layout widget.
1049
1050    Args:
1051        target_widget: The layout widget to add the other widget(s) to.
1052        *widgets: The widgets to add to the layout widget.
1053
1054    Returns:
1055        The layout widget with the children added.
1056    """
1057    target_widget.children += tuple(widgets)
1058    return target_widget
1059
1060
1061class NotebookMagicConsole(TerminalConsole):
1062    """
1063    Console to be used when using the magic notebook interface (`%<command>`).
1064    Generally reuses the Terminal console when possible by either directly outputing what it provides
1065    or capturing it and converting it into a widget.
1066    """
1067
1068    def __init__(
1069        self,
1070        display: t.Optional[t.Callable] = None,
1071        console: t.Optional[RichConsole] = None,
1072        **kwargs: t.Any,
1073    ) -> None:
1074        import ipywidgets as widgets
1075        from IPython import get_ipython
1076        from IPython.display import display as ipython_display
1077
1078        super().__init__(console, **kwargs)
1079
1080        self.display = display or get_ipython().user_ns.get("display", ipython_display)
1081        self.missing_dates_output = widgets.Output()
1082        self.dynamic_options_after_categorization_output = widgets.VBox()
1083
1084    def _show_missing_dates(self, plan: Plan, default_catalog: t.Optional[str]) -> None:
1085        self._add_to_dynamic_options(self.missing_dates_output)
1086        self.missing_dates_output.outputs = ()
1087        with self.missing_dates_output:
1088            super()._show_missing_dates(plan, default_catalog)
1089
1090    def _apply(self, button: widgets.Button) -> None:
1091        button.disabled = True
1092        with button.output:
1093            button.plan_builder.apply()
1094
1095    def _prompt_promote(self, plan_builder: PlanBuilder) -> None:
1096        import ipywidgets as widgets
1097
1098        button = widgets.Button(
1099            description="Apply - Virtual Update",
1100            disabled=False,
1101            button_style="success",
1102            # Auto will make the button really large.
1103            # Likely changing this soon anyways to be just `Apply` with description above
1104            layout={"width": "10rem"},
1105        )
1106        self._add_to_dynamic_options(button)
1107        output = widgets.Output()
1108        self._add_to_dynamic_options(output)
1109
1110        button.plan_builder = plan_builder
1111        button.on_click(self._apply)
1112        button.output = output
1113
1114    def _prompt_effective_from(
1115        self, plan_builder: PlanBuilder, auto_apply: bool, default_catalog: t.Optional[str]
1116    ) -> None:
1117        import ipywidgets as widgets
1118
1119        prompt = widgets.VBox()
1120
1121        def effective_from_change_callback(change: t.Dict[str, datetime.datetime]) -> None:
1122            plan_builder.set_effective_from(change["new"])
1123            self._show_options_after_categorization(plan_builder, auto_apply, default_catalog)
1124
1125        def going_forward_change_callback(change: t.Dict[str, bool]) -> None:
1126            checked = change["new"]
1127            plan_builder.set_effective_from(None if checked else yesterday_ds())
1128            self._show_options_after_categorization(
1129                plan_builder, auto_apply=auto_apply, default_catalog=default_catalog
1130            )
1131
1132        date_picker = widgets.DatePicker(
1133            disabled=plan_builder.build().effective_from is None,
1134            value=to_date(plan_builder.build().effective_from or yesterday_ds()),
1135            layout={"width": "auto"},
1136        )
1137        date_picker.observe(effective_from_change_callback, "value")
1138
1139        going_forward_checkbox = widgets.Checkbox(
1140            value=plan_builder.build().effective_from is None,
1141            description="Apply Going Forward Once Deployed To Prod",
1142            disabled=False,
1143            indent=False,
1144        )
1145        going_forward_checkbox.observe(going_forward_change_callback, "value")
1146
1147        add_to_layout_widget(
1148            prompt,
1149            widgets.HBox(
1150                [
1151                    widgets.Label("Effective From Date:", layout={"width": "8rem"}),
1152                    date_picker,
1153                    going_forward_checkbox,
1154                ]
1155            ),
1156        )
1157
1158        self._add_to_dynamic_options(prompt)
1159
1160    def _prompt_backfill(
1161        self, plan_builder: PlanBuilder, auto_apply: bool, default_catalog: t.Optional[str]
1162    ) -> None:
1163        import ipywidgets as widgets
1164
1165        prompt = widgets.VBox()
1166
1167        backfill_or_preview = (
1168            "Preview"
1169            if plan_builder.build().is_dev and plan_builder.build().forward_only
1170            else "Backfill"
1171        )
1172
1173        def _date_picker(
1174            plan_builder: PlanBuilder, value: t.Any, on_change: t.Callable, disabled: bool = False
1175        ) -> widgets.DatePicker:
1176            picker = widgets.DatePicker(
1177                disabled=disabled,
1178                value=value,
1179                layout={"width": "auto"},
1180            )
1181
1182            picker.observe(on_change, "value")
1183            return picker
1184
1185        def start_change_callback(change: t.Dict[str, datetime.datetime]) -> None:
1186            plan_builder.set_start(change["new"])
1187            self._show_options_after_categorization(plan_builder, auto_apply, default_catalog)
1188
1189        def end_change_callback(change: t.Dict[str, datetime.datetime]) -> None:
1190            plan_builder.set_end(change["new"])
1191            self._show_options_after_categorization(plan_builder, auto_apply, default_catalog)
1192
1193        if plan_builder.is_start_and_end_allowed:
1194            add_to_layout_widget(
1195                prompt,
1196                widgets.HBox(
1197                    [
1198                        widgets.Label(
1199                            f"Start {backfill_or_preview} Date:", layout={"width": "8rem"}
1200                        ),
1201                        _date_picker(
1202                            plan_builder, to_date(plan_builder.build().start), start_change_callback
1203                        ),
1204                    ]
1205                ),
1206            )
1207
1208            add_to_layout_widget(
1209                prompt,
1210                widgets.HBox(
1211                    [
1212                        widgets.Label(f"End {backfill_or_preview} Date:", layout={"width": "8rem"}),
1213                        _date_picker(
1214                            plan_builder,
1215                            to_date(plan_builder.build().end),
1216                            end_change_callback,
1217                        ),
1218                    ]
1219                ),
1220            )
1221
1222        self._add_to_dynamic_options(prompt)
1223
1224        if not auto_apply:
1225            button = widgets.Button(
1226                description=f"Apply - {backfill_or_preview} Tables",
1227                disabled=False,
1228                button_style="success",
1229            )
1230            self._add_to_dynamic_options(button)
1231            output = widgets.Output()
1232            self._add_to_dynamic_options(output)
1233
1234            button.plan_builder = plan_builder
1235            button.on_click(self._apply)
1236            button.output = output
1237
1238    def _show_options_after_categorization(
1239        self, plan_builder: PlanBuilder, auto_apply: bool, default_catalog: t.Optional[str]
1240    ) -> None:
1241        self.dynamic_options_after_categorization_output.children = ()
1242        self.display(self.dynamic_options_after_categorization_output)
1243        super()._show_options_after_categorization(plan_builder, auto_apply, default_catalog)
1244
1245    def _add_to_dynamic_options(self, *widgets: widgets.Widget) -> None:
1246        add_to_layout_widget(self.dynamic_options_after_categorization_output, *widgets)
1247
1248    def _get_snapshot_change_category(
1249        self,
1250        snapshot: Snapshot,
1251        plan_builder: PlanBuilder,
1252        auto_apply: bool,
1253        default_catalog: t.Optional[str],
1254    ) -> None:
1255        import ipywidgets as widgets
1256
1257        choice_mapping = self._snapshot_change_choices(
1258            snapshot,
1259            plan_builder.environment_naming_info,
1260            default_catalog,
1261            use_rich_formatting=False,
1262        )
1263        choices = list(choice_mapping)
1264        plan_builder.set_choice(snapshot, choices[0])
1265
1266        def radio_button_selected(change: t.Dict[str, t.Any]) -> None:
1267            plan_builder.set_choice(snapshot, choices[change["owner"].index])
1268            self._show_options_after_categorization(plan_builder, auto_apply, default_catalog)
1269
1270        radio = widgets.RadioButtons(
1271            options=choice_mapping.values(),
1272            layout={"width": "max-content"},
1273            disabled=False,
1274        )
1275        radio.observe(
1276            radio_button_selected,
1277            "value",
1278        )
1279        self.display(radio)
1280
1281    def log_test_results(
1282        self, result: unittest.result.TestResult, output: str, target_dialect: str
1283    ) -> None:
1284        import ipywidgets as widgets
1285
1286        divider_length = 70
1287        shared_style = {
1288            "font-size": "11px",
1289            "font-weight": "bold",
1290            "font-family": "Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace",
1291        }
1292        if result.wasSuccessful():
1293            success_color = {"color": "#008000"}
1294            header = str(h("span", {"style": shared_style}, "-" * divider_length))
1295            message = str(
1296                h(
1297                    "span",
1298                    {"style": {**shared_style, **success_color}},
1299                    f"Successfully Ran {str(result.testsRun)} Tests Against {target_dialect}",
1300                )
1301            )
1302            footer = str(h("span", {"style": shared_style}, "=" * divider_length))
1303            self.display(widgets.HTML("<br>".join([header, message, footer])))
1304        else:
1305            fail_color = {"color": "#db3737"}
1306            fail_shared_style = {**shared_style, **fail_color}
1307            header = str(h("span", {"style": fail_shared_style}, "-" * divider_length))
1308            message = str(h("span", {"style": fail_shared_style}, "Test Failure Summary"))
1309            num_success = str(
1310                h(
1311                    "span",
1312                    {"style": fail_shared_style},
1313                    f"Num Successful Tests: {result.testsRun - len(result.failures) - len(result.errors)}",
1314                )
1315            )
1316            failure_tests = []
1317            for test, _ in result.failures + result.errors:
1318                if isinstance(test, ModelTest):
1319                    failure_tests.append(
1320                        str(
1321                            h(
1322                                "span",
1323                                {"style": fail_shared_style},
1324                                f"Failure Test: {test.model.name} {test.test_name}",
1325                            )
1326                        )
1327                    )
1328            failures = "<br>".join(failure_tests)
1329            footer = str(h("span", {"style": fail_shared_style}, "=" * divider_length))
1330            error_output = widgets.Textarea(output, layout={"height": "300px", "width": "100%"})
1331            test_info = widgets.HTML(
1332                "<br>".join([header, message, footer, num_success, failures, footer])
1333            )
1334            self.display(widgets.VBox(children=[test_info, error_output], layout={"width": "100%"}))
1335
1336
1337class CaptureTerminalConsole(TerminalConsole):
1338    """
1339    Captures the output of the terminal console so that it can be extracted out and displayed within other interfaces.
1340    The captured output is cleared out after it is retrieved.
1341
1342    Note: `_prompt` and `_confirm` need to also be overriden to work with the custom interface if you want to use
1343    this console interactively.
1344    """
1345
1346    def __init__(self, console: t.Optional[RichConsole] = None, **kwargs: t.Any) -> None:
1347        super().__init__(console=console, **kwargs)
1348        self._captured_outputs: t.List[str] = []
1349        self._errors: t.List[str] = []
1350
1351    @property
1352    def captured_output(self) -> str:
1353        return "".join(self._captured_outputs)
1354
1355    @property
1356    def captured_errors(self) -> str:
1357        return "".join(self._errors)
1358
1359    def consume_captured_output(self) -> str:
1360        output = self.captured_output
1361        self.clear_captured_outputs()
1362        return output
1363
1364    def consume_captured_errors(self) -> str:
1365        errors = self.captured_errors
1366        self.clear_captured_errors()
1367        return errors
1368
1369    def clear_captured_outputs(self) -> None:
1370        self._captured_outputs = []
1371
1372    def clear_captured_errors(self) -> None:
1373        self._errors = []
1374
1375    def log_error(self, message: str) -> None:
1376        self._errors.append(message)
1377        super().log_error(message)
1378
1379    def _print(self, value: t.Any, **kwargs: t.Any) -> None:
1380        with self.console.capture() as capture:
1381            self.console.print(value, **kwargs)
1382        self._captured_outputs.append(capture.get())
1383
1384
1385class MarkdownConsole(CaptureTerminalConsole):
1386    """
1387    A console that outputs markdown. Currently this is only configured for non-interactive use so for use cases
1388    where you want to display a plan or test results in markdown.
1389    """
1390
1391    def show_model_difference_summary(
1392        self,
1393        context_diff: ContextDiff,
1394        environment_naming_info: EnvironmentNamingInfo,
1395        default_catalog: t.Optional[str],
1396        no_diff: bool = True,
1397        ignored_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None,
1398    ) -> None:
1399        """Shows a summary of the differences.
1400
1401        Args:
1402            context_diff: The context diff to use to print the summary.
1403            environment_naming_info: The environment naming info to reference when printing model names
1404            default_catalog: The default catalog to reference when deciding to remove catalog from display names
1405            no_diff: Hide the actual SQL differences.
1406            ignored_snapshot_ids: A set of snapshot names that are ignored
1407        """
1408        ignored_snapshot_ids = ignored_snapshot_ids or set()
1409        if context_diff.is_new_environment:
1410            self._print(
1411                f"**New environment `{context_diff.environment}` will be created from `{context_diff.create_from}`**\n"
1412            )
1413            if not context_diff.has_snapshot_changes:
1414                return
1415
1416        if not context_diff.has_changes:
1417            self._print(f"**No differences when compared to `{context_diff.environment}`**\n")
1418            return
1419
1420        self._print(f"**Summary of differences against `{context_diff.environment}`:**\n")
1421
1422        added_snapshots = {
1423            context_diff.snapshots[s_id]
1424            for s_id in context_diff.added
1425            if s_id not in ignored_snapshot_ids
1426        }
1427        added_snapshot_models = {s for s in added_snapshots if s.is_model}
1428        if added_snapshot_models:
1429            self._print("\n**Added Models:**")
1430            for snapshot in sorted(added_snapshot_models):
1431                self._print(
1432                    f"- `{snapshot.display_name(environment_naming_info, default_catalog)}`"
1433                )
1434
1435        added_snapshot_audits = {s for s in added_snapshots if s.is_audit}
1436        if added_snapshot_audits:
1437            self._print("\n**Added Standalone Audits:**")
1438            for snapshot in sorted(added_snapshot_audits):
1439                self._print(
1440                    f"- `{snapshot.display_name(environment_naming_info, default_catalog)}`"
1441                )
1442
1443        removed_snapshot_table_infos = {
1444            snapshot_table_info
1445            for s_id, snapshot_table_info in context_diff.removed_snapshots.items()
1446            if s_id not in ignored_snapshot_ids
1447        }
1448        removed_model_snapshot_table_infos = {s for s in removed_snapshot_table_infos if s.is_model}
1449        if removed_model_snapshot_table_infos:
1450            self._print("\n**Removed Models:**")
1451            for snapshot_table_info in sorted(removed_model_snapshot_table_infos):
1452                self._print(
1453                    f"- `{snapshot_table_info.display_name(environment_naming_info, default_catalog)}`"
1454                )
1455
1456        removed_audit_snapshot_table_infos = {s for s in removed_snapshot_table_infos if s.is_audit}
1457        if removed_audit_snapshot_table_infos:
1458            self._print("\n**Removed Standalone Audits:**")
1459            for snapshot_table_info in sorted(removed_audit_snapshot_table_infos):
1460                self._print(
1461                    f"- `{snapshot_table_info.display_name(environment_naming_info, default_catalog)}`"
1462                )
1463
1464        modified_snapshots = {
1465            current_snapshot
1466            for current_snapshot, _ in context_diff.modified_snapshots.values()
1467            if current_snapshot.snapshot_id not in ignored_snapshot_ids
1468        }
1469        if modified_snapshots:
1470            directly_modified = []
1471            indirectly_modified = []
1472            metadata_modified = []
1473            for snapshot in modified_snapshots:
1474                if context_diff.directly_modified(snapshot.name):
1475                    directly_modified.append(snapshot)
1476                elif context_diff.indirectly_modified(snapshot.name):
1477                    indirectly_modified.append(snapshot)
1478                elif context_diff.metadata_updated(snapshot.name):
1479                    metadata_modified.append(snapshot)
1480            if directly_modified:
1481                self._print("\n**Directly Modified:**")
1482                for snapshot in sorted(directly_modified):
1483                    self._print(
1484                        f"- `{snapshot.display_name(environment_naming_info, default_catalog)}`"
1485                    )
1486                    if not no_diff:
1487                        self._print(f"```diff\n{context_diff.text_diff(snapshot.name)}\n```")
1488            if indirectly_modified:
1489                self._print("\n**Indirectly Modified:**")
1490                for snapshot in sorted(indirectly_modified):
1491                    self._print(
1492                        f"- `{snapshot.display_name(environment_naming_info, default_catalog)}`"
1493                    )
1494            if metadata_modified:
1495                self._print("\n**Metadata Updated:**")
1496                for snapshot in sorted(metadata_modified):
1497                    self._print(
1498                        f"- `{snapshot.display_name(environment_naming_info, default_catalog)}`"
1499                    )
1500        if ignored_snapshot_ids:
1501            self._print("\n**Ignored Models (Expected Plan Start):**")
1502            for s_id in sorted(ignored_snapshot_ids):
1503                snapshot = context_diff.snapshots[s_id]
1504                self._print(
1505                    f"- `{snapshot.display_name(environment_naming_info, default_catalog)}` ({snapshot.get_latest(start_date(snapshot, context_diff.snapshots.values()))})"
1506                )
1507
1508    def _show_missing_dates(self, plan: Plan, default_catalog: t.Optional[str]) -> None:
1509        """Displays the models with missing dates."""
1510        missing_intervals = plan.missing_intervals
1511        if not missing_intervals:
1512            return
1513        self._print("\n**Models needing backfill (missing dates):**")
1514        for missing in missing_intervals:
1515            snapshot = plan.context_diff.snapshots[missing.snapshot_id]
1516            if not snapshot.is_model:
1517                continue
1518
1519            preview_modifier = ""
1520            if not plan.deployability_index.is_deployable(snapshot):
1521                preview_modifier = " (**preview**)"
1522
1523            self._print(
1524                f"* `{snapshot.display_name(plan.environment_naming_info, default_catalog)}`: {missing.format_intervals(snapshot.node.interval_unit)}{preview_modifier}"
1525            )
1526
1527    def _show_categorized_snapshots(self, plan: Plan, default_catalog: t.Optional[str]) -> None:
1528        context_diff = plan.context_diff
1529        for snapshot in plan.categorized:
1530            if not context_diff.directly_modified(snapshot.name):
1531                continue
1532
1533            category_str = SNAPSHOT_CHANGE_CATEGORY_STR[snapshot.change_category]
1534            tree = Tree(
1535                f"[bold][direct]Directly Modified: {snapshot.display_name(plan.environment_naming_info, default_catalog)} ({category_str})"
1536            )
1537            indirect_tree = None
1538            for child_sid in sorted(plan.indirectly_modified.get(snapshot.snapshot_id, set())):
1539                child_snapshot = context_diff.snapshots[child_sid]
1540                if not indirect_tree:
1541                    indirect_tree = Tree("[indirect]Indirectly Modified Children:")
1542                    tree.add(indirect_tree)
1543                child_category_str = SNAPSHOT_CHANGE_CATEGORY_STR[child_snapshot.change_category]
1544                indirect_tree.add(
1545                    f"[indirect]{child_snapshot.display_name(plan.environment_naming_info, default_catalog)} ({child_category_str})"
1546                )
1547            self._print(f"```diff\n{context_diff.text_diff(snapshot.name)}\n```\n")
1548            self._print("```\n")
1549            self._print(tree)
1550            self._print("\n```")
1551
1552    def log_test_results(
1553        self, result: unittest.result.TestResult, output: str, target_dialect: str
1554    ) -> None:
1555        # import ipywidgets as widgets
1556        if result.wasSuccessful():
1557            self._print(
1558                f"**Successfully Ran `{str(result.testsRun)}` Tests Against `{target_dialect}`**\n\n"
1559            )
1560        else:
1561            self._print(
1562                f"**Num Successful Tests: {result.testsRun - len(result.failures) - len(result.errors)}**\n\n"
1563            )
1564            for test, _ in result.failures + result.errors:
1565                if isinstance(test, ModelTest):
1566                    self._print(f"* Failure Test: `{test.model.name}` - `{test.test_name}`\n\n")
1567            self._print(f"```{output}```\n\n")
1568
1569    def log_error(self, message: str) -> None:
1570        super().log_error(f"```\n{message}```\n\n")
1571
1572
1573class DatabricksMagicConsole(CaptureTerminalConsole):
1574    """
1575    Note: Databricks Magic Console currently does not support progress bars while a plan is being applied. The
1576    NotebookMagicConsole does support progress bars, but they will time out after 5 minutes of execution
1577    and it makes it difficult to see the progress of the plan.
1578    """
1579
1580    def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
1581        super().__init__(*args, **kwargs)
1582        self.evaluation_batch_progress: t.Dict[SnapshotId, t.Tuple[str, int]] = {}
1583        self.promotion_status: t.Tuple[int, int] = (0, 0)
1584        self.model_creation_status: t.Tuple[int, int] = (0, 0)
1585        self.migration_status: t.Tuple[int, int] = (0, 0)
1586
1587    def _print(self, value: t.Any, **kwargs: t.Any) -> None:
1588        super()._print(value, **kwargs)
1589        for captured_output in self._captured_outputs:
1590            print(captured_output)
1591        self.clear_captured_outputs()
1592
1593    def _prompt(self, message: str, **kwargs: t.Any) -> t.Any:
1594        self._print(message)
1595        return super()._prompt("", **kwargs)
1596
1597    def _confirm(self, message: str, **kwargs: t.Any) -> bool:
1598        message = f"{message} [y/n]"
1599        self._print(message)
1600        return super()._confirm("", **kwargs)
1601
1602    def start_evaluation_progress(
1603        self,
1604        batches: t.Dict[Snapshot, int],
1605        environment_naming_info: EnvironmentNamingInfo,
1606        default_catalog: t.Optional[str],
1607    ) -> None:
1608        self.evaluation_batches = batches
1609        self.evaluation_environment_naming_info = environment_naming_info
1610        self.default_catalog = default_catalog
1611
1612    def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
1613        if not self.evaluation_batch_progress.get(snapshot.snapshot_id):
1614            display_name = snapshot.display_name(
1615                self.evaluation_environment_naming_info, self.default_catalog
1616            )
1617            self.evaluation_batch_progress[snapshot.snapshot_id] = (display_name, 0)
1618            print(f"Starting '{display_name}', Total batches: {self.evaluation_batches[snapshot]}")
1619
1620    def update_snapshot_evaluation_progress(
1621        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
1622    ) -> None:
1623        view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
1624        total_batches = self.evaluation_batches[snapshot]
1625
1626        loaded_batches += 1
1627        self.evaluation_batch_progress[snapshot.snapshot_id] = (view_name, loaded_batches)
1628
1629        finished_loading = loaded_batches == total_batches
1630        status = "Loaded" if finished_loading else "Loading"
1631        print(f"{status} '{view_name}', Completed Batches: {loaded_batches}/{total_batches}")
1632        if finished_loading:
1633            total_finished_loading = len(
1634                [
1635                    s
1636                    for s, total in self.evaluation_batches.items()
1637                    if self.evaluation_batch_progress.get(s.snapshot_id, (None, -1))[1] == total
1638                ]
1639            )
1640            total = len(self.evaluation_batch_progress)
1641            print(f"Completed Loading {total_finished_loading}/{total} Models")
1642
1643    def stop_evaluation_progress(self, success: bool = True) -> None:
1644        self.evaluation_batch_progress = {}
1645        super().stop_evaluation_progress(success)
1646        print(f"Loading {'succeeded' if success else 'failed'}")
1647
1648    def start_creation_progress(
1649        self,
1650        total_tasks: int,
1651        environment_naming_info: EnvironmentNamingInfo,
1652        default_catalog: t.Optional[str],
1653    ) -> None:
1654        """Indicates that a new creation progress has begun."""
1655        self.model_creation_status = (0, total_tasks)
1656        print("Starting Creating New Model Versions")
1657
1658    def update_creation_progress(self, snapshot: SnapshotInfoLike) -> None:
1659        """Update the snapshot creation progress."""
1660        num_creations, total_creations = self.model_creation_status
1661        num_creations += 1
1662        self.model_creation_status = (num_creations, total_creations)
1663        if num_creations % 5 == 0:
1664            print(f"Created New Model Versions: {num_creations}/{total_creations}")
1665
1666    def stop_creation_progress(self, success: bool = True) -> None:
1667        """Stop the snapshot creation progress."""
1668        self.model_creation_status = (0, 0)
1669        print(f"New Model Creation {'succeeded' if success else 'failed'}")
1670
1671    def start_promotion_progress(
1672        self,
1673        total_tasks: int,
1674        environment_naming_info: EnvironmentNamingInfo,
1675        default_catalog: t.Optional[str],
1676    ) -> None:
1677        """Indicates that a new snapshot promotion progress has begun."""
1678        self.promotion_status = (0, total_tasks)
1679        print(f"Virtually Updating '{environment_naming_info.name}'")
1680
1681    def update_promotion_progress(self, snapshot: SnapshotInfoLike, promoted: bool) -> None:
1682        """Update the snapshot promotion progress."""
1683        num_promotions, total_promotions = self.promotion_status
1684        num_promotions += 1
1685        self.promotion_status = (num_promotions, total_promotions)
1686        if num_promotions % 5 == 0:
1687            print(f"Virtually Updated {num_promotions}/{total_promotions}")
1688
1689    def stop_promotion_progress(self, success: bool = True) -> None:
1690        """Stop the snapshot promotion progress."""
1691        self.promotion_status = (0, 0)
1692        print(f"Virtual Update {'succeeded' if success else 'failed'}")
1693
1694    def start_migration_progress(self, total_tasks: int) -> None:
1695        """Indicates that a new migration progress has begun."""
1696        self.migration_status = (0, total_tasks)
1697        print("Starting Migration")
1698
1699    def update_migration_progress(self, num_tasks: int) -> None:
1700        """Update the migration progress."""
1701        num_migrations, total_migrations = self.migration_status
1702        num_migrations += num_tasks
1703        self.migration_status = (num_migrations, total_migrations)
1704        if num_migrations % 5 == 0:
1705            print(f"Migration Updated {num_migrations}/{total_migrations}")
1706
1707    def stop_migration_progress(self, success: bool = True) -> None:
1708        """Stop the migration progress."""
1709        self.migration_status = (0, 0)
1710        print(f"Migration {'succeeded' if success else 'failed'}")
1711
1712
1713class DebuggerTerminalConsole(TerminalConsole):
1714    """A terminal console to use while debugging with no fluff, progress bars, etc."""
1715
1716    def __init__(self, console: t.Optional[RichConsole], *args: t.Any, **kwargs: t.Any) -> None:
1717        self.console: RichConsole = console or srich.console
1718
1719    def _write(self, msg: t.Any, *args: t.Any, **kwargs: t.Any) -> None:
1720        self.console.log(msg, *args, **kwargs)
1721
1722    def start_plan_evaluation(self, plan: Plan) -> None:
1723        self._write("Starting plan", plan.plan_id)
1724
1725    def stop_plan_evaluation(self) -> None:
1726        self._write("Stopping plan")
1727
1728    def start_evaluation_progress(
1729        self,
1730        batches: t.Dict[Snapshot, int],
1731        environment_naming_info: EnvironmentNamingInfo,
1732        default_catalog: t.Optional[str],
1733    ) -> None:
1734        self._write(f"Starting evaluation for {len(batches)} snapshots")
1735
1736    def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
1737        self._write(f"Evaluating {snapshot.name}")
1738
1739    def update_snapshot_evaluation_progress(
1740        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
1741    ) -> None:
1742        self._write(f"Evaluating {snapshot.name} | batch={batch_idx} | duration={duration_ms}ms")
1743
1744    def stop_evaluation_progress(self, success: bool = True) -> None:
1745        self._write(f"Stopping evaluation with success={success}")
1746
1747    def start_creation_progress(
1748        self,
1749        total_tasks: int,
1750        environment_naming_info: EnvironmentNamingInfo,
1751        default_catalog: t.Optional[str],
1752    ) -> None:
1753        self._write(f"Starting creation for {total_tasks} snapshots")
1754
1755    def update_creation_progress(self, snapshot: SnapshotInfoLike) -> None:
1756        self._write(f"Creating {snapshot.name}")
1757
1758    def stop_creation_progress(self, success: bool = True) -> None:
1759        self._write(f"Stopping creation with success={success}")
1760
1761    def update_cleanup_progress(self, object_name: str) -> None:
1762        self._write(f"Cleaning up {object_name}")
1763
1764    def start_promotion_progress(
1765        self,
1766        total_tasks: int,
1767        environment_naming_info: EnvironmentNamingInfo,
1768        default_catalog: t.Optional[str],
1769    ) -> None:
1770        self._write(f"Starting promotion for {total_tasks} snapshots")
1771
1772    def update_promotion_progress(self, snapshot: SnapshotInfoLike, promoted: bool) -> None:
1773        self._write(f"Promoting {snapshot.name}")
1774
1775    def stop_promotion_progress(self, success: bool = True) -> None:
1776        self._write(f"Stopping promotion with success={success}")
1777
1778    def start_migration_progress(self, total_tasks: int) -> None:
1779        self._write(f"Starting migration for {total_tasks} snapshots")
1780
1781    def update_migration_progress(self, num_tasks: int) -> None:
1782        self._write(f"Migration {num_tasks}")
1783
1784    def stop_migration_progress(self, success: bool = True) -> None:
1785        self._write(f"Stopping migration with success={success}")
1786
1787    def show_model_difference_summary(
1788        self,
1789        context_diff: ContextDiff,
1790        environment_naming_info: EnvironmentNamingInfo,
1791        default_catalog: t.Optional[str],
1792        no_diff: bool = True,
1793        ignored_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None,
1794    ) -> None:
1795        self._write("Model Difference Summary:")
1796        for added in context_diff.new_snapshots:
1797            self._write(f"  Added: {added}")
1798        for removed in context_diff.removed_snapshots:
1799            self._write(f"  Removed: {removed}")
1800        for modified in context_diff.modified_snapshots:
1801            self._write(f"  Modified: {modified}")
1802
1803    def log_test_results(
1804        self, result: unittest.result.TestResult, output: str, target_dialect: str
1805    ) -> None:
1806        self._write("Test Results:", result)
1807
1808    def show_sql(self, sql: str) -> None:
1809        self._write(sql)
1810
1811    def log_status_update(self, message: str) -> None:
1812        self._write(message, style="bold blue")
1813
1814    def log_error(self, message: str) -> None:
1815        self._write(message, style="bold red")
1816
1817    def log_success(self, message: str) -> None:
1818        self._write(message, style="bold green")
1819
1820    def loading_start(self, message: t.Optional[str] = None) -> uuid.UUID:
1821        self._write(message)
1822        return uuid.uuid4()
1823
1824    def loading_stop(self, id: uuid.UUID) -> None:
1825        self._write("Done")
1826
1827    def show_schema_diff(self, schema_diff: SchemaDiff) -> None:
1828        self._write(schema_diff)
1829
1830    def show_row_diff(self, row_diff: RowDiff, show_sample: bool = True) -> None:
1831        self._write(row_diff)
1832
1833
1834def get_console(**kwargs: t.Any) -> TerminalConsole | DatabricksMagicConsole | NotebookMagicConsole:
1835    """
1836    Returns the console that is appropriate for the current runtime environment.
1837
1838    Note: Google Colab environment is untested and currently assumes is compatible with the base
1839    NotebookMagicConsole.
1840    """
1841    from sqlmesh import RuntimeEnv
1842
1843    runtime_env = RuntimeEnv.get()
1844
1845    runtime_env_mapping = {
1846        RuntimeEnv.DATABRICKS: DatabricksMagicConsole,
1847        RuntimeEnv.JUPYTER: NotebookMagicConsole,
1848        RuntimeEnv.TERMINAL: TerminalConsole,
1849        RuntimeEnv.GOOGLE_COLAB: NotebookMagicConsole,
1850        RuntimeEnv.DEBUGGER: DebuggerTerminalConsole,
1851    }
1852    rich_console_kwargs: t.Dict[str, t.Any] = {"theme": srich.theme}
1853    if runtime_env.is_jupyter or runtime_env.is_google_colab:
1854        rich_console_kwargs["force_jupyter"] = True
1855    return runtime_env_mapping[runtime_env](
1856        **{**{"console": RichConsole(**rich_console_kwargs)}, **kwargs}
1857    )