sqlmesh.magics

   1from __future__ import annotations
   2
   3from io import StringIO
   4
   5import functools
   6import logging
   7import typing as t
   8from argparse import Namespace, SUPPRESS
   9from collections import defaultdict
  10from copy import deepcopy
  11from pathlib import Path
  12
  13from hyperscript import h
  14
  15try:
  16    from IPython.core.display import display  # type: ignore
  17except ImportError:
  18    from IPython.display import display
  19
  20from IPython.core.magic import (
  21    Magics,
  22    cell_magic,
  23    line_cell_magic,
  24    line_magic,
  25    magics_class,
  26)
  27from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
  28from IPython.utils.process import arg_split
  29from rich.jupyter import JupyterRenderable
  30from sqlmesh.cli.project_init import ProjectTemplate, init_example_project
  31from sqlmesh.core import analytics
  32from sqlmesh.core.config import load_configs
  33from sqlmesh.core.config.connection import INIT_DISPLAY_INFO_TO_TYPE
  34from sqlmesh.core.console import create_console, set_console, configure_console
  35from sqlmesh.core.context import Context
  36from sqlmesh.core.dialect import format_model_expressions, parse
  37from sqlmesh.core.model import load_sql_based_model
  38from sqlmesh.core.test import ModelTestMetadata
  39from sqlmesh.utils import yaml, Verbosity, optional_import
  40from sqlmesh.utils.errors import MagicError, MissingContextException, SQLMeshError
  41
  42logger = logging.getLogger(__name__)
  43
  44CONTEXT_VARIABLE_NAMES = [
  45    "context",
  46    "ctx",
  47    "sqlmesh",
  48]
  49
  50
  51def pass_sqlmesh_context(func: t.Callable) -> t.Callable:
  52    @functools.wraps(func)
  53    def wrapper(self: SQLMeshMagics, *args: t.Any, **kwargs: t.Any) -> None:
  54        for variable_name in CONTEXT_VARIABLE_NAMES:
  55            context = self._shell.user_ns.get(variable_name)
  56            if isinstance(context, Context):
  57                break
  58        else:
  59            raise MissingContextException(
  60                f"Context must be defined and initialized with one of these names: {', '.join(CONTEXT_VARIABLE_NAMES)}"
  61            )
  62        old_console = context.console
  63        new_console = create_console(display=self.display)
  64        context.console = new_console
  65        set_console(new_console)
  66        context.refresh()
  67
  68        magic_name = func.__name__
  69        bound_method = getattr(self, magic_name, None)
  70        if bound_method:
  71            args_split = arg_split(args[0])
  72            parser = bound_method.parser
  73
  74            original_parser_actions = deepcopy(parser._actions)
  75            original_parser_defaults = parser._defaults
  76
  77            # Temporarily supress default values, otherwise any missing arg would be set and affect analytics
  78            parser._defaults = {}
  79            for action in parser._actions:
  80                action.default = SUPPRESS
  81
  82            parsed_args, _ = parser.parse_known_args(args_split, Namespace())
  83
  84            parser._actions = original_parser_actions
  85            parser._defaults = original_parser_defaults
  86
  87            command_args = {k for k, v in parsed_args.__dict__.items() if v is not None}
  88            analytics.collector.on_magic_command(command_name=magic_name, command_args=command_args)
  89
  90        func(self, context, *args, **kwargs)
  91
  92        context.console = old_console
  93        set_console(old_console)
  94
  95    return wrapper
  96
  97
  98def parse_expand(value: str) -> t.Union[bool, t.List[str]]:
  99    if value.lower() == "true":
 100        return True
 101    if value.lower() == "false":
 102        return False
 103    return [name.strip() for name in value.split(",") if name.strip()]
 104
 105
 106def format_arguments(func: t.Callable) -> t.Callable:
 107    """Decorator to add common format arguments to magic commands."""
 108    func = argument(
 109        "--normalize",
 110        action="store_true",
 111        help="Whether or not to normalize identifiers to lowercase.",
 112        default=None,
 113    )(func)
 114    func = argument(
 115        "--pad",
 116        type=int,
 117        help="Determines the pad size in a formatted string.",
 118    )(func)
 119    func = argument(
 120        "--indent",
 121        type=int,
 122        help="Determines the indentation size in a formatted string.",
 123    )(func)
 124    func = argument(
 125        "--normalize-functions",
 126        type=str,
 127        help="Whether or not to normalize all function names. Possible values are: 'upper', 'lower'",
 128    )(func)
 129    func = argument(
 130        "--leading-comma",
 131        action="store_true",
 132        help="Determines whether or not the comma is leading or trailing in select expressions. Default is trailing.",
 133        default=None,
 134    )(func)
 135    func = argument(
 136        "--max-text-width",
 137        type=int,
 138        help="The max number of characters in a segment before creating new lines in pretty mode.",
 139    )(func)
 140    return func
 141
 142
 143@magics_class
 144class SQLMeshMagics(Magics):
 145    @property
 146    def display(self) -> t.Callable:
 147        from sqlmesh import RuntimeEnv
 148
 149        if RuntimeEnv.get().is_databricks:
 150            # Use Databricks' special display instead of the normal IPython display
 151            return self._shell.user_ns["display"]
 152        return display
 153
 154    @property
 155    def _shell(self) -> t.Any:
 156        # Make mypy happy.
 157        if not self.shell:
 158            raise RuntimeError("IPython Magics are in invalid state")
 159        return self.shell
 160
 161    @magic_arguments()
 162    @argument(
 163        "paths",
 164        type=str,
 165        nargs="+",
 166        default="",
 167        help="The path(s) to the SQLMesh project(s).",
 168    )
 169    @argument(
 170        "--config",
 171        type=str,
 172        help="Name of the config object. Only applicable to configuration defined using Python script.",
 173    )
 174    @argument("--gateway", type=str, help="The name of the gateway.")
 175    @argument("--ignore-warnings", action="store_true", help="Ignore warnings.")
 176    @argument("--debug", action="store_true", help="Enable debug mode.")
 177    @argument("--log-file-dir", type=str, help="The directory to write the log file to.")
 178    @argument(
 179        "--dotenv", type=str, help="Path to a custom .env file to load environment variables from."
 180    )
 181    @line_magic
 182    def context(self, line: str) -> None:
 183        """Sets the context in the user namespace."""
 184        from sqlmesh import configure_logging, remove_excess_logs
 185
 186        args = parse_argstring(self.context, line)
 187        log_file_dir = args.log_file_dir
 188
 189        configure_logging(
 190            args.debug,
 191            log_file_dir=log_file_dir,
 192            ignore_warnings=args.ignore_warnings,
 193        )
 194        configure_console(ignore_warnings=args.ignore_warnings)
 195
 196        dotenv_path = Path(args.dotenv) if args.dotenv else None
 197        configs = load_configs(
 198            args.config, Context.CONFIG_TYPE, args.paths, dotenv_path=dotenv_path
 199        )
 200        log_limit = list(configs.values())[0].log_limit
 201
 202        remove_excess_logs(log_file_dir, log_limit)
 203
 204        try:
 205            context = Context(paths=args.paths, config=configs, gateway=args.gateway)
 206            self._shell.user_ns["context"] = context
 207        except Exception:
 208            if args.debug:
 209                logger.exception("Failed to initialize SQLMesh context")
 210            raise
 211
 212        context.console.log_success(f"SQLMesh project context set to: {', '.join(args.paths)}")
 213
 214    @magic_arguments()
 215    @argument("path", type=str, help="The path where the new SQLMesh project should be created.")
 216    @argument(
 217        "engine",
 218        type=str,
 219        help=f"Project SQL engine. Supported values: '{', '.join([info[1] for info in sorted(INIT_DISPLAY_INFO_TO_TYPE.values(), key=lambda x: x[0])])}'.",  # type: ignore
 220    )
 221    @argument(
 222        "--template",
 223        "-t",
 224        type=str,
 225        help="Project template. Supported values: dbt, default, empty.",
 226    )
 227    @argument(
 228        "--dlt-pipeline",
 229        type=str,
 230        help="DLT pipeline for which to generate a SQLMesh project. Use alongside template: dlt",
 231    )
 232    @argument(
 233        "--dlt-path",
 234        type=str,
 235        help="The DLT pipelines working directory, where DLT stores pipeline state (by default ~/.dlt/pipelines). Use alongside template: dlt",
 236    )
 237    @line_magic
 238    def init(self, line: str) -> None:
 239        """Creates a SQLMesh project scaffold with a default SQL dialect."""
 240        args = parse_argstring(self.init, line)
 241        try:
 242            project_template = ProjectTemplate(
 243                args.template.lower() if args.template else "default"
 244            )
 245        except ValueError:
 246            raise MagicError(f"Invalid project template '{args.template}'")
 247        init_example_project(
 248            path=args.path,
 249            engine_type=args.engine,
 250            dialect=None,
 251            template=project_template,
 252            pipeline=args.dlt_pipeline,
 253            dlt_path=args.dlt_path,
 254        )
 255        html = str(
 256            h(
 257                "div",
 258                h(
 259                    "span",
 260                    {"style": {"color": "green", "font-weight": "bold"}},
 261                    "SQLMesh project scaffold created",
 262                ),
 263            )
 264        )
 265        self.display(JupyterRenderable(html=html, text=""))
 266
 267    @magic_arguments()
 268    @argument("model", type=str, help="The model.")
 269    @argument("--start", "-s", type=str, help="Start date to render.")
 270    @argument("--end", "-e", type=str, help="End date to render.")
 271    @argument("--execution-time", type=str, help="Execution time.")
 272    @argument("--dialect", "-d", type=str, help="The rendered dialect.")
 273    @line_cell_magic
 274    @pass_sqlmesh_context
 275    def model(self, context: Context, line: str, sql: t.Optional[str] = None) -> None:
 276        """Renders the model and automatically fills in an editable cell with the model definition."""
 277        args = parse_argstring(self.model, line)
 278
 279        model = context.get_model(args.model, raise_if_missing=True)
 280        config = context.config_for_node(model)
 281
 282        if sql:
 283            expressions = parse(sql, default_dialect=config.dialect)
 284            loaded = load_sql_based_model(
 285                expressions,
 286                macros=context._macros,
 287                jinja_macros=context._jinja_macros,
 288                path=model._path,
 289                dialect=config.dialect,
 290                time_column_format=config.time_column_format,
 291                physical_schema_mapping=context.config.physical_schema_mapping,
 292                default_catalog=context.default_catalog,
 293            )
 294
 295            if loaded.name == args.model:
 296                model = loaded
 297        else:
 298            if model._path:
 299                with open(model._path, "r", encoding="utf-8") as file:
 300                    expressions = parse(file.read(), default_dialect=config.dialect)
 301
 302        formatted = format_model_expressions(
 303            expressions,
 304            model.dialect,
 305            rewrite_casts=not config.format.no_rewrite_casts,
 306            **config.format.generator_options,
 307        )
 308
 309        self._shell.set_next_input(
 310            "\n".join(
 311                [
 312                    " ".join(["%%model", line]),
 313                    formatted,
 314                ]
 315            ),
 316            replace=True,
 317        )
 318
 319        if model._path:
 320            with open(model._path, "w", encoding="utf-8") as file:
 321                file.write(formatted)
 322
 323        if sql:
 324            context.console.log_success(f"Model `{args.model}` updated")
 325
 326        context.upsert_model(model)
 327        context.console.show_sql(
 328            context.render(
 329                model.name,
 330                start=args.start,
 331                end=args.end,
 332                execution_time=args.execution_time,
 333            ).sql(pretty=True, dialect=args.dialect or model.dialect)
 334        )
 335
 336    @magic_arguments()
 337    @argument("model", type=str, help="The model.")
 338    @argument("test_name", type=str, nargs="?", default=None, help="The test name to display")
 339    @argument("--ls", action="store_true", help="List tests associated with a model")
 340    @line_cell_magic
 341    @pass_sqlmesh_context
 342    def test(self, context: Context, line: str, test_def_raw: t.Optional[str] = None) -> None:
 343        """Allow the user to list tests for a model, output a specific test, and then write their changes back"""
 344        args = parse_argstring(self.test, line)
 345        if not args.test_name and not args.ls:
 346            raise MagicError("Must provide either test name or `--ls` to list tests")
 347
 348        test_meta = context.select_tests()
 349
 350        tests: t.Dict[str, t.Dict[str, ModelTestMetadata]] = defaultdict(dict)
 351        for model_test_metadata in test_meta:
 352            model = model_test_metadata.body.get("model")
 353            if not model:
 354                context.console.log_error(
 355                    f"Test found that does not have `model` defined: {model_test_metadata.path}"
 356                )
 357            else:
 358                tests[model][model_test_metadata.test_name] = model_test_metadata
 359
 360        model = context.get_model(args.model, raise_if_missing=True)
 361
 362        if args.ls:
 363            # TODO: Provide better UI for displaying tests
 364            for test_name in tests[model.name]:
 365                context.console.log_status_update(test_name)
 366            return
 367
 368        test = tests[model.name][args.test_name]
 369        test_def = yaml.load(test_def_raw) if test_def_raw else test.body
 370        test_def_output = yaml.dump(test_def)
 371
 372        self._shell.set_next_input(
 373            "\n".join(
 374                [
 375                    " ".join(["%%test", line]),
 376                    test_def_output,
 377                ]
 378            ),
 379            replace=True,
 380        )
 381
 382        with open(test.path, "r+", encoding="utf-8") as file:
 383            content = yaml.load(file.read())
 384            content[args.test_name] = test_def
 385            file.seek(0)
 386            yaml.dump(content, file)
 387            file.truncate()
 388
 389    @magic_arguments()
 390    @argument(
 391        "environment",
 392        nargs="?",
 393        type=str,
 394        help="The environment to run the plan against",
 395    )
 396    @argument("--start", "-s", type=str, help="Start date to backfill.")
 397    @argument("--end", "-e", type=str, help="End date to backfill.")
 398    @argument("--execution-time", type=str, help="Execution time.")
 399    @argument(
 400        "--create-from",
 401        type=str,
 402        help="The environment to create the target environment from if it doesn't exist. Default: prod.",
 403    )
 404    @argument(
 405        "--skip-tests",
 406        "-t",
 407        action="store_true",
 408        help="Skip the unit tests defined for the model.",
 409    )
 410    @argument(
 411        "--skip-linter",
 412        action="store_true",
 413        help="Skip the linter for the model.",
 414    )
 415    @argument(
 416        "--restate-model",
 417        "-r",
 418        type=str,
 419        nargs="*",
 420        help="Restate data for specified models (and models downstream from the one specified). For production environment, all related model versions will have their intervals wiped, but only the current versions will be backfilled. For development environment, only the current model versions will be affected.",
 421    )
 422    @argument(
 423        "--no-gaps",
 424        "-g",
 425        action="store_true",
 426        help="Ensure that new snapshots have no data gaps when comparing to existing snapshots for matching models in the target environment.",
 427    )
 428    @argument(
 429        "--skip-backfill",
 430        "--dry-run",
 431        action="store_true",
 432        help="Skip the backfill step and only create a virtual update for the plan.",
 433    )
 434    @argument(
 435        "--empty-backfill",
 436        action="store_true",
 437        help="Produce empty backfill. Like --skip-backfill no models will be backfilled, unlike --skip-backfill missing intervals will be recorded as if they were backfilled.",
 438    )
 439    @argument(
 440        "--forward-only",
 441        action="store_true",
 442        help="Create a plan for forward-only changes.",
 443        default=None,
 444    )
 445    @argument(
 446        "--effective-from",
 447        type=str,
 448        help="The effective date from which to apply forward-only changes on production.",
 449    )
 450    @argument(
 451        "--no-prompts",
 452        action="store_true",
 453        help="Disables interactive prompts for the backfill time range. Please note that if this flag is set and there are uncategorized changes, plan creation will fail.",
 454        default=None,
 455    )
 456    @argument(
 457        "--auto-apply",
 458        action="store_true",
 459        help="Automatically applies the new plan after creation.",
 460        default=None,
 461    )
 462    @argument(
 463        "--no-auto-categorization",
 464        action="store_true",
 465        help="Disable automatic change categorization.",
 466        default=None,
 467    )
 468    @argument(
 469        "--include-unmodified",
 470        action="store_true",
 471        help="Include unmodified models in the target environment.",
 472        default=None,
 473    )
 474    @argument(
 475        "--select-model",
 476        type=str,
 477        nargs="*",
 478        help="Select specific model changes that should be included in the plan.",
 479    )
 480    @argument(
 481        "--backfill-model",
 482        type=str,
 483        nargs="*",
 484        help="Backfill only the models whose names match the expression.",
 485    )
 486    @argument(
 487        "--no-diff",
 488        action="store_true",
 489        help="Hide text differences for changed models.",
 490        default=None,
 491    )
 492    @argument(
 493        "--run",
 494        action="store_true",
 495        help="Run latest intervals as part of the plan application (prod environment only).",
 496    )
 497    @argument(
 498        "--ignore-cron",
 499        action="store_true",
 500        help="Run for all missing intervals, ignoring individual cron schedules. Only applies if --run is set.",
 501        default=None,
 502    )
 503    @argument(
 504        "--enable-preview",
 505        action="store_true",
 506        help="Enable preview for forward-only models when targeting a development environment.",
 507        default=None,
 508    )
 509    @argument(
 510        "--diff-rendered",
 511        action="store_true",
 512        help="Output text differences for the rendered versions of the models and standalone audits",
 513    )
 514    @argument(
 515        "--verbose",
 516        "-v",
 517        action="count",
 518        default=0,
 519        help="Verbose output. Use -vv for very verbose.",
 520    )
 521    @line_magic
 522    @pass_sqlmesh_context
 523    def plan(self, context: Context, line: str) -> None:
 524        """Goes through a set of prompts to both establish a plan and apply it"""
 525        args = parse_argstring(self.plan, line)
 526
 527        setattr(context.console, "verbosity", Verbosity(args.verbose))
 528
 529        context.plan(
 530            args.environment,
 531            start=args.start,
 532            end=args.end,
 533            execution_time=args.execution_time,
 534            create_from=args.create_from,
 535            skip_tests=args.skip_tests,
 536            restate_models=args.restate_model,
 537            backfill_models=args.backfill_model,
 538            no_gaps=args.no_gaps,
 539            skip_backfill=args.skip_backfill,
 540            empty_backfill=args.empty_backfill,
 541            forward_only=args.forward_only,
 542            no_prompts=args.no_prompts,
 543            auto_apply=args.auto_apply,
 544            no_auto_categorization=args.no_auto_categorization,
 545            effective_from=args.effective_from,
 546            include_unmodified=args.include_unmodified,
 547            select_models=args.select_model,
 548            no_diff=args.no_diff,
 549            run=args.run,
 550            ignore_cron=args.run,
 551            enable_preview=args.enable_preview,
 552            diff_rendered=args.diff_rendered,
 553        )
 554
 555    @magic_arguments()
 556    @argument(
 557        "environment",
 558        nargs="?",
 559        type=str,
 560        help="The environment to run against",
 561    )
 562    @argument("--start", "-s", type=str, help="Start date to evaluate.")
 563    @argument("--end", "-e", type=str, help="End date to evaluate.")
 564    @argument("--skip-janitor", action="store_true", help="Skip the janitor task.")
 565    @argument(
 566        "--ignore-cron",
 567        action="store_true",
 568        help="Run for all missing intervals, ignoring individual cron schedules.",
 569    )
 570    @argument(
 571        "--select-model",
 572        type=str,
 573        nargs="*",
 574        help="Select specific models to run. Note: this always includes upstream dependencies.",
 575    )
 576    @argument(
 577        "--exit-on-env-update",
 578        type=int,
 579        help="If set, the command will exit with the specified code if the run is interrupted by an update to the target environment.",
 580    )
 581    @argument(
 582        "--no-auto-upstream",
 583        action="store_true",
 584        help="Do not automatically include upstream models. Only applicable when --select-model is used. Note: this may result in missing / invalid data for the selected models.",
 585    )
 586    @line_magic
 587    @pass_sqlmesh_context
 588    def run_dag(self, context: Context, line: str) -> None:
 589        """Evaluate the DAG of models using the built-in scheduler."""
 590        args = parse_argstring(self.run_dag, line)
 591
 592        completion_status = context.run(
 593            args.environment,
 594            start=args.start,
 595            end=args.end,
 596            skip_janitor=args.skip_janitor,
 597            ignore_cron=args.ignore_cron,
 598            select_models=args.select_model,
 599            exit_on_env_update=args.exit_on_env_update,
 600            no_auto_upstream=args.no_auto_upstream,
 601        )
 602        if completion_status.is_failure:
 603            raise SQLMeshError("Error Running DAG. Check logs for details.")
 604
 605    @magic_arguments()
 606    @argument("model", type=str, help="The model.")
 607    @argument("--start", "-s", type=str, help="Start date to render.")
 608    @argument("--end", "-e", type=str, help="End date to render.")
 609    @argument("--execution-time", type=str, help="Execution time.")
 610    @argument(
 611        "--limit",
 612        type=int,
 613        help="The number of rows which the query should be limited to.",
 614    )
 615    @line_magic
 616    @pass_sqlmesh_context
 617    def evaluate(self, context: Context, line: str) -> None:
 618        """Evaluate a model query and fetches a dataframe."""
 619        context.refresh()
 620
 621        snowpark = optional_import("snowflake.snowpark")
 622        args = parse_argstring(self.evaluate, line)
 623
 624        df = context.evaluate(
 625            args.model,
 626            start=args.start,
 627            end=args.end,
 628            execution_time=args.execution_time,
 629            limit=args.limit,
 630        )
 631
 632        if snowpark and isinstance(df, snowpark.DataFrame):
 633            df = df.limit(args.limit or 100).to_pandas()
 634
 635        self.display(df)
 636
 637    @magic_arguments()
 638    @argument("model", type=str, help="The model.")
 639    @argument("--start", "-s", type=str, help="Start date to render.")
 640    @argument("--end", "-e", type=str, help="End date to render.")
 641    @argument("--execution-time", type=str, help="Execution time.")
 642    @argument(
 643        "--expand",
 644        type=parse_expand,
 645        help="Whether or not to use expand materialized models, defaults to False. If 'true', all referenced models are expanded as raw queries. If a comma-separated list of model names, only those models are expanded as raw queries.",
 646    )
 647    @argument("--dialect", type=str, help="SQL dialect to render.")
 648    @argument("--no-format", action="store_true", help="Disable fancy formatting of the query.")
 649    @format_arguments
 650    @line_magic
 651    @pass_sqlmesh_context
 652    def render(self, context: Context, line: str) -> None:
 653        """Renders a model's query, optionally expanding referenced models."""
 654        context.refresh()
 655        render_opts = vars(parse_argstring(self.render, line))
 656        model = render_opts.pop("model")
 657        dialect = render_opts.pop("dialect", None)
 658        expand = render_opts.pop("expand", False)
 659
 660        model = context.get_model(model, raise_if_missing=True)
 661
 662        query = context.render(
 663            model,
 664            start=render_opts.pop("start", None),
 665            end=render_opts.pop("end", None),
 666            execution_time=render_opts.pop("execution_time", None),
 667            expand=expand,
 668        )
 669
 670        no_format = render_opts.pop("no_format", False)
 671
 672        format_config = context.config_for_node(model).format
 673        format_options = {
 674            **format_config.generator_options,
 675            **{k: v for k, v in render_opts.items() if v is not None},
 676        }
 677
 678        sql = query.sql(
 679            pretty=True,
 680            dialect=context.config.dialect if dialect is None else dialect,
 681            **format_options,
 682        )
 683
 684        if no_format:
 685            context.console.log_status_update(sql)
 686        else:
 687            context.console.show_sql(sql)
 688
 689    @magic_arguments()
 690    @argument(
 691        "df_var",
 692        default=None,
 693        nargs="?",
 694        type=str,
 695        help="An optional variable name to store the resulting dataframe.",
 696    )
 697    @cell_magic
 698    @pass_sqlmesh_context
 699    def fetchdf(self, context: Context, line: str, sql: str) -> None:
 700        """Fetches a dataframe from sql, optionally storing it in a variable."""
 701        args = parse_argstring(self.fetchdf, line)
 702        df = context.fetchdf(sql)
 703        if args.df_var:
 704            self._shell.user_ns[args.df_var] = df
 705        self.display(df)
 706
 707    @magic_arguments()
 708    @argument("--file", "-f", type=str, help="An optional file path to write the HTML output to.")
 709    @argument(
 710        "--select-model",
 711        type=str,
 712        nargs="*",
 713        help="Select specific models to include in the dag.",
 714    )
 715    @line_magic
 716    @pass_sqlmesh_context
 717    def dag(self, context: Context, line: str) -> None:
 718        """Displays the HTML DAG."""
 719        args = parse_argstring(self.dag, line)
 720        dag = context.get_dag(args.select_model)
 721        if args.file:
 722            with open(args.file, "w", encoding="utf-8") as file:
 723                file.write(str(dag))
 724        # TODO: Have this go through console instead of calling display directly
 725        self.display(dag)
 726
 727    @magic_arguments()
 728    @line_magic
 729    @pass_sqlmesh_context
 730    def migrate(self, context: Context, line: str) -> None:
 731        """Migrate SQLMesh to the current running version."""
 732        context.migrate()
 733        context.console.log_success("Migration complete")
 734
 735    @magic_arguments()
 736    @argument(
 737        "--strict",
 738        action="store_true",
 739        help="Raise an error if the external model is missing in the database",
 740    )
 741    @line_magic
 742    @pass_sqlmesh_context
 743    def create_external_models(self, context: Context, line: str) -> None:
 744        """Create a schema file containing external model schemas."""
 745        args = parse_argstring(self.create_external_models, line)
 746        context.create_external_models(strict=args.strict)
 747
 748    @magic_arguments()
 749    @argument(
 750        "source_to_target",
 751        type=str,
 752        metavar="SOURCE:TARGET",
 753        help="Source and target in `SOURCE:TARGET` format",
 754    )
 755    @argument(
 756        "--on",
 757        type=str,
 758        nargs="*",
 759        help="The column to join on. Can be specified multiple times. The model grain will be used if not specified.",
 760    )
 761    @argument(
 762        "--skip-columns",
 763        type=str,
 764        nargs="*",
 765        help="The column(s) to skip when comparing the source and target table.",
 766    )
 767    @argument(
 768        "--model",
 769        type=str,
 770        help="The model to diff against when source and target are environments and not tables.",
 771    )
 772    @argument(
 773        "--where",
 774        type=str,
 775        help="An optional where statement to filter results.",
 776    )
 777    @argument(
 778        "--limit",
 779        type=int,
 780        default=20,
 781        help="The limit of the sample dataframe.",
 782    )
 783    @argument(
 784        "--show-sample",
 785        action="store_true",
 786        help="Show a sample of the rows that differ. With many columns, the output can be very wide.",
 787    )
 788    @argument(
 789        "--decimals",
 790        type=int,
 791        default=3,
 792        help="The number of decimal places to keep when comparing floating point columns. Default: 3",
 793    )
 794    @argument(
 795        "--select-model",
 796        type=str,
 797        nargs="*",
 798        help="Specify one or more models to data diff. Use wildcards to diff multiple models. Ex: '*' (all models with applied plan diffs), 'demo.model+' (this and downstream models), 'git:feature_branch' (models with direct modifications in this branch only)",
 799    )
 800    @argument(
 801        "--skip-grain-check",
 802        action="store_true",
 803        help="Disable the check for a primary key (grain) that is missing or is not unique.",
 804    )
 805    @argument(
 806        "--warn-grain-check",
 807        action="store_true",
 808        help="Warn if any selected model is missing a grain, and compute diffs for the remaining models.",
 809    )
 810    @argument(
 811        "--schema-diff-ignore-case",
 812        action="store_true",
 813        help="If set, when performing a schema diff the case of column names is ignored when matching between the two schemas. For example, 'col_a' in the source schema and 'COL_A' in the target schema will be treated as the same column.",
 814    )
 815    @line_magic
 816    @pass_sqlmesh_context
 817    def table_diff(self, context: Context, line: str) -> None:
 818        """Show the diff between two tables.
 819
 820        Can either be two tables or two environments and a model.
 821        """
 822        args = parse_argstring(self.table_diff, line)
 823        source, target = args.source_to_target.split(":")
 824        select_models = {args.model} if args.model else args.select_model or None
 825        context.table_diff(
 826            source=source,
 827            target=target,
 828            on=args.on,
 829            skip_columns=args.skip_columns,
 830            select_models=select_models,
 831            where=args.where,
 832            limit=args.limit,
 833            show_sample=args.show_sample,
 834            decimals=args.decimals,
 835            skip_grain_check=args.skip_grain_check,
 836            warn_grain_check=args.warn_grain_check,
 837            schema_diff_ignore_case=args.schema_diff_ignore_case,
 838        )
 839
 840    @magic_arguments()
 841    @argument(
 842        "model_name",
 843        nargs="?",
 844        type=str,
 845        help="The name of the model to get the table name for.",
 846    )
 847    @argument(
 848        "--environment",
 849        type=str,
 850        help="The environment to source the model version from.",
 851    )
 852    @argument(
 853        "--prod",
 854        action="store_true",
 855        help="If set, return the name of the physical table that will be used in production for the model version promoted in the target environment.",
 856    )
 857    @line_magic
 858    @pass_sqlmesh_context
 859    def table_name(self, context: Context, line: str) -> None:
 860        """Prints the name of the physical table for the given model."""
 861        args = parse_argstring(self.table_name, line)
 862        context.console.log_status_update(
 863            context.table_name(args.model_name, args.environment, args.prod)
 864        )
 865
 866    @magic_arguments()
 867    @argument(
 868        "pipeline",
 869        nargs="?",
 870        type=str,
 871        help="The dlt pipeline to attach for this SQLMesh project.",
 872    )
 873    @argument(
 874        "--table",
 875        "-t",
 876        type=str,
 877        nargs="*",
 878        help="The specific dlt tables to refresh in the SQLMesh models.",
 879    )
 880    @argument(
 881        "--force",
 882        "-f",
 883        action="store_true",
 884        help="If set, existing models are overwritten with the new DLT tables.",
 885    )
 886    @argument(
 887        "--dlt-path",
 888        type=str,
 889        help="The DLT pipelines working directory, where DLT stores pipeline state (by default ~/.dlt/pipelines).",
 890    )
 891    @line_magic
 892    @pass_sqlmesh_context
 893    def dlt_refresh(self, context: Context, line: str) -> None:
 894        """Attaches to a DLT pipeline with the option to update specific or all missing tables in the SQLMesh project."""
 895        from sqlmesh.integrations.dlt import generate_dlt_models
 896
 897        args = parse_argstring(self.dlt_refresh, line)
 898        sqlmesh_models = generate_dlt_models(
 899            context, args.pipeline, list(args.table or []), args.force, args.dlt_path
 900        )
 901        if sqlmesh_models:
 902            model_names = "\n".join([f"- {model_name}" for model_name in sqlmesh_models])
 903            context.console.log_success(f"Updated SQLMesh project with models:\n{model_names}")
 904        else:
 905            context.console.log_success("All SQLMesh models are up to date.")
 906
 907    @magic_arguments()
 908    @argument(
 909        "--read",
 910        type=str,
 911        default="",
 912        help="The input dialect of the sql string.",
 913    )
 914    @argument(
 915        "--write",
 916        type=str,
 917        default="",
 918        help="The output dialect of the sql string.",
 919    )
 920    @line_cell_magic
 921    @pass_sqlmesh_context
 922    def rewrite(self, context: Context, line: str, sql: str) -> None:
 923        """Rewrite a sql expression with semantic references into an executable query.
 924
 925        https://sqlmesh.readthedocs.io/en/latest/concepts/metrics/overview/
 926        """
 927        args = parse_argstring(self.rewrite, line)
 928        context.console.show_sql(
 929            context.rewrite(sql, args.read).sql(
 930                dialect=args.write or context.config.dialect, pretty=True
 931            )
 932        )
 933
 934    @magic_arguments()
 935    @argument(
 936        "--transpile",
 937        "-t",
 938        type=str,
 939        help="Transpile project models to the specified dialect.",
 940    )
 941    @argument(
 942        "--check",
 943        action="store_true",
 944        help="Whether or not to check formatting (but not actually format anything).",
 945        default=None,
 946    )
 947    @argument(
 948        "--append-newline",
 949        action="store_true",
 950        help="Include a newline at the end of the output.",
 951        default=None,
 952    )
 953    @argument(
 954        "--no-rewrite-casts",
 955        action="store_true",
 956        help="Preserve the existing casts, without rewriting them to use the :: syntax.",
 957        default=None,
 958    )
 959    @format_arguments
 960    @line_magic
 961    @pass_sqlmesh_context
 962    def format(self, context: Context, line: str) -> bool:
 963        """Format all SQL models and audits."""
 964        format_opts = vars(parse_argstring(self.format, line))
 965        if format_opts.pop("no_rewrite_casts", None):
 966            format_opts["rewrite_casts"] = False
 967
 968        return context.format(**{k: v for k, v in format_opts.items() if v is not None})
 969
 970    @magic_arguments()
 971    @argument("environment", type=str, help="The environment to diff local state against.")
 972    @line_magic
 973    @pass_sqlmesh_context
 974    def diff(self, context: Context, line: str) -> None:
 975        """Show the diff between the local state and the target environment."""
 976        args = parse_argstring(self.diff, line)
 977        context.diff(args.environment)
 978
 979    @magic_arguments()
 980    @argument("environment", type=str, help="The environment to invalidate.")
 981    @line_magic
 982    @pass_sqlmesh_context
 983    def invalidate(self, context: Context, line: str) -> None:
 984        """Invalidate the target environment, forcing its removal during the next run of the janitor process."""
 985        args = parse_argstring(self.invalidate, line)
 986        context.invalidate_environment(args.environment)
 987
 988    @magic_arguments()
 989    @argument(
 990        "--ignore-ttl",
 991        action="store_true",
 992        help="Cleanup snapshots that are not referenced in any environment, regardless of when they're set to expire",
 993    )
 994    @line_magic
 995    @pass_sqlmesh_context
 996    def janitor(self, context: Context, line: str) -> None:
 997        """Run the janitor process to clean up old environments and expired snapshots."""
 998        args = parse_argstring(self.janitor, line)
 999        context.run_janitor(ignore_ttl=args.ignore_ttl)
