Feat: Add dlt_refresh command to update models from a dlt pipeline (#3340)

Author: themisvaltinos

docs/integrations/dlt.md

@@ -28,6 +28,28 @@ This will create the configuration file and directories, which are found in all
 
 SQLMesh will also automatically generate models to ingest data from the pipeline incrementally. Incremental loading is ideal for large datasets where recomputing entire tables is resource-intensive. In this case utilizing the [`INCREMENTAL_BY_TIME_RANGE` model kind](../concepts/models/model_kinds.md#incremental_by_time_range). However, these model definitions can be customized to meet your specific project needs.
 
+### Generating models on demand
+
+To update the models in your SQLMesh project on demand, use the `dlt_refresh` command. This allows you to either specify individual tables to generate incremental models from or update all models at once.
+
+- **Generate all missing tables**:
+  
+```bash
+$ sqlmesh dlt_refresh <pipeline-name>
+```
+
+- **Generate all missing tables and overwrite existing ones** (use with `--force` or `-f`):
+  
+```bash
+$ sqlmesh dlt_refresh <pipeline-name> --force
+```
+
+- **Generate specific dlt tables** (using `--table` or `-t`):
+
+```bash
+$ sqlmesh dlt_refresh <pipeline-name> --table <dlt-table>
+```
+
 #### Configuration
 
 SQLMesh will retrieve the data warehouse connection credentials from your dlt project to configure the `config.yaml` file. This configuration can be modified or customized as needed. For more details, refer to the [configuration guide](../guides/configuration.md).

docs/reference/cli.md

@@ -121,6 +121,18 @@ Options:
   --help               Show this message and exit.
 ```
 
+## dlt_refresh
+
+```
+Usage: dlt_refresh PIPELINE [OPTIONS]
+
+  Attaches to a DLT pipeline with the option to update specific or all models of the SQLMesh project.
+
+Options:
+  -t, --table TEXT  The DLT tables to generate SQLMesh models from. When none specified, all new missing tables will be generated.
+  -f, --force       If set it will overwrite existing models with the new generated models from the DLT tables.
+```
+
 ## diff
 
 ```

docs/reference/notebook.md

@@ -230,6 +230,17 @@ options:
   --file FILE, -f FILE  An optional file path to write the HTML output to.
 ```
 
+#### dlt_refresh
+```
+%dlt_refresh PIPELINE [--table] TABLE [--force]
+
+Attaches to a DLT pipeline with the option to update specific or all models of the SQLMesh project.
+
+options:
+  --table TABLE, -t TABLE  The DLT tables to generate SQLMesh models from. When none specified, all new missing tables will be generated.
+  --force, -f              If set it will overwrite existing models with the new generated models from the DLT tables.
+```
+
 #### fetchdf
 ```
 %%fetchdf [df_var]

sqlmesh/cli/main.py

@@ -916,3 +916,39 @@ def clean(obj: Context) -> None:
 def table_name(obj: Context, model_name: str, dev: bool) -> None:
     """Prints the name of the physical table for the given model."""
     print(obj.table_name(model_name, dev))
+
+
+@cli.command("dlt_refresh")
+@click.argument("pipeline", required=True)
+@click.option(
+    "-t",
+    "--table",
+    type=str,
+    multiple=True,
+    help="The specific dlt tables to refresh in the SQLMesh models.",
+)
+@click.option(
+    "-f",
+    "--force",
+    is_flag=True,
+    default=False,
+    help="If set, existing models are overwritten with the new DLT tables.",
+)
+@click.pass_context
+@error_handler
+@cli_analytics
+def dlt_refresh(
+    ctx: click.Context,
+    pipeline: str,
+    force: bool,
+    table: t.List[str] = [],
+) -> None:
+    """Attaches to a DLT pipeline with the option to update specific or all missing tables in the SQLMesh project."""
+    from sqlmesh.integrations.dlt import generate_dlt_models
+
+    sqlmesh_models = generate_dlt_models(ctx.obj, pipeline, list(table or []), force)
+    if sqlmesh_models:
+        model_names = "\n".join([f"- {model_name}" for model_name in sqlmesh_models])
+        ctx.obj.console.log_success(f"Updatde SQLMesh project with models:\n{model_names}")
+    else:
+        ctx.obj.console.log_success("All SQLMesh models are up to date.")

sqlmesh/integrations/dlt.py

@@ -4,11 +4,12 @@
 from pydantic import ValidationError
 from sqlglot import exp, parse_one
 from sqlmesh.core.config.connection import parse_connection_config
+from sqlmesh.core.context import Context
 from sqlmesh.utils.date import yesterday_ds
 
 
 def generate_dlt_models_and_settings(
-    pipeline_name: str, dialect: str
+    pipeline_name: str, dialect: str, tables: t.Optional[t.List[str]] = None
 ) -> t.Tuple[t.Set[t.Tuple[str, str]], str, str]:
     """This function attaches to a DLT pipeline and retrieves the connection configs and
     SQLMesh models based on the tables present in the pipeline's default schema.
@@ -42,8 +43,11 @@ def generate_dlt_models_and_settings(
     dlt_tables = {
         name: table
         for name, table in schema.tables.items()
-        if (has_table_seen_data(table) and not name.startswith(schema._dlt_tables_prefix))
-        or name == schema.loads_table_name
+        if (
+            (has_table_seen_data(table) and not name.startswith(schema._dlt_tables_prefix))
+            or name == schema.loads_table_name
+        )
+        and (name in tables if tables else True)
     }
 
     sqlmesh_models = set()
@@ -85,6 +89,27 @@ def generate_dlt_models_and_settings(
     return sqlmesh_models, format_config(configs, db_type), get_start_date(storage_ids)
 
 
+def generate_dlt_models(
+    context: Context, pipeline_name: str, tables: t.List[str], force: bool
+) -> t.List[str]:
+    from sqlmesh.cli.example_project import _create_models
+
+    sqlmesh_models, _, _ = generate_dlt_models_and_settings(
+        pipeline_name=pipeline_name,
+        dialect=context.config.dialect or "",
+        tables=tables if tables else None,
+    )
+
+    if not tables and not force:
+        existing_models = [m.name for m in context.models.values()]
+        sqlmesh_models = {model for model in sqlmesh_models if model[0] not in existing_models}
+
+    if sqlmesh_models:
+        _create_models(models_path=context.path / "models", models=sqlmesh_models)
+        return [model[0] for model in sqlmesh_models]
+    return []
+
+
 def generate_incremental_model(
     model_name: str,
     model_def_columns: str,

sqlmesh/magics.py

@@ -690,6 +690,42 @@ def table_name(self, context: Context, line: str) -> None:
         args = parse_argstring(self.table_name, line)
         context.console.log_status_update(context.table_name(args.model_name, args.dev))
 
+    @magic_arguments()
+    @argument(
+        "pipeline",
+        nargs="?",
+        type=str,
+        help="The dlt pipeline to attach for this SQLMesh project.",
+    )
+    @argument(
+        "--table",
+        "-t",
+        type=str,
+        nargs="*",
+        help="The specific dlt tables to refresh in the SQLMesh models.",
+    )
+    @argument(
+        "--force",
+        "-f",
+        action="store_true",
+        help="If set, existing models are overwritten with the new DLT tables.",
+    )
+    @line_magic
+    @pass_sqlmesh_context
+    def dlt_refresh(self, context: Context, line: str) -> None:
+        """Attaches to a DLT pipeline with the option to update specific or all missing tables in the SQLMesh project."""
+        from sqlmesh.integrations.dlt import generate_dlt_models
+
+        args = parse_argstring(self.dlt_refresh, line)
+        sqlmesh_models = generate_dlt_models(
+            context, args.pipeline, list(args.table or []), args.force
+        )
+        if sqlmesh_models:
+            model_names = "\n".join([f"- {model_name}" for model_name in sqlmesh_models])
+            context.console.log_success(f"Updated SQLMesh project with models:\n{model_names}")
+        else:
+            context.console.log_success("All SQLMesh models are up to date.")
+
     @magic_arguments()
     @argument(
         "--read",

tests/cli/test_cli.py

@@ -10,6 +10,7 @@
 from sqlmesh.cli.example_project import ProjectTemplate, init_example_project
 from sqlmesh.cli.main import cli
 from sqlmesh.core.context import Context
+from sqlmesh.integrations.dlt import generate_dlt_models
 from sqlmesh.utils.date import yesterday_ds
 
 FREEZE_TIME = "2023-01-01 00:00:00"
@@ -731,7 +732,11 @@ def test_plan_dlt(runner, tmp_path):
   TO_TIMESTAMP(CAST(_dlt_load_id AS DOUBLE)) BETWEEN @start_ds AND @end_ds
 """
 
-    with open(tmp_path / "models/incremental_sushi_types.sql") as file:
+    dlt_sushi_types_model_path = tmp_path / "models/incremental_sushi_types.sql"
+    dlt_loads_model_path = tmp_path / "models/incremental__dlt_loads.sql"
+    dlt_waiters_model_path = tmp_path / "models/incremental_waiters.sql"
+
+    with open(dlt_sushi_types_model_path) as file:
         incremental_model = file.read()
 
     expected_dlt_loads_model = """MODEL (
@@ -762,11 +767,14 @@ def test_plan_dlt(runner, tmp_path):
   TO_TIMESTAMP(CAST(load_id AS DOUBLE)) BETWEEN @start_ds AND @end_ds
 """
 
-    with open(tmp_path / "models/incremental__dlt_loads.sql") as file:
+    with open(dlt_loads_model_path) as file:
         dlt_loads_model = file.read()
 
     # Validate generated config and models
     assert config == expected_config
+    assert dlt_loads_model_path.exists()
+    assert dlt_sushi_types_model_path.exists()
+    assert dlt_waiters_model_path.exists()
     assert dlt_loads_model == expected_dlt_loads_model
     assert incremental_model == expected_incremental_model
 
@@ -777,4 +785,37 @@ def test_plan_dlt(runner, tmp_path):
 
     assert result.exit_code == 0
     assert_backfill_success(result)
+
+    # Remove and update with missing model
+    remove(dlt_waiters_model_path)
+    assert not dlt_waiters_model_path.exists()
+
+    # Update with force = False will generate only the missing model
+    context = Context(paths=tmp_path)
+    assert generate_dlt_models(context, "sushi", [], False) == [
+        "sushi_dataset_sqlmesh.incremental_waiters"
+    ]
+    assert dlt_waiters_model_path.exists()
+
+    # Remove all models
+    remove(dlt_waiters_model_path)
+    remove(dlt_loads_model_path)
+    remove(dlt_sushi_types_model_path)
+
+    # Update to generate a specific model: sushi_types
+    assert generate_dlt_models(context, "sushi", ["sushi_types"], False) == [
+        "sushi_dataset_sqlmesh.incremental_sushi_types"
+    ]
+
+    # Only the sushi_types should be generated now
+    assert not dlt_waiters_model_path.exists()
+    assert not dlt_loads_model_path.exists()
+    assert dlt_sushi_types_model_path.exists()
+
+    # Update with force = True will generate all models and overwrite existing ones
+    generate_dlt_models(context, "sushi", [], True)
+    assert dlt_loads_model_path.exists()
+    assert dlt_sushi_types_model_path.exists()
+    assert dlt_waiters_model_path.exists()
+
     remove(dataset_path)