Edit on GitHub

sqlmesh.magics

  1from __future__ import annotations
  2
  3import typing as t
  4from collections import defaultdict
  5
  6from IPython.core.display import HTML, display
  7from IPython.core.magic import Magics, line_cell_magic, line_magic, magics_class
  8from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
  9from ruamel.yaml import YAML
 10
 11from sqlmesh.core.console import get_console
 12from sqlmesh.core.context import Context
 13from sqlmesh.core.dialect import format_model_expressions, parse
 14from sqlmesh.core.model import load_model
 15from sqlmesh.core.test import ModelTestMetadata, get_all_model_tests
 16from sqlmesh.utils.errors import MagicError, MissingContextException, SQLMeshError
 17from sqlmesh.utils.yaml import dumps as yaml_dumps
 18from sqlmesh.utils.yaml import load as yaml_load
 19
 20CONTEXT_VARIABLE_NAMES = [
 21    "context",
 22    "ctx",
 23    "sqlmesh",
 24]
 25
 26
 27@magics_class
 28class SQLMeshMagics(Magics):
 29    @property
 30    def display(self) -> t.Callable:
 31        from sqlmesh import runtime_env
 32
 33        if runtime_env.is_databricks:
 34            # Use Databrick's special display instead of the normal IPython display
 35            return self._shell.user_ns["display"]
 36        return display
 37
 38    @property
 39    def _context(self) -> Context:
 40        for variable_name in CONTEXT_VARIABLE_NAMES:
 41            context = self._shell.user_ns.get(variable_name)
 42            if context:
 43                return context
 44        raise MissingContextException(
 45            f"Context must be defined and initialized with one of these names: {', '.join(CONTEXT_VARIABLE_NAMES)}"
 46        )
 47
 48    @magic_arguments()
 49    @argument("path", type=str, help="The path to the SQLMesh project.")
 50    @line_magic
 51    def context(self, path: str) -> None:
 52        """Sets the context in the user namespace."""
 53        self._shell.user_ns["context"] = Context(path=path)
 54
 55    @magic_arguments()
 56    @argument("model", type=str, help="The model.")
 57    @argument("--start", "-s", type=str, help="Start date to render.")
 58    @argument("--end", "-e", type=str, help="End date to render.")
 59    @argument("--latest", "-l", type=str, help="Latest date to render.")
 60    @argument("--dialect", "-d", type=str, help="The rendered dialect.")
 61    @line_cell_magic
 62    def model(self, line: str, sql: t.Optional[str] = None) -> None:
 63        """Renders the model and automatically fills in an editable cell with the model definition."""
 64        args = parse_argstring(self.model, line)
 65        model = self._context.get_model(args.model)
 66
 67        if not model:
 68            raise SQLMeshError(f"Cannot find {model}")
 69
 70        if sql:
 71            loaded = load_model(
 72                parse(sql, default_dialect=self._context.dialect),
 73                macros=self._context._macros,
 74                hooks=self._context._hooks,
 75                path=model._path,
 76                dialect=self._context.dialect,
 77                time_column_format=self._context.config.time_column_format,
 78            )
 79
 80            if loaded.name == args.model:
 81                model = loaded
 82
 83        self._context.upsert_model(model)
 84        expressions = model.render_definition(include_python=False)
 85
 86        formatted = format_model_expressions(expressions, model.dialect)
 87
 88        self._shell.set_next_input(
 89            "\n".join(
 90                [
 91                    " ".join(["%%model", line]),
 92                    formatted,
 93                ]
 94            ),
 95            replace=True,
 96        )
 97
 98        with open(model._path, "w", encoding="utf-8") as file:
 99            file.write(formatted)
