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