Skip to content

Commit

Permalink
Feat: Extend support of project wide model properties (#3832)
Browse files Browse the repository at this point in the history
  • Loading branch information
themisvaltinos authored Feb 14, 2025
1 parent f24f04f commit beaebe3
Show file tree
Hide file tree
Showing 4 changed files with 101 additions and 3 deletions.
3 changes: 3 additions & 0 deletions docs/reference/model_configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,9 @@ The SQLMesh project-level `model_defaults` key supports the following options, d
- audits (described [here](../concepts/audits.md#generic-audits))
- optimize_query
- validate_query
- allow_partials
- enabled
- interval_unit


### Model Naming
Expand Down
11 changes: 10 additions & 1 deletion sqlmesh/core/config/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
model_kind_validator,
on_destructive_change_validator,
)
from sqlmesh.core.node import IntervalUnit
from sqlmesh.utils.date import TimeLike
from sqlmesh.core.model.meta import FunctionCall
from sqlmesh.utils.pydantic import field_validator
Expand All @@ -36,7 +37,12 @@ class ModelDefaultsConfig(BaseConfig):
virtual_properties: A key-value mapping of arbitrary properties that are applied to the model view in the virtual layer.
session_properties: A key-value mapping of properties specific to the target engine that are applied to the engine session.
audits: The audits to be applied globally to all models in the project.
optimize_query: Whether the SQL models should be optimized
optimize_query: Whether the SQL models should be optimized.
validate_query: Whether the SQL models should be validated at compile time.
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.
"""

kind: t.Optional[ModelKind] = None
Expand All @@ -53,6 +59,9 @@ class ModelDefaultsConfig(BaseConfig):
audits: t.Optional[t.List[FunctionCall]] = None
optimize_query: t.Optional[bool] = None
validate_query: t.Optional[bool] = None
allow_partials: t.Optional[bool] = None
interval_unit: t.Optional[IntervalUnit] = None
enabled: t.Optional[bool] = None

_model_kind_validator = model_kind_validator
_on_destructive_change_validator = on_destructive_change_validator
Expand Down
6 changes: 5 additions & 1 deletion sqlmesh/core/model/definition.py
Original file line number Diff line number Diff line change
Expand Up @@ -2153,7 +2153,11 @@ def _create_model(
) -> Model:
_validate_model_fields(klass, {"name", *kwargs} - {"grain", "table_properties"}, path)

for prop in ["session_properties", "physical_properties", "virtual_properties"]:
for prop in [
"session_properties",
"physical_properties",
"virtual_properties",
]:
kwargs[prop] = _resolve_properties((defaults or {}).get(prop), kwargs.get(prop))

dialect = dialect or ""
Expand Down
84 changes: 83 additions & 1 deletion tests/core/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -3397,16 +3397,23 @@ def test_project_level_properties(sushi_context):
"target_lag": "1 hour",
},
virtual_properties={"creatable_type": "SECURE"},
enabled=False,
allow_partials=True,
interval_unit="quarter_hour",
validate_query=True,
optimize_query=True,
cron="@hourly",
)

model = load_sql_based_model(
d.parse(
"""
MODEL (
name test_schema.test_model,
kind FULL,
virtual_properties (
creatable_type = None
)
),
);
SELECT a FROM tbl;
""",
Expand All @@ -3415,6 +3422,14 @@ def test_project_level_properties(sushi_context):
defaults=model_defaults.dict(),
)

# Validate use of project wide defaults
assert not model.enabled
assert model.allow_partials
assert model.interval_unit == IntervalUnit.QUARTER_HOUR
assert model.optimize_query
assert model.validate_query
assert model.cron == "@hourly"

assert model.session_properties == {
"some_bool": False,
"some_float": 0.1,
Expand All @@ -3429,6 +3444,73 @@ def test_project_level_properties(sushi_context):
# Validate disabling global property
assert not model.virtual_properties

model_2 = load_sql_based_model(
d.parse(
"""
MODEL (
name test_schema.test_model_2,
kind FULL,
allow_partials False,
interval_unit hour,
cron '@daily'
);
SELECT a, b FROM tbl;
""",
default_dialect="snowflake",
),
defaults=model_defaults.dict(),
)

# Validate overriding of project wide defaults
assert not model_2.allow_partials
assert model_2.interval_unit == IntervalUnit.HOUR
assert model_2.cron == "@daily"


def test_project_level_properties_python_model():
model_defaults = {
"physical_properties": {
"partition_expiration_days": 13,
"creatable_type": "TRANSIENT",
},
"description": "general model description",
"enabled": False,
"allow_partials": True,
"interval_unit": "quarter_hour",
"validate_query": True,
"optimize_query": True,
}

@model(
name="python_model_t_defaults",
kind="full",
columns={"some_col": "int"},
physical_properties={"partition_expiration_days": 7},
)
def python_model_prop(context, **kwargs):
context.resolve_table("foo")

m = model.get_registry()["python_model_t_defaults"].model(
module_path=Path("."),
path=Path("."),
dialect="duckdb",
defaults=model_defaults,
)

assert m.physical_properties == {
"partition_expiration_days": exp.convert(7),
"creatable_type": exp.convert("TRANSIENT"),
}

# Even if in the project wide defaults these are ignored for python models
assert not m.optimize_query
assert not m.validate_query

assert not m.enabled
assert m.allow_partials
assert m.interval_unit == IntervalUnit.QUARTER_HOUR


def test_model_session_properties(sushi_context):
assert sushi_context.models['"memory"."sushi"."items"'].session_properties == {
Expand Down

0 comments on commit beaebe3

Please sign in to comment.