100
101        self._context.upsert_model(model)
102        self._context.console.show_sql(
103            self._context.render(
104                model.name,
105                start=args.start,
106                end=args.end,
107                latest=args.latest,
108            ).sql(pretty=True, dialect=args.dialect or model.dialect)
109        )
110
111    @magic_arguments()
112    @argument("model", type=str, help="The model.")
113    @argument("test_name", type=str, nargs="?", default=None, help="The test name to display")
114    @argument("--ls", action="store_true", help="List tests associated with a model")
115    @line_cell_magic
116    def test(self, line: str, test_def_raw: t.Optional[str] = None) -> None:
117        """Allow the user to list tests for a model, output a specific test, and then write their changes back"""
118        args = parse_argstring(self.test, line)
119        if not args.test_name and not args.ls:
120            raise MagicError("Must provide either test name or `--ls` to list tests")
121
122        model_test_metadatas = get_all_model_tests(
123            self._context.test_directory_path,
124            ignore_patterns=self._context.ignore_patterns,
125        )
126        tests: t.Dict[str, t.Dict[str, ModelTestMetadata]] = defaultdict(dict)
127        for model_test_metadata in model_test_metadatas:
128            model = model_test_metadata.body.get("model")
129            if not model:
130                self._context.console.log_error(
131                    f"Test found that does not have `model` defined: {model_test_metadata.path}"
132                )
133            tests[model][model_test_metadata.test_name] = model_test_metadata
134        if args.ls:
135            # TODO: Provide better UI for displaying tests
136            for test_name in tests[args.model]:
137                self._context.console.log_status_update(test_name)
138            return
139
140        test = tests[args.model][args.test_name]
141        test_def = yaml_load(test_def_raw) if test_def_raw else test.body
142        test_def_output = yaml_dumps(test_def)
143
144        self._shell.set_next_input(
145            "\n".join(
146                [
147                    " ".join(["%%test", line]),
148                    test_def_output,
149                ]
150            ),
151            replace=True,
152        )
153
154        with open(test.path, "r+", encoding="utf-8") as file:
155            content = yaml_load(file.read())
156            content[args.test_name] = test_def
157            file.seek(0)
158            YAML().dump(content, file)
159            file.truncate()
160
161    @magic_arguments()
162    @argument(
163        "environment",
164        nargs="?",
165        type=str,
166        help="The environment to run the plan against",
167    )
168    @argument("--start", "-s", type=str, help="Start date to backfill.")
169    @argument("--end", "-e", type=str, help="End date to backfill.")
170    @argument(
171        "--create-from",
172        type=str,
173        help="The environment to create the target environment from if it doesn't exist. Default: prod.",
174    )
175    @argument(
176        "--skip-tests",
177        "-t",
178        action="store_true",
179        help="Skip the unit tests defined for the model.",
180    )
181    @argument(
182        "--restate-model",
183        "-r",
184        type=str,
185        nargs="*",
186        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.",
187    )
188    @argument(
189        "--no-gaps",
190        "-g",
191        action="store_true",
192        help="Ensure that new snapshots have no data gaps when comparing to existing snapshots for matching models in the target environment.",
193    )
194    @argument(
195        "--skip-backfill",
196        action="store_true",
197        help="Skip the backfill step.",
198    )
199    @argument(
200        "--forward-only",
201        action="store_true",
202        help="Create a plan for forward-only changes.",
203    )
204    @argument(
205        "--no-prompts",
206        action="store_true",
207        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.",
208    )
209    @argument(
210        "--auto-apply",
211        action="store_true",
212        help="Automatically applies the new plan after creation.",
213    )
214    @argument(
215        "--no-auto-categorization",
216        action="store_true",
217        help="Disable automatic change categorization.",
218        default=None,
219    )
220    @line_magic
221    def plan(self, line: str) -> None:
222        """Goes through a set of prompts to both establish a plan and apply it"""
223        self._context.refresh()
224        args = parse_argstring(self.plan, line)
225
226        # Since the magics share a context we want to clear out any state before generating a new plan
227        console = self._context.console
228        self._context.console = get_console(display=self.display)
229
230        self._context.plan(
231            args.environment,
232            start=args.start,
233            end=args.end,
234            create_from=args.create_from,
235            skip_tests=args.skip_tests,
236            restate_models=args.restate_model,
237            no_gaps=args.no_gaps,
238            skip_backfill=args.skip_backfill,
239            forward_only=args.forward_only,
240            no_prompts=args.no_prompts,
241            auto_apply=args.auto_apply,
242            no_auto_categorization=args.no_auto_categorization,
243        )
244        self._context.console = console
245
246    @magic_arguments()
247    @argument(
248        "environment",
249        nargs="?",
250        type=str,
251        help="The environment to run against",
252    )
253    @argument("--start", "-s", type=str, help="Start date to evaluate.")
254    @argument("--end", "-e", type=str, help="End date to evaluate.")
255    @argument("--skip-janitor", action="store_true", help="Skip the jantitor task.")
256    @line_magic
257    def run_dag(self, line: str) -> None:
258        """Evaluate the DAG of models using the built-in scheduler."""
259        args = parse_argstring(self.run_dag, line)
260
261        # Since the magics share a context we want to clear out any state before generating a new plan
262        console = self._context.console
263        self._context.console = get_console(display=self.display)
264
265        self._context.run(
266            args.environment,
267            start=args.start,
268            end=args.end,
269            skip_janitor=args.skip_janitor,
270        )
271        self._context.console = console
272
273    @magic_arguments()
274    @argument("model", type=str, help="The model.")
275    @argument("--start", "-s", type=str, help="Start date to render.")
276    @argument("--end", "-e", type=str, help="End date to render.")
277    @argument("--latest", "-l", type=str, help="Latest date to render.")
278    @argument(
279        "--limit",
280        type=int,
281        help="The number of rows which the query should be limited to.",
282    )
283    @line_magic
284    def evaluate(self, line: str) -> None:
285        """Evaluate a model query and fetches a dataframe."""
286        self._context.refresh()
287        args = parse_argstring(self.evaluate, line)
288
289        df = self._context.evaluate(
290            args.model,
291            start=args.start,
292            end=args.end,
293            latest=args.latest,
294            limit=args.limit,
295        )
296        self.display(df)
297
298    @magic_arguments()
299    @argument("model", type=str, help="The model.")
300    @argument("--start", "-s", type=str, help="Start date to render.")
301    @argument("--end", "-e", type=str, help="End date to render.")
302    @argument("--latest", "-l", type=str, help="Latest date to render.")
303    @argument(
304        "--expand",
305        type=t.Union[bool, t.Iterable[str]],
306        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.",
307    )
308    @argument("--dialect", type=str, help="SQL dialect to render.")
309    @line_magic
310    def render(self, line: str) -> None:
311        """Renders a model's query, optionally expanding referenced models."""
312        self._context.refresh()
313        args = parse_argstring(self.render, line)
314
315        query = self._context.render(
316            args.model,
317            start=args.start,
318            end=args.end,
319            latest=args.latest,
320            expand=args.expand,
321        )
322
323        self._context.console.show_sql(query.sql(pretty=True, dialect=args.dialect))
324
325    @magic_arguments()
326    @argument(
327        "df_var",
328        default=None,
329        nargs="?",
330        type=str,
331        help="An optional variable name to store the resulting dataframe.",
332    )
333    @line_cell_magic
334    def fetchdf(self, line: str, sql: str) -> None:
335        """Fetches a dataframe from sql, optionally storing it in a variable."""
336        args = parse_argstring(self.fetchdf, line)
337        df = self._context.fetchdf(sql)
338        if args.df_var:
339            self._shell.user_ns[args.df_var] = df
340        self.display(df)
341
342    @magic_arguments()
343    @line_magic
344    def dag(self, line: str) -> None:
345        """Displays the dag"""
346        self._context.refresh()
347        dag = self._context.get_dag()
348        self.display(HTML(dag.pipe().decode("utf-8")))
349
350    @property
351    def _shell(self) -> t.Any:
352        # Make mypy happy.
353        if not self.shell:
354            raise RuntimeError("IPython Magics are in invalid state")
355        return self.shell
356
357
358def register_magics() -> None:
359    try:
360        shell = get_ipython()  # type: ignore
361        shell.register_magics(SQLMeshMagics)
362    except NameError:
363        pass
@magics_class
class SQLMeshMagics(IPython.core.magic.Magics):
 28@magics_class
 29class SQLMeshMagics(Magics):
 30    @property
 31    def display(self) -> t.Callable:
 32        from sqlmesh import runtime_env
 33
 34        if runtime_env.is_databricks:
 35            # Use Databrick's special display instead of the normal IPython display
 36            return self._shell.user_ns["display"]
 37        return display
 38
 39    @property
 40    def _context(self) -> Context:
 41        for variable_name in CONTEXT_VARIABLE_NAMES:
 42            context = self._shell.user_ns.get(variable_name)
 43            if context:
 44                return context
 45        raise MissingContextException(
 46            f"Context must be defined and initialized with one of these names: {', '.join(CONTEXT_VARIABLE_NAMES)}"
 47        )
 48
 49    @magic_arguments()
 50    @argument("path", type=str, help="The path to the SQLMesh project.")
 51    @line_magic
 52    def context(self, path: str) -> None:
 53        """Sets the context in the user namespace."""
 54        self._shell.user_ns["context"] = Context(path=path)
 55
 56    @magic_arguments()
 57    @argument("model", type=str, help="The model.")
 58    @argument("--start", "-s", type=str, help="Start date to render.")
 59    @argument("--end", "-e", type=str, help="End date to render.")
 60    @argument("--latest", "-l", type=str, help="Latest date to render.")
 61    @argument("--dialect", "-d", type=str, help="The rendered dialect.")
 62    @line_cell_magic
 63    def model(self, line: str, sql: t.Optional[str] = None) -> None:
 64        """Renders the model and automatically fills in an editable cell with the model definition."""
 65        args = parse_argstring(self.model, line)
 66        model = self._context.get_model(args.model)
 67
 68        if not model:
 69            raise SQLMeshError(f"Cannot find {model}")
 70
 71        if sql:
 72            loaded = load_model(
 73                parse(sql, default_dialect=self._context.dialect),
 74                macros=self._context._macros,
 75                hooks=self._context._hooks,
 76                path=model._path,
 77                dialect=self._context.dialect,
 78                time_column_format=self._context.config.time_column_format,
 79            )
 80
 81            if loaded.name == args.model:
 82                model = loaded
 83
 84        self._context.upsert_model(model)
 85        expressions = model.render_definition(include_python=False)
 86
 87        formatted = format_model_expressions(expressions, model.dialect)
 88
 89        self._shell.set_next_input(
 90            "\n".join(
 91                [
 92                    " ".join(["%%model", line]),
 93                    formatted,
 94                ]
 95            ),
 96            replace=True,
 97        )
 98
 99        with open(model._path, "w", encoding="utf-8") as file:
