sqlmesh.magics

  1from __future__ import annotations
  2
  3import functools
  4import logging
  5import typing as t
  6from collections import defaultdict
  7
  8from hyperscript import h
  9from IPython.core.display import display
 10from IPython.core.magic import (
 11    Magics,
 12    cell_magic,
 13    line_cell_magic,
 14    line_magic,
 15    magics_class,
 16)
 17from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
 18from rich.jupyter import JupyterRenderable
 19
 20from sqlmesh.cli.example_project import ProjectTemplate, init_example_project
 21from sqlmesh.core import constants as c
 22from sqlmesh.core.config import load_configs
 23from sqlmesh.core.console import get_console
 24from sqlmesh.core.context import Context
 25from sqlmesh.core.dialect import format_model_expressions, parse
 26from sqlmesh.core.model import load_sql_based_model
 27from sqlmesh.core.test import ModelTestMetadata, get_all_model_tests
 28from sqlmesh.utils import sqlglot_dialects, yaml
 29from sqlmesh.utils.errors import MagicError, MissingContextException, SQLMeshError
 30
 31logger = logging.getLogger(__name__)
 32
 33CONTEXT_VARIABLE_NAMES = [
 34    "context",
 35    "ctx",
 36    "sqlmesh",
 37]
 38
 39
 40def pass_sqlmesh_context(func: t.Callable) -> t.Callable:
 41    @functools.wraps(func)
 42    def wrapper(self: SQLMeshMagics, *args: t.Any, **kwargs: t.Any) -> None:
 43        for variable_name in CONTEXT_VARIABLE_NAMES:
 44            context = self._shell.user_ns.get(variable_name)
 45            if isinstance(context, Context):
 46                break
 47        else:
 48            raise MissingContextException(
 49                f"Context must be defined and initialized with one of these names: {', '.join(CONTEXT_VARIABLE_NAMES)}"
 50            )
 51        old_console = context.console
 52        context.console = get_console(display=self.display)
 53        context.refresh()
 54        func(self, context, *args, **kwargs)
 55        context.console = old_console
 56
 57    return wrapper
 58
 59
 60@magics_class
 61class SQLMeshMagics(Magics):
 62    @property
 63    def display(self) -> t.Callable:
 64        from sqlmesh import RuntimeEnv
 65
 66        if RuntimeEnv.get().is_databricks:
 67            # Use Databricks' special display instead of the normal IPython display
 68            return self._shell.user_ns["display"]
 69        return display
 70
 71    @property
 72    def _shell(self) -> t.Any:
 73        # Make mypy happy.
 74        if not self.shell:
 75            raise RuntimeError("IPython Magics are in invalid state")
 76        return self.shell
 77
 78    @magic_arguments()
 79    @argument(
 80        "paths",
 81        type=str,
 82        nargs="+",
 83        default="",
 84        help="The path(s) to the SQLMesh project(s).",
 85    )
 86    @argument(
 87        "--config",
 88        type=str,
 89        help="Name of the config object. Only applicable to configuration defined using Python script.",
 90    )
 91    @argument("--gateway", type=str, help="The name of the gateway.")
 92    @argument("--ignore-warnings", action="store_true", help="Ignore warnings.")
 93    @argument("--debug", action="store_true", help="Enable debug mode.")
 94    @argument("--log-file-dir", type=str, help="The directory to write the log file to.")
 95    @line_magic
 96    def context(self, line: str) -> None:
 97        """Sets the context in the user namespace."""
 98        from sqlmesh import configure_logging
 99