1000
1001    @magic_arguments()
1002    @argument("model", type=str)
1003    @argument(
1004        "--query",
1005        "-q",
1006        type=str,
1007        nargs="+",
1008        default=[],
1009        help="Queries that will be used to generate data for the model's dependencies.",
1010    )
1011    @argument(
1012        "--overwrite",
1013        "-o",
1014        action="store_true",
1015        help="When true, the fixture file will be overwritten in case it already exists.",
1016    )
1017    @argument(
1018        "--var",
1019        "-v",
1020        type=str,
1021        nargs="+",
1022        help="Key-value pairs that will define variables needed by the model.",
1023    )
1024    @argument(
1025        "--path",
1026        "-p",
1027        type=str,
1028        help="The file path corresponding to the fixture, relative to the test directory. "
1029        "By default, the fixture will be created under the test directory and the file "
1030        "name will be inferred based on the test's name.",
1031    )
1032    @argument(
1033        "--name",
1034        "-n",
1035        type=str,
1036        help="The name of the test that will be created. By default, it's inferred based on the model's name.",
1037    )
1038    @argument(
1039        "--include-ctes",
1040        action="store_true",
1041        help="When true, CTE fixtures will also be generated.",
1042    )
1043    @line_magic
1044    @pass_sqlmesh_context
1045    def create_test(self, context: Context, line: str) -> None:
1046        """Generate a unit test fixture for a given model."""
1047        args = parse_argstring(self.create_test, line)
1048        queries = iter(args.query)
1049        variables = iter(args.var) if args.var else None
1050        context.create_test(
1051            args.model,
1052            input_queries={k: v.strip('"') for k, v in dict(zip(queries, queries)).items()},
1053            overwrite=args.overwrite,
1054            variables=dict(zip(variables, variables)) if variables else None,
1055            path=args.path,
1056            name=args.name,
1057            include_ctes=args.include_ctes,
1058        )
1059
1060    @magic_arguments()
1061    @argument("tests", nargs="*", type=str)
1062    @argument(
1063        "--pattern",
1064        "-k",
1065        nargs="*",
1066        type=str,
1067        help="Only run tests that match the pattern of substring.",
1068    )
1069    @argument(
1070        "--verbose",
1071        "-v",
1072        action="count",
1073        default=0,
1074        help="Verbose output. Use -vv for very verbose.",
1075    )
1076    @argument(
1077        "--preserve-fixtures",
1078        action="store_true",
1079        help="Preserve the fixture tables in the testing database, useful for debugging.",
1080    )
1081    @line_magic
1082    @pass_sqlmesh_context
1083    def run_test(self, context: Context, line: str) -> None:
1084        """Run unit test(s)."""
1085        args = parse_argstring(self.run_test, line)
1086
1087        context.test(
1088            match_patterns=args.pattern,
1089            tests=args.tests,
1090            verbosity=Verbosity(args.verbose),
1091            preserve_fixtures=args.preserve_fixtures,
1092            stream=StringIO(),  # consume the output instead of redirecting to stdout
1093        )
1094
1095    @magic_arguments()
1096    @argument(
1097        "models", type=str, nargs="*", help="A model to audit. Multiple models can be audited."
1098    )
1099    @argument("--start", "-s", type=str, help="Start date to audit.")
1100    @argument("--end", "-e", type=str, help="End date to audit.")
1101    @argument("--execution-time", type=str, help="Execution time.")
1102    @line_magic
1103    @pass_sqlmesh_context
1104    def audit(self, context: Context, line: str) -> bool:
1105        """Run audit(s)"""
1106        args = parse_argstring(self.audit, line)
1107        return context.audit(
1108            models=args.models, start=args.start, end=args.end, execution_time=args.execution_time
1109        )
1110
1111    @magic_arguments()
1112    @argument("environment", nargs="?", type=str, help="The environment to check intervals for.")
1113    @argument(
1114        "--no-signals",
1115        action="store_true",
1116        help="Disable signal checks and only show missing intervals.",
1117        default=False,
1118    )
1119    @argument(
1120        "--select-model",
1121        type=str,
1122        nargs="*",
1123        help="Select specific model changes that should be included in the plan.",
1124    )
1125    @argument("--start", "-s", type=str, help="Start date of intervals to check for.")
1126    @argument("--end", "-e", type=str, help="End date of intervals to check for.")
1127    @line_magic
1128    @pass_sqlmesh_context
1129    def check_intervals(self, context: Context, line: str) -> None:
1130        """Show missing intervals in an environment, respecting signals."""
1131        args = parse_argstring(self.check_intervals, line)
1132
1133        context.console.show_intervals(
1134            context.check_intervals(
1135                environment=args.environment,
1136                no_signals=args.no_signals,
1137                select_models=args.select_model,
1138                start=args.start,
1139                end=args.end,
1140            )
1141        )
1142
1143    @magic_arguments()
1144    @argument(
1145        "--skip-connection",
1146        action="store_true",
1147        help="Skip the connection test.",
1148        default=False,
1149    )
1150    @argument(
1151        "--verbose",
1152        "-v",
1153        action="count",
1154        default=0,
1155        help="Verbose output. Use -vv for very verbose.",
1156    )
1157    @line_magic
1158    @pass_sqlmesh_context
1159    def info(self, context: Context, line: str) -> None:
1160        """Display SQLMesh project information."""
1161        args = parse_argstring(self.info, line)
1162        context.print_info(skip_connection=args.skip_connection, verbosity=Verbosity(args.verbose))
1163
1164    @magic_arguments()
1165    @line_magic
1166    @pass_sqlmesh_context
1167    def rollback(self, context: Context, line: str) -> None:
1168        """Rollback SQLMesh to the previous migration."""
1169        context.rollback()
1170
1171    @magic_arguments()
1172    @line_magic
1173    @pass_sqlmesh_context
1174    def clean(self, context: Context, line: str) -> None:
1175        """Clears the SQLMesh cache and any build artifacts."""
1176        context.clear_caches()
1177        context.console.log_success("SQLMesh cache and build artifacts cleared")
1178
1179    @magic_arguments()
1180    @line_magic
1181    @pass_sqlmesh_context
1182    def environments(self, context: Context, line: str) -> None:
1183        """Prints the list of SQLMesh environments with its expiry datetime."""
1184        context.print_environment_names()
1185
1186    @magic_arguments()
1187    @argument(
1188        "--models",
1189        "--model",
1190        type=str,
1191        nargs="*",
1192        help="A model to lint. Multiple models can be linted. If no models are specified, every model will be linted.",
1193    )
1194    @line_magic
1195    @pass_sqlmesh_context
1196    def lint(self, context: Context, line: str) -> None:
1197        """Run linter for target model(s)"""
1198        args = parse_argstring(self.lint, line)
1199        context.lint_models(args.models)
1200
1201    @magic_arguments()
1202    @line_magic
1203    @pass_sqlmesh_context
1204    def destroy(self, context: Context, line: str) -> None:
1205        """Removes all project resources, engine-managed objects, state tables and clears the SQLMesh cache."""
1206        context.destroy()
1207
1208
1209def register_magics() -> None:
1210    try:
1211        shell = get_ipython()  # type: ignore
1212        shell.register_magics(SQLMeshMagics)
1213    except NameError:
1214        pass