100            file.write(formatted)
101
102        self._context.upsert_model(model)
103        self._context.console.show_sql(
104            self._context.render(
105                model.name,
106                start=args.start,
107                end=args.end,
108                latest=args.latest,
109            ).sql(pretty=True, dialect=args.dialect or model.dialect)
110        )
111
112    @magic_arguments()
113    @argument("model", type=str, help="The model.")
114    @argument("test_name", type=str, nargs="?", default=None, help="The test name to display")
115    @argument("--ls", action="store_true", help="List tests associated with a model")
116    @line_cell_magic
117    def test(self, line: str, test_def_raw: t.Optional[str] = None) -> None:
118        """Allow the user to list tests for a model, output a specific test, and then write their changes back"""
119        args = parse_argstring(self.test, line)
120        if not args.test_name and not args.ls:
121            raise MagicError("Must provide either test name or `--ls` to list tests")
122
123        model_test_metadatas = get_all_model_tests(
124            self._context.test_directory_path,
125            ignore_patterns=self._context.ignore_patterns,
126        )
127        tests: t.Dict[str, t.Dict[str, ModelTestMetadata]] = defaultdict(dict)
128        for model_test_metadata in model_test_metadatas:
129            model = model_test_metadata.body.get("model")
130            if not model:
131                self._context.console.log_error(
132                    f"Test found that does not have `model` defined: {model_test_metadata.path}"
133                )
134            tests[model][model_test_metadata.test_name] = model_test_metadata
135        if args.ls:
136            # TODO: Provide better UI for displaying tests
137            for test_name in tests[args.model]:
138                self._context.console.log_status_update(test_name)
139            return
140
141        test = tests[args.model][args.test_name]
142        test_def = yaml_load(test_def_raw) if test_def_raw else test.body
143        test_def_output = yaml_dumps(test_def)
144
145        self._shell.set_next_input(
146            "\n".join(
147                [
148                    " ".join(["%%test", line]),
149                    test_def_output,
150                ]
151            ),
152            replace=True,
153        )
154
155        with open(test.path, "r+", encoding="utf-8") as file:
156            content = yaml_load(file.read())
157            content[args.test_name] = test_def
158            file.seek(0)
159            YAML().dump(content, file)
160            file.truncate()
161
162    @magic_arguments()
163    @argument(
164        "environment",
165        nargs="?",
166        type=str,
167        help="The environment to run the plan against",
168    )
169    @argument("--start", "-s", type=str, help="Start date to backfill.")
170    @argument("--end", "-e", type=str, help="End date to backfill.")
171    @argument(
172        "--create-from",
173        type=str,
174        help="The environment to create the target environment from if it doesn't exist. Default: prod.",
175    )
176    @argument(
177        "--skip-tests",
178        "-t",
179        action="store_true",
180        help="Skip the unit tests defined for the model.",
181    )
182    @argument(
183        "--restate-model",
184        "-r",
185        type=str,
186        nargs="*",
187        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.",
188    )
189    @argument(
190        "--no-gaps",
191        "-g",
192        action="store_true",
193        help="Ensure that new snapshots have no data gaps when comparing to existing snapshots for matching models in the target environment.",
194    )
195    @argument(
196        "--skip-backfill",
197        action="store_true",
198        help="Skip the backfill step.",
199    )
200    @argument(
201        "--forward-only",
202        action="store_true",
203        help="Create a plan for forward-only changes.",
204    )
205    @argument(
206        "--no-prompts",
207        action="store_true",
208        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.",
209    )
210    @argument(
211        "--auto-apply",
212        action="store_true",
213        help="Automatically applies the new plan after creation.",
214    )
215    @argument(
216        "--no-auto-categorization",
217        action="store_true",
218        help="Disable automatic change categorization.",
219        default=None,
220    )
221    @line_magic
222    def plan(self, line: str) -> None:
223        """Goes through a set of prompts to both establish a plan and apply it"""
224        self._context.refresh()
225        args = parse_argstring(self.plan, line)
226
227        # Since the magics share a context we want to clear out any state before generating a new plan
228        console = self._context.console
229        self._context.console = get_console(display=self.display)
230
231        self._context.plan(
232            args.environment,
233            start=args.start,
234            end=args.end,
235            create_from=args.create_from,
236            skip_tests=args.skip_tests,
237            restate_models=args.restate_model,
238            no_gaps=args.no_gaps,
239            skip_backfill=args.skip_backfill,
240            forward_only=args.forward_only,
241            no_prompts=args.no_prompts,
242            auto_apply=args.auto_apply,
243            no_auto_categorization=args.no_auto_categorization,
244        )
245        self._context.console = console
246
247    @magic_arguments()
248    @argument(
249        "environment",
250        nargs="?",
251        type=str,
252        help="The environment to run against",
253    )
254    @argument("--start", "-s", type=str, help="Start date to evaluate.")
255    @argument("--end", "-e", type=str, help="End date to evaluate.")
256    @argument("--skip-janitor", action="store_true", help="Skip the jantitor task.")
257    @line_magic
258    def run_dag(self, line: str) -> None:
259        """Evaluate the DAG of models using the built-in scheduler."""
260        args = parse_argstring(self.run_dag, line)
261
262        # Since the magics share a context we want to clear out any state before generating a new plan
263        console = self._context.console
264        self._context.console = get_console(display=self.display)
265
266        self._context.run(
267            args.environment,
268            start=args.start,
269            end=args.end,
270            skip_janitor=args.skip_janitor,
271        )
272        self._context.console = console
273
274    @magic_arguments()
275    @argument("model", type=str, help="The model.")
276    @argument("--start", "-s", type=str, help="Start date to render.")
277    @argument("--end", "-e", type=str, help="End date to render.")
278    @argument("--latest", "-l", type=str, help="Latest date to render.")
279    @argument(
280        "--limit",
281        type=int,
282        help="The number of rows which the query should be limited to.",
283    )
284    @line_magic
285    def evaluate(self, line: str) -> None:
286        """Evaluate a model query and fetches a dataframe."""
287        self._context.refresh()
288        args = parse_argstring(self.evaluate, line)
289
290        df = self._context.evaluate(
291            args.model,
292            start=args.start,
293            end=args.end,
294            latest=args.latest,
295            limit=args.limit,
296        )
297        self.display(df)
298
299    @magic_arguments()
300    @argument("model", type=str, help="The model.")
301    @argument("--start", "-s", type=str, help="Start date to render.")
302    @argument("--end", "-e", type=str, help="End date to render.")
303    @argument("--latest", "-l", type=str, help="Latest date to render.")
304    @argument(
305        "--expand",
306        type=t.Union[bool, t.Iterable[str]],
307        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.",
308    )
309    @argument("--dialect", type=str, help="SQL dialect to render.")
310    @line_magic
311    def render(self, line: str) -> None:
312        """Renders a model's query, optionally expanding referenced models."""
313        self._context.refresh()
314        args = parse_argstring(self.render, line)
315
316        query = self._context.render(
317            args.model,
318            start=args.start,
319            end=args.end,
320            latest=args.latest,
321            expand=args.expand,
322        )
323
324        self._context.console.show_sql(query.sql(pretty=True, dialect=args.dialect))
325
326    @magic_arguments()
327    @argument(
328        "df_var",
329        default=None,
330        nargs="?",
331        type=str,
332        help="An optional variable name to store the resulting dataframe.",
333    )
334    @line_cell_magic
335    def fetchdf(self, line: str, sql: str) -> None:
336        """Fetches a dataframe from sql, optionally storing it in a variable."""
337        args = parse_argstring(self.fetchdf, line)
338        df = self._context.fetchdf(sql)
339        if args.df_var:
340            self._shell.user_ns[args.df_var] = df
341        self.display(df)
342
343    @magic_arguments()
344    @line_magic
345    def dag(self, line: str) -> None:
346        """Displays the dag"""
347        self._context.refresh()
348        dag = self._context.get_dag()
349        self.display(HTML(dag.pipe().decode("utf-8")))
350
351    @property
352    def _shell(self) -> t.Any:
353        # Make mypy happy.
354        if not self.shell:
355            raise RuntimeError("IPython Magics are in invalid state")
356        return self.shell

