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
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, ANDUse 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.
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.
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.
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
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
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.
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.
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.
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.
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