sqlmesh.magics API documentation

Edit on GitHub

# sqlmesh.magics

View Source

      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): View Source

     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: View Source

    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: View Source

     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: View Source

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    359def register_magics() -> None:
    360    try:
    361        shell = get_ipython()  # type: ignore
    362        shell.register_magics(SQLMeshMagics)
    363    except NameError:
    364        pass