Base class for implementing magic functions.

Shell functions which can be reached as %function_name. All magic functions should accept a string, which they can parse for their own needs. This can make some functions easier to type, eg %cd ../ vs. %cd("../")

Classes providing magic functions need to subclass this class, and they MUST:

  • Use the method decorators @line_magic and @cell_magic to decorate individual methods as magic functions, AND

  • Use the class decorator @magics_class to ensure that the magic methods are properly registered at the instance level upon instance initialization.

See magic_functions for examples of actual implementation classes.

@magic_arguments()
@argument('path', type=str, help='The path to the SQLMesh project.')
@line_magic
def context(self, path: str) -> None:
49    @magic_arguments()
50    @argument("path", type=str, help="The path to the SQLMesh project.")
51    @line_magic
52    def context(self, path: str) -> None:
53        """Sets the context in the user namespace."""
54        self._shell.user_ns["context"] = Context(path=path)

Sets the context in the user namespace.

@magic_arguments()
@argument('model', type=str, help='The model.')
@argument('--start', '-s', type=str, help='Start date to render.')
@argument('--end', '-e', type=str, help='End date to render.')
@argument('--latest', '-l', type=str, help='Latest date to render.')
@argument('--dialect', '-d', type=str, help='The rendered dialect.')
@line_cell_magic
def model(self, line: str, sql: Optional[str] = None) -> None:
 56    @magic_arguments()
 57    @argument("model", type=str, help="The model.")
 58    @argument("--start", "-s", type=str, help="Start date to render.")
 59    @argument("--end", "-e", type=str, help="End date to render.")
 60    @argument("--latest", "-l", type=str, help="Latest date to render.")
 61    @argument("--dialect", "-d", type=str, help="The rendered dialect.")
 62    @line_cell_magic
 63    def model(self, line: str, sql: t.Optional[str] = None) -> None:
 64        """Renders the model and automatically fills in an editable cell with the model definition."""
 65        args = parse_argstring(self.model, line)
 66        model = self._context.get_model(args.model)
 67
 68        if not model:
 69            raise SQLMeshError(f"Cannot find {model}")
 70
 71        if sql:
 72            loaded = load_model(
 73                parse(sql, default_dialect=self._context.dialect),
 74                macros=self._context._macros,
 75                hooks=self._context._hooks,
 76                path=model._path,
 77                dialect=self._context.dialect,
 78                time_column_format=self._context.config.time_column_format,
 79            )
 80
 81            if loaded.name == args.model:
 82                model = loaded
 83
 84        self._context.upsert_model(model)
 85        expressions = model.render_definition(include_python=False)
 86
 87        formatted = format_model_expressions(expressions, model.dialect)
 88
 89        self._shell.set_next_input(
 90            "\n".join(
 91                [
 92                    " ".join(["%%model", line]),
 93                    formatted,
 94                ]
 95            ),
 96            replace=True,
 97        )
 98
 99        with open(model._path, "w", encoding="utf-8") as file:
