sqlmesh 0.202.2.dev15__tar.gz → 0.203.1.dev1__tar.gz

{sqlmesh-0.202.2.dev15 → sqlmesh-0.203.1.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.202.2.dev15
+Version: 0.203.1.dev1
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.202.2.dev15 → sqlmesh-0.203.1.dev1}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.202.2.dev15'
-__version_tuple__ = version_tuple = (0, 202, 2, 'dev15')
+__version__ = version = '0.203.1.dev1'
+__version_tuple__ = version_tuple = (0, 203, 1, 'dev1')

{sqlmesh-0.202.2.dev15 → sqlmesh-0.203.1.dev1}/sqlmesh/core/config/model.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import typing as t
 
+from sqlglot import exp
 from sqlmesh.core.dialect import parse_one, extract_func_call
 from sqlmesh.core.config.base import BaseConfig
 from sqlmesh.core.model.kind import (
@@ -41,6 +42,9 @@ class ModelDefaultsConfig(BaseConfig):
         allow_partials: Whether the models can process partial (incomplete) data intervals.
         enabled: Whether the models are enabled.
         interval_unit: The temporal granularity of the models data intervals. By default computed from cron.
+        pre_statements: The list of SQL statements that get executed before a model runs.
+        post_statements: The list of SQL statements that get executed before a model runs.
+        on_virtual_update: The list of SQL statements to be executed after the virtual update.
 
     """
 
@@ -61,6 +65,9 @@ class ModelDefaultsConfig(BaseConfig):
     interval_unit: t.Optional[t.Union[str, IntervalUnit]] = None
     enabled: t.Optional[t.Union[str, bool]] = None
     formatting: t.Optional[t.Union[str, bool]] = None
+    pre_statements: t.Optional[t.List[t.Union[str, exp.Expression]]] = None
+    post_statements: t.Optional[t.List[t.Union[str, exp.Expression]]] = None
+    on_virtual_update: t.Optional[t.List[t.Union[str, exp.Expression]]] = None
 
     _model_kind_validator = model_kind_validator
     _on_destructive_change_validator = on_destructive_change_validator

{sqlmesh-0.202.2.dev15 → sqlmesh-0.203.1.dev1}/sqlmesh/core/model/definition.py

@@ -2472,6 +2472,24 @@ def _create_model(
 
     statements: t.List[t.Union[exp.Expression, t.Tuple[exp.Expression, bool]]] = []
 
+    # Merge default pre_statements with model-specific pre_statements
+    if "pre_statements" in defaults:
+        kwargs["pre_statements"] = [
+            exp.maybe_parse(stmt, dialect=dialect) for stmt in defaults["pre_statements"]
+        ] + kwargs.get("pre_statements", [])
+
+    # Merge default post_statements with model-specific post_statements
+    if "post_statements" in defaults:
+        kwargs["post_statements"] = [
+            exp.maybe_parse(stmt, dialect=dialect) for stmt in defaults["post_statements"]
+        ] + kwargs.get("post_statements", [])
+
+    # Merge default on_virtual_update with model-specific on_virtual_update
+    if "on_virtual_update" in defaults:
+        kwargs["on_virtual_update"] = [
+            exp.maybe_parse(stmt, dialect=dialect) for stmt in defaults["on_virtual_update"]
+        ] + kwargs.get("on_virtual_update", [])
+
     if "pre_statements" in kwargs:
         statements.extend(kwargs["pre_statements"])
     if "query" in kwargs:

{sqlmesh-0.202.2.dev15 → sqlmesh-0.203.1.dev1}/sqlmesh.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.202.2.dev15
+Version: 0.203.1.dev1
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.202.2.dev15 → sqlmesh-0.203.1.dev1}/tests/core/test_config.py

@@ -676,6 +676,65 @@ model_defaults:
     assert config.model_defaults.audits[1][1]["threshold"].this == "1000"
 
 
+def test_load_model_defaults_statements(tmp_path):
+    config_path = tmp_path / "config_model_defaults_statements.yaml"
+    with open(config_path, "w", encoding="utf-8") as fd:
+        fd.write(
+            """
+model_defaults:
+    dialect: duckdb
+    pre_statements:
+        - SET memory_limit = '10GB'
+        - CREATE TEMP TABLE temp_data AS SELECT 1 as id
+    post_statements:
+        - DROP TABLE IF EXISTS temp_data
+        - ANALYZE @this_model
+        - SET memory_limit = '5GB'
+    on_virtual_update:
+        - UPDATE stats_table SET last_update = CURRENT_TIMESTAMP
+        """
+        )
+
+    config = load_config_from_paths(
+        Config,
+        project_paths=[config_path],
+    )
+
+    assert config.model_defaults.pre_statements is not None
+    assert len(config.model_defaults.pre_statements) == 2
+    assert isinstance(exp.maybe_parse(config.model_defaults.pre_statements[0]), exp.Set)
+    assert isinstance(exp.maybe_parse(config.model_defaults.pre_statements[1]), exp.Create)
+
+    assert config.model_defaults.post_statements is not None
+    assert len(config.model_defaults.post_statements) == 3
+    assert isinstance(exp.maybe_parse(config.model_defaults.post_statements[0]), exp.Drop)
+    assert isinstance(exp.maybe_parse(config.model_defaults.post_statements[1]), exp.Analyze)
+    assert isinstance(exp.maybe_parse(config.model_defaults.post_statements[2]), exp.Set)
+
+    assert config.model_defaults.on_virtual_update is not None
+    assert len(config.model_defaults.on_virtual_update) == 1
+    assert isinstance(exp.maybe_parse(config.model_defaults.on_virtual_update[0]), exp.Update)
+
+
+def test_load_model_defaults_validation_statements(tmp_path):
+    config_path = tmp_path / "config_model_defaults_statements_wrong.yaml"
+    with open(config_path, "w", encoding="utf-8") as fd:
+        fd.write(
+            """
+model_defaults:
+    dialect: duckdb
+    pre_statements:
+        - 313
+        """
+        )
+
+    with pytest.raises(TypeError, match=r"expected str instance, int found"):
+        config = load_config_from_paths(
+            Config,
+            project_paths=[config_path],
+        )
+
+
 def test_scheduler_config(tmp_path_factory):
     config_path = tmp_path_factory.mktemp("yaml_config") / "config.yaml"
     with open(config_path, "w", encoding="utf-8") as fd:

{sqlmesh-0.202.2.dev15 → sqlmesh-0.203.1.dev1}/tests/core/test_context.py

@@ -2731,3 +2731,139 @@ def test_plan_min_intervals_adjusted_for_downstream(tmp_path: Path):
     assert context.engine_adapter.fetchall(
         "select min(start_dt), max(end_dt) from sqlmesh_example__pr_env.unrelated_monthly_model"
     ) == [(to_datetime("2020-01-01 00:00:00"), to_datetime("2020-01-31 23:59:59.999999"))]
+
+
+def test_defaults_pre_post_statements(tmp_path: Path):
+    config_path = tmp_path / "config.yaml"
+    models_path = tmp_path / "models"
+    models_path.mkdir()
+
+    # Create config with default statements
+    config_path.write_text(
+        """
+model_defaults:
+    dialect: duckdb
+    pre_statements:
+        - SET memory_limit = '10GB'
+        - SET threads = @var1
+    post_statements:
+        - ANALYZE @this_model
+variables:
+    var1: 4
+"""
+    )
+
+    # Create a model
+    model_path = models_path / "test_model.sql"
+    model_path.write_text(
+        """
+MODEL (
+    name test_model,
+    kind FULL
+);
+
+SELECT 1 as id, 'test' as status;
+"""
+    )
+
+    ctx = Context(paths=[tmp_path])
+
+    # Initial plan and apply
+    initial_plan = ctx.plan(auto_apply=True, no_prompts=True)
+    assert len(initial_plan.new_snapshots) == 1
+
+    snapshot = list(initial_plan.new_snapshots)[0]
+    model = snapshot.model
+
+    # Verify statements are in the model and python environment has been popuplated
+    assert len(model.pre_statements) == 2
+    assert len(model.post_statements) == 1
+    assert model.python_env[c.SQLMESH_VARS].payload == "{'var1': 4}"
+
+    # Verify the statements contain the expected SQL
+    assert model.pre_statements[0].sql() == "SET memory_limit = '10GB'"
+    assert model.render_pre_statements()[0].sql() == "SET \"memory_limit\" = '10GB'"
+    assert model.pre_statements[1].sql() == "SET threads = @var1"
+    assert model.render_pre_statements()[1].sql() == 'SET "threads" = 4'
+
+    # Update config to change pre_statement
+    config_path.write_text(
+        """
+model_defaults:
+    dialect: duckdb
+    pre_statements:
+        - SET memory_limit = '5GB'  # Changed value
+    post_statements:
+        - ANALYZE @this_model
+"""
+    )
+
+    # Reload context and create new plan
+    ctx = Context(paths=[tmp_path])
+    updated_plan = ctx.plan(no_prompts=True)
+
+    # Should detect a change due to different pre_statements
+    assert len(updated_plan.directly_modified) == 1
+
+    # Apply the plan
+    ctx.apply(updated_plan)
+
+    # Reload the models to get the updated version
+    ctx.load()
+    new_model = ctx.models['"test_model"']
+
+    # Verify updated statements
+    assert len(new_model.pre_statements) == 1
+    assert new_model.pre_statements[0].sql() == "SET memory_limit = '5GB'"
+    assert new_model.render_pre_statements()[0].sql() == "SET \"memory_limit\" = '5GB'"
+
+    # Verify the change was detected by the plan
+    assert len(updated_plan.directly_modified) == 1
+
+
+def test_model_defaults_statements_with_on_virtual_update(tmp_path: Path):
+    config_path = tmp_path / "config.yaml"
+    models_path = tmp_path / "models"
+    models_path.mkdir()
+
+    # Create config with on_virtual_update
+    config_path.write_text(
+        """
+model_defaults:
+    dialect: duckdb
+    on_virtual_update:
+        - SELECT  'Model-defailt virtual update' AS message
+"""
+    )
+
+    # Create a model with its own on_virtual_update as wel
+    model_path = models_path / "test_model.sql"
+    model_path.write_text(
+        """
+MODEL (
+    name test_model,
+    kind FULL
+);
+
+SELECT 1 as id, 'test' as name;
+
+ON_VIRTUAL_UPDATE_BEGIN;
+SELECT 'Model-specific update' AS message;
+ON_VIRTUAL_UPDATE_END;
+"""
+    )
+
+    ctx = Context(paths=[tmp_path])
+
+    # Plan and apply
+    plan = ctx.plan(auto_apply=True, no_prompts=True)
+
+    snapshot = list(plan.new_snapshots)[0]
+    model = snapshot.model
+
+    # Verify both default and model-specific on_virtual_update statements
+    assert len(model.on_virtual_update) == 2
+
+    # Default statements should come first
+    assert model.on_virtual_update[0].sql() == "SELECT 'Model-defailt virtual update' AS message"
+    assert model.on_virtual_update[1].sql() == "SELECT 'Model-specific update' AS message"

{sqlmesh-0.202.2.dev15 → sqlmesh-0.203.1.dev1}/tests/core/test_model.py

@@ -1691,6 +1691,115 @@ def test_description(sushi_context):
     assert sushi_context.models['"memory"."sushi"."orders"'].description == "Table of sushi orders."
 
 
+def test_model_defaults_statements_merge():
+    model_defaults = ModelDefaultsConfig(
+        dialect="duckdb",
+        pre_statements=[
+            "SET enable_progress_bar = true",
+            "CREATE TEMP TABLE default_temp AS SELECT 1",
+        ],
+        post_statements=[
+            "DROP TABLE IF EXISTS default_temp",
+            "grant select on @this_model to group reporter",
+        ],
+        on_virtual_update=["ANALYZE"],
+    )
+
+    # Create a model with its own statements as well
+    expressions = parse(
+        """
+    MODEL (
+        name test_model,
+        kind FULL
+    );
+
+    CREATE TEMP TABLE model_temp AS SELECT 2;
+
+    SELECT * FROM test_table;
+
+    DROP TABLE IF EXISTS model_temp;
+
+    ON_VIRTUAL_UPDATE_BEGIN;
+    UPDATE stats_table SET last_update = CURRENT_TIMESTAMP;
+    ON_VIRTUAL_UPDATE_END;
+    """
+    )
+
+    model = load_sql_based_model(
+        expressions,
+        path=Path("./test_model.sql"),
+        defaults=model_defaults.dict(),
+    )
+
+    # Check that pre_statements contains both default and model-specific statements
+    assert len(model.pre_statements) == 3
+    assert model.pre_statements[0].sql() == "SET enable_progress_bar = TRUE"
+    assert model.pre_statements[1].sql() == "CREATE TEMPORARY TABLE default_temp AS SELECT 1"
+    assert model.pre_statements[2].sql() == "CREATE TEMPORARY TABLE model_temp AS SELECT 2"
+
+    # Check that post_statements contains both default and model-specific statements
+    assert len(model.post_statements) == 3
+    assert model.post_statements[0].sql() == "DROP TABLE IF EXISTS default_temp"
+    assert model.post_statements[1].sql() == "GRANT SELECT ON @this_model TO GROUP reporter"
+    assert model.post_statements[2].sql() == "DROP TABLE IF EXISTS model_temp"
+
+    # Check that the query is rendered correctly with @this_model resolved to table name
+    assert (
+        model.render_post_statements()[1].sql()
+        == 'GRANT SELECT ON "test_model" TO GROUP "reporter"'
+    )
+
+    # Check that on_virtual_update contains both default and model-specific statements
+    assert len(model.on_virtual_update) == 2
+    assert model.on_virtual_update[0].sql() == "ANALYZE"
+    assert (
+        model.on_virtual_update[1].sql()
+        == "UPDATE stats_table SET last_update = CURRENT_TIMESTAMP()"
+    )
+
+
+def test_model_defaults_statements_integration():
+    config = Config(
+        model_defaults=ModelDefaultsConfig(
+            dialect="postgres",
+            pre_statements=["SET memory_limit = '10GB'"],
+            post_statements=["VACUUM ANALYZE"],
+            on_virtual_update=["GRANT SELECT ON @this_model TO GROUP public"],
+        )
+    )
+
+    expressions = parse(
+        """
+    MODEL (
+        name test_model,
+        kind FULL
+    );
+
+    SELECT * FROM source_table;
+    """
+    )
+
+    model = load_sql_based_model(
+        expressions,
+        path=Path("./test_model.sql"),
+        defaults=config.model_defaults.dict(),
+    )
+
+    # Verify defaults were applied
+    assert len(model.pre_statements) == 1
+    assert model.pre_statements[0].sql() == "SET memory_limit = '10GB'"
+
+    assert len(model.post_statements) == 1
+    assert isinstance(model.post_statements[0], exp.Command)
+
+    assert len(model.on_virtual_update) == 1
+    assert model.on_virtual_update[0].sql() == "GRANT SELECT ON @this_model TO GROUP public"
+    assert (
+        model.render_on_virtual_update()[0].sql()
+        == 'GRANT SELECT ON "test_model" TO GROUP "public"'
+    )
+
+
 def test_render_definition():
     expressions = d.parse(
         """
@@ -5568,7 +5677,7 @@ def test_when_matched_normalization() -> None:
             when_matched (
                 WHEN MATCHED THEN UPDATE SET
                     target.key_a = source.key_a,
-                    target.key_b = source.key_b,            
+                    target.key_b = source.key_b,
             )
           )
         );