Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions examples/sushi/models/orders.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
"end_ts": "int",
"event_date": "date",
},
signals=[("test_signal", {"arg": 1})],
)
def execute(
context: ExecutionContext,
Expand Down
1 change: 1 addition & 0 deletions sqlmesh/core/loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -672,6 +672,7 @@ def _load_python_models(
default_catalog=self.context.default_catalog,
infer_names=self.config.model_naming.infer_names,
audit_definitions=audits,
signal_definitions=signals,
default_catalog_per_gateway=self.context.default_catalog_per_gateway,
):
if model.enabled:
Expand Down
36 changes: 20 additions & 16 deletions sqlmesh/core/model/decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from sqlglot.dialects.dialect import DialectType

from sqlmesh.core.macros import MacroRegistry
from sqlmesh.core.signal import SignalRegistry
from sqlmesh.utils.jinja import JinjaMacroRegistry
from sqlmesh.core import constants as c
from sqlmesh.core.dialect import MacroFunc, parse_one
Expand Down Expand Up @@ -48,23 +49,24 @@ def __init__(self, name: t.Optional[str] = None, is_sql: bool = False, **kwargs:
self.kwargs = kwargs

# Make sure that argument values are expressions in order to pass validation in ModelMeta.
calls = self.kwargs.pop("audits", [])
self.kwargs["audits"] = [
(
(call, {})
if isinstance(call, str)
else (
call[0],
{
arg_key: exp.convert(
tuple(arg_value) if isinstance(arg_value, list) else arg_value
)
for arg_key, arg_value in call[1].items()
},
for function_call_attribute in ("audits", "signals"):
calls = self.kwargs.pop(function_call_attribute, [])
self.kwargs[function_call_attribute] = [
(
(call, {})
if isinstance(call, str)
else (
call[0],
{
arg_key: exp.convert(
tuple(arg_value) if isinstance(arg_value, list) else arg_value
)
for arg_key, arg_value in call[1].items()
},
)
)
)
for call in calls
]
for call in calls
]

if "default_catalog" in kwargs:
raise ConfigError("`default_catalog` cannot be set on a per-model basis.")
Expand Down Expand Up @@ -142,6 +144,7 @@ def model(
defaults: t.Optional[t.Dict[str, t.Any]] = None,
macros: t.Optional[MacroRegistry] = None,
jinja_macros: t.Optional[JinjaMacroRegistry] = None,
signal_definitions: t.Optional[SignalRegistry] = None,
audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None,
dialect: t.Optional[str] = None,
time_column_format: str = c.DEFAULT_TIME_COLUMN_FORMAT,
Expand Down Expand Up @@ -223,6 +226,7 @@ def model(
"macros": macros,
"jinja_macros": jinja_macros,
"audit_definitions": audit_definitions,
"signal_definitions": signal_definitions,
"blueprint_variables": blueprint_variables,
**rendered_fields,
}
Expand Down
2 changes: 1 addition & 1 deletion tests/core/test_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -2122,7 +2122,7 @@ def test_check_intervals(sushi_context, mocker):
intervals = sushi_context.check_intervals(environment=None, no_signals=False, select_models=[])

min_intervals = 19
assert spy.call_count == 1
assert spy.call_count == 2
assert len(intervals) >= min_intervals

for i in intervals.values():
Expand Down
24 changes: 24 additions & 0 deletions tests/core/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -5303,6 +5303,30 @@ def my_signal(batch):
)


def test_load_python_model_with_signals():
@signal()
def always_true(batch):
return True

@model(
name="model_with_signal",
kind="full",
columns={'"COL"': "int"},
signals=[("always_true", {})],
)
def model_with_signal(context, **kwargs):
return pd.DataFrame([{"COL": 1}])

models = model.get_registry()["model_with_signal"].models(
get_variables=lambda _: {},
path=Path("."),
module_path=Path("."),
signal_definitions=signal.get_registry(),
)
assert len(models) == 1
assert models[0].signals == [("always_true", {})]


def test_null_column_type():
expressions = d.parse(
"""
Expand Down