100            file.write(formatted)
101
102        self._context.upsert_model(model)
103        self._context.console.show_sql(
104            self._context.render(
105                model.name,
106                start=args.start,
107                end=args.end,
108                latest=args.latest,
109            ).sql(pretty=True, dialect=args.dialect or model.dialect)
110        )

Renders the model and automatically fills in an editable cell with the model definition.

@magic_arguments()
@argument('model', type=str, help='The model.')
@argument('test_name', type=str, nargs='?', default=None, help='The test name to display')
@argument('--ls', action='store_true', help='List tests associated with a model')
@line_cell_magic
def test(self, line: str, test_def_raw: Optional[str] = None) -> None:
112    @magic_arguments()
113    @argument("model", type=str, help="The model.")
114    @argument("test_name", type=str, nargs="?", default=None, help="The test name to display")
115    @argument("--ls", action="store_true", help="List tests associated with a model")
116    @line_cell_magic
117    def test(self, line: str, test_def_raw: t.Optional[str] = None) -> None:
118        """Allow the user to list tests for a model, output a specific test, and then write their changes back"""
119        args = parse_argstring(self.test, line)
120        if not args.test_name and not args.ls:
121            raise MagicError("Must provide either test name or `--ls` to list tests")
122
123        model_test_metadatas = get_all_model_tests(
124            self._context.test_directory_path,
125            ignore_patterns=self._context.ignore_patterns,
126        )
127        tests: t.Dict[str, t.Dict[str, ModelTestMetadata]] = defaultdict(dict)
128        for model_test_metadata in model_test_metadatas:
129            model = model_test_metadata.body.get("model")
130            if not model:
131                self._context.console.log_error(
132                    f"Test found that does not have `model` defined: {model_test_metadata.path}"
133                )
134            tests[model][model_test_metadata.test_name] = model_test_metadata
135        if args.ls:
136            # TODO: Provide better UI for displaying tests
137            for test_name in tests[args.model]:
138                self._context.console.log_status_update(test_name)
139            return
140
141        test = tests[args.model][args.test_name]
142        test_def = yaml_load(test_def_raw) if test_def_raw else test.body
143        test_def_output = yaml_dumps(test_def)
144
145        self._shell.set_next_input(
146            "\n".join(
147                [
148                    " ".join(["%%test", line]),
149                    test_def_output,
150                ]
151            ),
152            replace=True,
153        )
154
155        with open(test.path, "r+", encoding="utf-8") as file:
156            content = yaml_load(file.read())
157            content[args.test_name] = test_def
158            file.seek(0)
159            YAML().dump(content, file)
160            file.truncate()