100        args = parse_argstring(self.context, line)
101        configs = load_configs(args.config, Context.CONFIG_TYPE, args.paths)
102        log_limit = list(configs.values())[0].log_limit
103        configure_logging(
104            args.debug, args.ignore_warnings, log_limit=log_limit, log_file_dir=args.log_file_dir
105        )
106        try:
107            context = Context(paths=args.paths, config=configs, gateway=args.gateway)
108            self._shell.user_ns["context"] = context
109        except Exception:
110            if args.debug:
111                logger.exception("Failed to initialize SQLMesh context")
112            raise
113        context.console.log_success(f"SQLMesh project context set to: {', '.join(args.paths)}")
114
115    @magic_arguments()
116    @argument("path", type=str, help="The path where the new SQLMesh project should be created.")
117    @argument(
118        "sql_dialect",
119        type=str,
120        help=f"Default model SQL dialect. Supported values: {sqlglot_dialects()}.",
121    )
122    @argument(
123        "--template",
124        "-t",
125        type=str,
126        help="Project template. Supported values: airflow, dbt, default, empty.",
127    )
128    @line_magic
129    def init(self, line: str) -> None:
130        """Creates a SQLMesh project scaffold with a default SQL dialect."""
131        args = parse_argstring(self.init, line)
132        try:
133            project_template = ProjectTemplate(
134                args.template.lower() if args.template else "default"
135            )
136        except ValueError:
137            raise MagicError(f"Invalid project template '{args.template}'")
138        init_example_project(args.path, args.sql_dialect, project_template)
139        html = str(
140            h(
141                "div",
142                h(
143                    "span",
144                    {"style": {"color": "green", "font-weight": "bold"}},
145                    "SQLMesh project scaffold created",
146                ),
147            )
148        )
149        self.display(JupyterRenderable(html=html, text=""))
150
151    @magic_arguments()
152    @argument("model", type=str, help="The model.")
153    @argument("--start", "-s", type=str, help="Start date to render.")
154    @argument("--end", "-e", type=str, help="End date to render.")
155    @argument("--execution-time", type=str, help="Execution time.")
156    @argument("--dialect", "-d", type=str, help="The rendered dialect.")
157    @line_cell_magic
158    @pass_sqlmesh_context
159    def model(self, context: Context, line: str, sql: t.Optional[str] = None) -> None:
160        """Renders the model and automatically fills in an editable cell with the model definition."""
161        args = parse_argstring(self.model, line)
162
163        model = context.get_model(args.model, raise_if_missing=True)
164        config = context.config_for_node(model)
165
166        if sql:
167            expressions = parse(sql, default_dialect=config.dialect)
168            loaded = load_sql_based_model(
169                expressions,
170                macros=context._macros,
171                jinja_macros=context._jinja_macros,
172                path=model._path,
173                dialect=config.dialect,
174                time_column_format=config.time_column_format,
175                physical_schema_override=context.config.physical_schema_override,
176                default_catalog=context.default_catalog,
177            )
178
179            if loaded.name == args.model:
180                model = loaded
181        else:
182            with open(model._path, "r", encoding="utf-8") as file:
183                expressions = parse(file.read(), default_dialect=config.dialect)
184
185        formatted = format_model_expressions(
186            expressions, model.dialect, **config.format.generator_options
187        )
188
189        self._shell.set_next_input(
190            "\n".join(
191                [
192                    " ".join(["%%model", line]),
193                    formatted,
194                ]
195            ),
196            replace=True,
197        )
198
199        with open(model._path, "w", encoding="utf-8") as file:
200            file.write(formatted)
201
202        if sql:
203            context.console.log_success(f"Model `{args.model}` updated")
204
205        context.upsert_model(model)
206        context.console.show_sql(
207            context.render(
208                model.name,
209                start=args.start,
210                end=args.end,
211                execution_time=args.execution_time,
212            ).sql(pretty=True, dialect=args.dialect or model.dialect)
213        )
214
215    @magic_arguments()
216    @argument("model", type=str, help="The model.")
217    @argument("test_name", type=str, nargs="?", default=None, help="The test name to display")
218    @argument("--ls", action="store_true", help="List tests associated with a model")
219    @line_cell_magic
220    @pass_sqlmesh_context
221    def test(self, context: Context, line: str, test_def_raw: t.Optional[str] = None) -> None:
222        """Allow the user to list tests for a model, output a specific test, and then write their changes back"""
223        args = parse_argstring(self.test, line)
224        if not args.test_name and not args.ls:
225            raise MagicError("Must provide either test name or `--ls` to list tests")
226
227        test_meta = []
228
229        for path, config in context.configs.items():
230            test_meta.extend(
231                get_all_model_tests(
232                    path / c.TESTS,
233                    ignore_patterns=config.ignore_patterns,
234                )
235            )
236
237        tests: t.Dict[str, t.Dict[str, ModelTestMetadata]] = defaultdict(dict)
238        for model_test_metadata in test_meta:
239            model = model_test_metadata.body.get("model")
240            if not model:
241                context.console.log_error(
242                    f"Test found that does not have `model` defined: {model_test_metadata.path}"
243                )
244            else:
245                tests[model][model_test_metadata.test_name] = model_test_metadata
246
247        model = context.get_model(args.model, raise_if_missing=True)
248
249        if args.ls:
250            # TODO: Provide better UI for displaying tests
251            for test_name in tests[model.name]:
252                context.console.log_status_update(test_name)
253            return
254
255        test = tests[model.name][args.test_name]
256        test_def = yaml.load(test_def_raw) if test_def_raw else test.body
257        test_def_output = yaml.dump(test_def)
258
259        self._shell.set_next_input(
260            "\n".join(
261                [
262                    " ".join(["%%test", line]),
263                    test_def_output,
264                ]
265            ),
266            replace=True,
267        )
268
269        with open(test.path, "r+", encoding="utf-8") as file:
270            content = yaml.load(file.read())
271            content[args.test_name] = test_def
272            file.seek(0)
273            yaml.dump(content, file)
274            file.truncate()
275
276    @magic_arguments()
277    @argument(
278        "environment",
279        nargs="?",
280        type=str,
281        help="The environment to run the plan against",
282    )
283    @argument("--start", "-s", type=str, help="Start date to backfill.")
284    @argument("--end", "-e", type=str, help="End date to backfill.")
285    @argument("--execution-time", type=str, help="Execution time.")
286    @argument(
287        "--create-from",
288        type=str,
289        help="The environment to create the target environment from if it doesn't exist. Default: prod.",
290    )
291    @argument(
292        "--skip-tests",
293        "-t",
294        action="store_true",
295        help="Skip the unit tests defined for the model.",
296    )
297    @argument(
298        "--restate-model",
299        "-r",
300        type=str,
301        nargs="*",
302        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.",
303    )
304    @argument(
305        "--no-gaps",
306        "-g",
307        action="store_true",
308        help="Ensure that new snapshots have no data gaps when comparing to existing snapshots for matching models in the target environment.",
309    )
310    @argument(
311        "--skip-backfill",
312        action="store_true",
313        help="Skip the backfill step.",
314    )
315    @argument(
316        "--forward-only",
317        action="store_true",
318        help="Create a plan for forward-only changes.",
319        default=None,
320    )
321    @argument(
322        "--effective-from",
323        type=str,
324        help="The effective date from which to apply forward-only changes on production.",
325    )
326    @argument(
327        "--no-prompts",
328        action="store_true",
329        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.",
330        default=None,
331    )
332    @argument(
333        "--auto-apply",
334        action="store_true",
335        help="Automatically applies the new plan after creation.",
336        default=None,
337    )
338    @argument(
339        "--no-auto-categorization",
340        action="store_true",
341        help="Disable automatic change categorization.",
342        default=None,
343    )
344    @argument(
345        "--include-unmodified",
346        action="store_true",
347        help="Include unmodified models in the target environment.",
348        default=None,
349    )
350    @argument(
351        "--select-model",
352        type=str,
353        nargs="*",
354        help="Select specific model changes that should be included in the plan.",
355    )
356    @argument(
357        "--backfill-model",
358        type=str,
359        nargs="*",
360        help="Backfill only the models whose names match the expression. This is supported only when targeting a development environment.",
361    )
362    @argument(
363        "--no-diff",
364        action="store_true",
365        help="Hide text differences for changed models.",
366        default=None,
367    )
368    @argument(
369        "--run",
370        action="store_true",
371        help="Run latest intervals as part of the plan application (prod environment only).",
372    )
373    @argument(
374        "--enable-preview",
375        action="store_true",
376        help="Enable preview for forward-only models when targeting a development environment.",
377        default=None,
378    )
379    @line_magic
380    @pass_sqlmesh_context
381    def plan(self, context: Context, line: str) -> None:
382        """Goes through a set of prompts to both establish a plan and apply it"""
383        args = parse_argstring(self.plan, line)
384
385        context.plan(
386            args.environment,
387            start=args.start,
388            end=args.end,
389            execution_time=args.execution_time,
390            create_from=args.create_from,
391            skip_tests=args.skip_tests,
392            restate_models=args.restate_model,
393            backfill_models=args.backfill_model,
394            no_gaps=args.no_gaps,
395            skip_backfill=args.skip_backfill,
396            forward_only=args.forward_only,
397            no_prompts=args.no_prompts,
398            auto_apply=args.auto_apply,
399            no_auto_categorization=args.no_auto_categorization,
400            effective_from=args.effective_from,
401            include_unmodified=args.include_unmodified,
402            select_models=args.select_model,
403            no_diff=args.no_diff,
404            run=args.run,
405            enable_preview=args.enable_preview,
406        )
407
408    @magic_arguments()
409    @argument(
410        "environment",
411        nargs="?",
412        type=str,
413        help="The environment to run against",
414    )
415    @argument("--start", "-s", type=str, help="Start date to evaluate.")
416    @argument("--end", "-e", type=str, help="End date to evaluate.")
417    @argument("--skip-janitor", action="store_true", help="Skip the janitor task.")
418    @argument(
419        "--ignore-cron",
420        action="store_true",
421        help="Run for all missing intervals, ignoring individual cron schedules.",
422    )
423    @line_magic
424    @pass_sqlmesh_context
425    def run_dag(self, context: Context, line: str) -> None:
426        """Evaluate the DAG of models using the built-in scheduler."""
427        args = parse_argstring(self.run_dag, line)
428
429        success = context.run(
430            args.environment,
431            start=args.start,
432            end=args.end,
433            skip_janitor=args.skip_janitor,
434            ignore_cron=args.ignore_cron,
435        )
436        if not success:
437            raise SQLMeshError("Error Running DAG. Check logs for details.")
438
439    @magic_arguments()
440    @argument("model", type=str, help="The model.")
441    @argument("--start", "-s", type=str, help="Start date to render.")
442    @argument("--end", "-e", type=str, help="End date to render.")
443    @argument("--execution-time", type=str, help="Execution time.")
444    @argument(
445        "--limit",
446        type=int,
447        help="The number of rows which the query should be limited to.",
448    )
449    @line_magic
450    @pass_sqlmesh_context
451    def evaluate(self, context: Context, line: str) -> None:
452        """Evaluate a model query and fetches a dataframe."""
453        context.refresh()
454        args = parse_argstring(self.evaluate, line)
455
456        df = context.evaluate(
457            args.model,
458            start=args.start,
459            end=args.end,
460            execution_time=args.execution_time,
461            limit=args.limit,
462        )
463        self.display(df)
464
465    @magic_arguments()
466    @argument("model", type=str, help="The model.")
467    @argument("--start", "-s", type=str, help="Start date to render.")
468    @argument("--end", "-e", type=str, help="End date to render.")
469    @argument("--execution-time", type=str, help="Execution time.")
470    @argument(
471        "--expand",
472        type=t.Union[bool, t.Iterable[str]],
473        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.",
474    )
475    @argument("--dialect", type=str, help="SQL dialect to render.")
476    @argument("--no-format", action="store_true", help="Disable fancy formatting of the query.")
477    @line_magic
478    @pass_sqlmesh_context
479    def render(self, context: Context, line: str) -> None:
480        """Renders a model's query, optionally expanding referenced models."""
481        context.refresh()
482        args = parse_argstring(self.render, line)
483
484        query = context.render(
485            args.model,
486            start=args.start,
487            end=args.end,
488            execution_time=args.execution_time,
489            expand=args.expand,
490        )
491
492        sql = query.sql(pretty=True, dialect=args.dialect or context.config.dialect)
493        if args.no_format:
494            context.console.log_status_update(sql)
495        else:
496            context.console.show_sql(sql)
497
498    @magic_arguments()
499    @argument(
500        "df_var",
501        default=None,
502        nargs="?",
503        type=str,
504        help="An optional variable name to store the resulting dataframe.",
505    )
506    @cell_magic
507    @pass_sqlmesh_context
508    def fetchdf(self, context: Context, line: str, sql: str) -> None:
509        """Fetches a dataframe from sql, optionally storing it in a variable."""
510        args = parse_argstring(self.fetchdf, line)
511        df = context.fetchdf(sql)
512        if args.df_var:
513            self._shell.user_ns[args.df_var] = df
514        self.display(df)
515
516    @magic_arguments()
517    @argument("--file", "-f", type=str, help="An optional file path to write the HTML output to.")
518    @argument(
519        "--select-model",
520        type=str,
521        nargs="*",
522        help="Select specific models to include in the dag.",
523    )
524    @line_magic
525    @pass_sqlmesh_context
526    def dag(self, context: Context, line: str) -> None:
527        """Displays the HTML DAG."""
528        args = parse_argstring(self.dag, line)
529        dag = context.get_dag(args.select_model)
530        if args.file:
531            with open(args.file, "w", encoding="utf-8") as file:
532                file.write(str(dag))
533        # TODO: Have this go through console instead of calling display directly
534        self.display(dag)
535
536    @magic_arguments()
537    @line_magic
538    @pass_sqlmesh_context
539    def migrate(self, context: Context, line: str) -> None:
540        """Migrate SQLMesh to the current running version."""
541        context.migrate()
542        context.console.log_success("Migration complete")
543
544    @magic_arguments()
545    @line_magic
546    @pass_sqlmesh_context
547    def create_external_models(self, context: Context, line: str) -> None:
548        """Create a schema file containing external model schemas."""
549        context.create_external_models()
550
551    @magic_arguments()
552    @argument(
553        "source_to_target",
554        type=str,
555        metavar="SOURCE:TARGET",
556        help="Source and target in `SOURCE:TARGET` format",
557    )
558    @argument(
559        "--on",
560        type=str,
561        nargs="*",
562        help="The column to join on. Can be specified multiple times. The model grain will be used if not specified.",
563    )
564    @argument(
565        "--model",
566        type=str,
567        help="The model to diff against when source and target are environments and not tables.",
568    )
569    @argument(
570        "--where",
571        type=str,
572        help="An optional where statement to filter results.",
573    )
574    @argument(
575        "--limit",
576        type=int,
577        default=20,
578        help="The limit of the sample dataframe.",
579    )
580    @argument(
581        "--show-sample",
582        action="store_true",
583        help="Show a sample of the rows that differ. With many columns, the output can be very wide.",
584    )
585    @line_magic
586    @pass_sqlmesh_context
587    def table_diff(self, context: Context, line: str) -> None:
588        """Show the diff between two tables.
589
590        Can either be two tables or two environments and a model.
591        """
592        args = parse_argstring(self.table_diff, line)
593        source, target = args.source_to_target.split(":")
594        context.table_diff(
595            source=source,
596            target=target,
597            on=args.on,
598            model_or_snapshot=args.model,
599            where=args.where,
600            limit=args.limit,
601            show_sample=args.show_sample,
602        )
603
604    @magic_arguments()
605    @argument(
606        "model_name",
607        nargs="?",
608        type=str,
609        help="The name of the model to get the table name for.",
610    )
611    @argument(
612        "--dev",
613        action="store_true",
614        help="Print the name of the snapshot table used for previews in development environments.",
615    )
616    @line_magic
617    @pass_sqlmesh_context
618    def table_name(self, context: Context, line: str) -> None:
619        """Prints the name of the physical table for the given model."""
620        args = parse_argstring(self.table_name, line)
621        context.console.log_status_update(context.table_name(args.model_name, args.dev))
622
623    @magic_arguments()
624    @argument(
625        "--read",
626        type=str,
627        default="",
628        help="The input dialect of the sql string.",
629    )
630    @argument(
631        "--write",
632        type=str,
633        default="",
634        help="The output dialect of the sql string.",
635    )
636    @line_cell_magic
637    @pass_sqlmesh_context
638    def rewrite(self, context: Context, line: str, sql: str) -> None:
639        """Rewrite a sql expression with semantic references into an executable query.
640
641        https://sqlmesh.readthedocs.io/en/latest/concepts/metrics/overview/
642        """
643        args = parse_argstring(self.rewrite, line)
644        context.console.show_sql(
645            context.rewrite(sql, args.read).sql(
646                dialect=args.write or context.config.dialect, pretty=True
647            )
648        )
649
650    @magic_arguments()
651    @argument(
652        "--transpile",
653        "-t",
654        type=str,
655        help="Transpile project models to the specified dialect.",
656    )
657    @argument(
658        "--append-newline",
659        action="store_true",
660        help="Whether or not to append a newline to the end of the file.",
661        default=None,
662    )
663    @argument(
664        "--normalize",
665        action="store_true",
666        help="Whether or not to normalize identifiers to lowercase.",
667        default=None,
668    )
669    @argument(
670        "--pad",
671        type=int,
672        help="Determines the pad size in a formatted string.",
673    )
674    @argument(
675        "--indent",
676        type=int,
677        help="Determines the indentation size in a formatted string.",
678    )
679    @argument(
680        "--normalize-functions",
681        type=str,
682        help="Whether or not to normalize all function names. Possible values are: 'upper', 'lower'",
683    )
684    @argument(
685        "--leading-comma",
686        action="store_true",
687        help="Determines whether or not the comma is leading or trailing in select expressions. Default is trailing.",
688        default=None,
689    )
690    @argument(
691        "--max-text-width",
692        type=int,
693        help="The max number of characters in a segment before creating new lines in pretty mode.",
694    )
695    @line_magic
696    @pass_sqlmesh_context
697    def format(self, context: Context, line: str) -> None:
698        """Format all SQL models."""
699        args = parse_argstring(self.format, line)
700        context.format(**{k: v for k, v in vars(args).items() if v is not None})
701
702    @magic_arguments()
703    @argument("environment", type=str, help="The environment to diff local state against.")
704    @line_magic
705    @pass_sqlmesh_context
706    def diff(self, context: Context, line: str) -> None:
707        """Show the diff between the local state and the target environment."""
708        args = parse_argstring(self.diff, line)
709        context.diff(args.environment)
710
711    @magic_arguments()
712    @argument("environment", type=str, help="The environment to invalidate.")
713    @line_magic
714    @pass_sqlmesh_context
715    def invalidate(self, context: Context, line: str) -> None:
716        """Invalidate the target environment, forcing its removal during the next run of the janitor process."""
717        args = parse_argstring(self.invalidate, line)
718        context.invalidate_environment(args.environment)
719
720    @magic_arguments()
721    @argument("model", type=str)
722    @argument(
723        "--query",
724        "-q",
725        type=str,
726        nargs="+",
727        required=True,
728        help="Queries that will be used to generate data for the model's dependencies.",
729    )
730    @argument(
731        "--overwrite",
732        "-o",
733        action="store_true",
734        help="When true, the fixture file will be overwritten in case it already exists.",
735    )
736    @argument(
737        "--var",
738        "-v",
739        type=str,
740        nargs="+",
741        help="Key-value pairs that will define variables needed by the model.",
742    )
743    @argument(
744        "--path",
745        "-p",
746        type=str,
747        help="The file path corresponding to the fixture, relative to the test directory. "
748        "By default, the fixture will be created under the test directory and the file "
749        "name will be inferred based on the test's name.",
750    )
751    @argument(
752        "--name",
753        "-n",
754        type=str,
755        help="The name of the test that will be created. By default, it's inferred based on the model's name.",
756    )
757    @argument(
758        "--include-ctes",
759        action="store_true",
760        help="When true, CTE fixtures will also be generated.",
761    )
762    @line_magic
763    @pass_sqlmesh_context
764    def create_test(self, context: Context, line: str) -> None:
765        """Generate a unit test fixture for a given model."""
766        args = parse_argstring(self.create_test, line)
767        queries = iter(args.query)
768        variables = iter(args.var) if args.var else None
769        context.create_test(
770            args.model,
771            input_queries={k: v.strip('"') for k, v in dict(zip(queries, queries)).items()},
772            overwrite=args.overwrite,
773            variables=dict(zip(variables, variables)) if variables else None,
774            path=args.path,
775            name=args.name,
776            include_ctes=args.include_ctes,
777        )
778
779    @magic_arguments()
780    @argument("tests", nargs="*", type=str)
781    @argument(
782        "--pattern",
783        "-k",
784        nargs="*",
785        type=str,
786        help="Only run tests that match the pattern of substring.",
787    )
788    @argument("--verbose", "-v", action="store_true", help="Verbose output.")
789    @argument(
790        "--preserve-fixtures",
791        action="store_true",
792        help="Preserve the fixture tables in the testing database, useful for debugging.",
793    )
794    @line_magic
795    @pass_sqlmesh_context
796    def run_test(self, context: Context, line: str) -> None:
797        """Run unit test(s)."""
798        args = parse_argstring(self.run_test, line)
799        context.test(
800            match_patterns=args.pattern,
801            tests=args.tests,
802            verbose=args.verbose,
803            preserve_fixtures=args.preserve_fixtures,
804        )
805
806    @magic_arguments()
807    @argument(
808        "models", type=str, nargs="*", help="A model to audit. Multiple models can be audited."
809    )
810    @argument("--start", "-s", type=str, help="Start date to audit.")
811    @argument("--end", "-e", type=str, help="End date to audit.")
812    @argument("--execution-time", type=str, help="Execution time.")
813    @line_magic
814    @pass_sqlmesh_context
815    def audit(self, context: Context, line: str) -> None:
816        """Run audit(s)"""
817        args = parse_argstring(self.audit, line)
818        context.audit(
819            models=args.models, start=args.start, end=args.end, execution_time=args.execution_time
820        )
821
822    @magic_arguments()
823    @line_magic
824    @pass_sqlmesh_context
825    def info(self, context: Context, line: str) -> None:
826        """Display SQLMesh project information."""
827        context.print_info()
828
829    @magic_arguments()
830    @line_magic
831    @pass_sqlmesh_context
832    def rollback(self, context: Context, line: str) -> None:
833        """Rollback SQLMesh to the previous migration."""
834        context.rollback()
835
836    @magic_arguments()
837    @line_magic
838    @pass_sqlmesh_context
839    def clean(self, context: Context, line: str) -> None:
840        """Clears the SQLMesh cache and any build artifacts."""
841        context.clear_caches()
842        context.console.log_success("SQLMesh cache and build artifacts cleared")
843
844
845def register_magics() -> None:
846    try:
847        shell = get_ipython()  # type: ignore
848        shell.register_magics(SQLMeshMagics)
849    except NameError:
850        pass