Allow the user to list tests for a model, output a specific test, and then write their changes back

@magic_arguments()
@argument('environment', nargs='?', type=str, help='The environment to run the plan against')
@argument('--start', '-s', type=str, help='Start date to backfill.')
@argument('--end', '-e', type=str, help='End date to backfill.')
@argument('--create-from', type=str, help="The environment to create the target environment from if it doesn't exist. Default: prod.")
@argument('--skip-tests', '-t', action='store_true', help='Skip the unit tests defined for the model.')
@argument('--restate-model', '-r', type=str, nargs='*', 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.')
@argument('--no-gaps', '-g', action='store_true', help='Ensure that new snapshots have no data gaps when comparing to existing snapshots for matching models in the target environment.')
@argument('--skip-backfill', action='store_true', help='Skip the backfill step.')
@argument('--forward-only', action='store_true', help='Create a plan for forward-only changes.')
@argument('--no-prompts', action='store_true', 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.')
@argument('--auto-apply', action='store_true', help='Automatically applies the new plan after creation.')
@argument('--no-auto-categorization', action='store_true', help='Disable automatic change categorization.', default=None)
@line_magic
def plan(self, line: str) -> None:
162    @magic_arguments()
163    @argument(
164        "environment",
165        nargs="?",
166        type=str,
167        help="The environment to run the plan against",
168    )
169    @argument("--start", "-s", type=str, help="Start date to backfill.")
170    @argument("--end", "-e", type=str, help="End date to backfill.")
171    @argument(
172        "--create-from",
173        type=str,
174        help="The environment to create the target environment from if it doesn't exist. Default: prod.",
175    )
176    @argument(
177        "--skip-tests",
178        "-t",
179        action="store_true",
180        help="Skip the unit tests defined for the model.",
181    )
182    @argument(
183        "--restate-model",
184        "-r",
185        type=str,
186        nargs="*",
187        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.",
188    )
189    @argument(
190        "--no-gaps",
191        "-g",
192        action="store_true",
193        help="Ensure that new snapshots have no data gaps when comparing to existing snapshots for matching models in the target environment.",
194    )
195    @argument(
196        "--skip-backfill",
197        action="store_true",
198        help="Skip the backfill step.",
199    )
200    @argument(
201        "--forward-only",
202        action="store_true",
203        help="Create a plan for forward-only changes.",
204    )
205    @argument(
206        "--no-prompts",
207        action="store_true",
208        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.",
209    )
210    @argument(
211        "--auto-apply",
212        action="store_true",
213        help="Automatically applies the new plan after creation.",
214    )
215    @argument(
216        "--no-auto-categorization",
217        action="store_true",
218        help="Disable automatic change categorization.",
219        default=None,
220    )
221    @line_magic
222    def plan(self, line: str) -> None:
223        """Goes through a set of prompts to both establish a plan and apply it"""
224        self._context.refresh()
225        args = parse_argstring(self.plan, line)
226
227        # Since the magics share a context we want to clear out any state before generating a new plan
228        console = self._context.console
229        self._context.console = get_console(display=self.display)
230
231        self._context.plan(
232            args.environment,
233            start=args.start,
234            end=args.end,
235            create_from=args.create_from,
236            skip_tests=args.skip_tests,
237            restate_models=args.restate_model,
238            no_gaps=args.no_gaps,
239            skip_backfill=args.skip_backfill,
240            forward_only=args.forward_only,
241            no_prompts=args.no_prompts,
242            auto_apply=args.auto_apply,
243            no_auto_categorization=args.no_auto_categorization,
244        )
245        self._context.console = console

Goes through a set of prompts to both establish a plan and apply it

@magic_arguments()
@argument('environment', nargs='?', type=str, help='The environment to run against')
@argument('--start', '-s', type=str, help='Start date to evaluate.')
@argument('--end', '-e', type=str, help='End date to evaluate.')
@argument('--skip-janitor', action='store_true', help='Skip the jantitor task.')
@line_magic
def run_dag(self, line: str) -> None:
247    @magic_arguments()
248    @argument(
249        "environment",
250        nargs="?",
251        type=str,
252        help="The environment to run against",
253    )
254    @argument("--start", "-s", type=str, help="Start date to evaluate.")
255    @argument("--end", "-e", type=str, help="End date to evaluate.")
256    @argument("--skip-janitor", action="store_true", help="Skip the jantitor task.")
257    @line_magic
258    def run_dag(self, line: str) -> None:
259        """Evaluate the DAG of models using the built-in scheduler."""
260        args = parse_argstring(self.run_dag, line)
261
262        # Since the magics share a context we want to clear out any state before generating a new plan
263        console = self._context.console
264        self._context.console = get_console(display=self.display)
265
266        self._context.run(
267            args.environment,
268            start=args.start,
269            end=args.end,
270            skip_janitor=args.skip_janitor,
271        )
272        self._context.console = console

Evaluate the DAG of models using the built-in scheduler.

@magic_arguments()
@argument('model', type=str, help='The model.')
@argument('--start', '-s', type=str, help='Start date to render.')
@argument('--end', '-e', type=str, help='End date to render.')
@argument('--latest', '-l', type=str, help='Latest date to render.')
@argument('--limit', type=int, help='The number of rows which the query should be limited to.')
@line_magic
def evaluate(self, line: str) -> None:
274    @magic_arguments()
275    @argument("model", type=str, help="The model.")
276    @argument("--start", "-s", type=str, help="Start date to render.")
277    @argument("--end", "-e", type=str, help="End date to render.")
278    @argument("--latest", "-l", type=str, help="Latest date to render.")
279    @argument(
280        "--limit",
281        type=int,
282        help="The number of rows which the query should be limited to.",
283    )
284    @line_magic
285    def evaluate(self, line: str) -> None:
286        """Evaluate a model query and fetches a dataframe."""
287        self._context.refresh()
288        args = parse_argstring(self.evaluate, line)
289
290        df = self._context.evaluate(
291            args.model,
292            start=args.start,
293            end=args.end,
294            latest=args.latest,
295            limit=args.limit,
296        )
297        self.display(df)

Evaluate a model query and fetches a dataframe.

@magic_arguments()
@argument('model', type=str, help='The model.')
@argument('--start', '-s', type=str, help='Start date to render.')
@argument('--end', '-e', type=str, help='End date to render.')
@argument('--latest', '-l', type=str, help='Latest date to render.')
@argument('--expand', type=t.Union[bool, t.Iterable[str]], 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.')
@argument('--dialect', type=str, help='SQL dialect to render.')
@line_magic
def render(self, line: str) -> None:
299    @magic_arguments()
300    @argument("model", type=str, help="The model.")
301    @argument("--start", "-s", type=str, help="Start date to render.")
302    @argument("--end", "-e", type=str, help="End date to render.")
303    @argument("--latest", "-l", type=str, help="Latest date to render.")
304    @argument(
305        "--expand",
306        type=t.Union[bool, t.Iterable[str]],
307        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.",
308    )
309    @argument("--dialect", type=str, help="SQL dialect to render.")
310    @line_magic
311    def render(self, line: str) -> None:
312        """Renders a model's query, optionally expanding referenced models."""
313        self._context.refresh()
314        args = parse_argstring(self.render, line)
315
316        query = self._context.render(
317            args.model,
318            start=args.start,
319            end=args.end,
320            latest=args.latest,
321            expand=args.expand,
322        )
323
324        self._context.console.show_sql(query.sql(pretty=True, dialect=args.dialect))

Renders a model's query, optionally expanding referenced models.

@magic_arguments()
@argument('df_var', default=None, nargs='?', type=str, help='An optional variable name to store the resulting dataframe.')
@line_cell_magic
def fetchdf(self, line: str, sql: str) -> None:
326    @magic_arguments()
327    @argument(
328        "df_var",
329        default=None,
330        nargs="?",
331        type=str,
332        help="An optional variable name to store the resulting dataframe.",
333    )
334    @line_cell_magic
335    def fetchdf(self, line: str, sql: str) -> None:
336        """Fetches a dataframe from sql, optionally storing it in a variable."""
337        args = parse_argstring(self.fetchdf, line)
338        df = self._context.fetchdf(sql)
339        if args.df_var:
340            self._shell.user_ns[args.df_var] = df
341        self.display(df)

Fetches a dataframe from sql, optionally storing it in a variable.

@magic_arguments()
@line_magic
def dag(self, line: str) -> None:
343    @magic_arguments()
344    @line_magic
345    def dag(self, line: str) -> None:
346        """Displays the dag"""
347        self._context.refresh()
348        dag = self._context.get_dag()
349        self.display(HTML(dag.pipe().decode("utf-8")))

Displays the dag

Inherited Members
IPython.core.magic.Magics
Magics
arg_err
format_latex
parse_options
default_option
traitlets.config.configurable.Configurable
config
parent
section_names
update_config
class_get_help
class_get_trait_help
class_print_help
class_config_section
class_config_rst_doc
traitlets.traitlets.HasTraits
setup_instance
cross_validation_lock
hold_trait_notifications
notify_change
on_trait_change
observe
unobserve
unobserve_all
add_traits
set_trait
class_trait_names
class_traits
class_own_traits
has_trait
trait_has_value
trait_values
trait_defaults
trait_names
traits
trait_metadata
class_own_trait_events
trait_events
def register_magics() -> None:
359def register_magics() -> None:
360    try:
361        shell = get_ipython()  # type: ignore
362        shell.register_magics(SQLMeshMagics)
363    except NameError:
364        pass