Skip to content

Commit 51b475e

Browse files
committed
feat: add cron_tz to model_defaults
1 parent d5ceeb2 commit 51b475e

File tree

4 files changed

+88
-22
lines changed

4 files changed

+88
-22
lines changed

docs/reference/model_configuration.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,7 @@ The SQLMesh project-level `model_defaults` key supports the following options, d
178178
- kind
179179
- dialect
180180
- cron
181+
- cron_tz
181182
- owner
182183
- start
183184
- table_format

sqlmesh/core/config/model.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
OnAdditiveChange,
1515
)
1616
from sqlmesh.core.model.meta import FunctionCall
17-
from sqlmesh.core.node import IntervalUnit
17+
from sqlmesh.core.node import IntervalUnit, cron_tz_validator
1818
from sqlmesh.utils.date import TimeLike
1919
from sqlmesh.utils.pydantic import field_validator
2020

@@ -27,6 +27,7 @@ class ModelDefaultsConfig(BaseConfig):
2727
dialect: The SQL dialect that the model's query is written in.
2828
cron: A cron string specifying how often the model should be refreshed, leveraging the
2929
[croniter](https://github.com/kiorky/croniter) library.
30+
cron_tz: The timezone for the cron expression, defaults to UTC. [IANA time zones](https://docs.python.org/3/library/zoneinfo.html).
3031
owner: The owner of the model.
3132
start: The earliest date that the model will be backfilled for. If this is None,
3233
then the date is inferred by taking the most recent start date of its ancestors.
@@ -55,6 +56,7 @@ class ModelDefaultsConfig(BaseConfig):
5556
kind: t.Optional[ModelKind] = None
5657
dialect: t.Optional[str] = None
5758
cron: t.Optional[str] = None
59+
cron_tz: t.Any = None
5860
owner: t.Optional[str] = None
5961
start: t.Optional[TimeLike] = None
6062
table_format: t.Optional[str] = None
@@ -78,6 +80,7 @@ class ModelDefaultsConfig(BaseConfig):
7880
_model_kind_validator = model_kind_validator
7981
_on_destructive_change_validator = on_destructive_change_validator
8082
_on_additive_change_validator = on_additive_change_validator
83+
_cron_tz_validator = cron_tz_validator
8184

8285
@field_validator("audits", mode="before")
8386
def _audits_validator(cls, v: t.Any) -> t.Any:

sqlmesh/core/node.py

Lines changed: 24 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -259,6 +259,28 @@ def dbt_fqn(self) -> t.Optional[str]:
259259
IntervalUnit.FIVE_MINUTE: 60 * 5,
260260
}
261261

262+
def _cron_tz_validator(cls, v: t.Any) -> t.Optional[zoneinfo.ZoneInfo]:
263+
if not v or v == "UTC":
264+
return None
265+
266+
v = str_or_exp_to_str(v)
267+
268+
try:
269+
return zoneinfo.ZoneInfo(v)
270+
except Exception as e:
271+
available_timezones = zoneinfo.available_timezones()
272+
273+
if available_timezones:
274+
raise ConfigError(f"{e}. {v} must be in {available_timezones}.")
275+
else:
276+
raise ConfigError(
277+
f"{e}. IANA time zone data is not available on your system. `pip install tzdata` to leverage cron time zones or remove this field which will default to UTC."
278+
)
279+
280+
return None
281+
282+
cron_tz_validator = field_validator("cron_tz", mode="before")(_cron_tz_validator)
283+
262284

263285
class _Node(DbtInfoMixin, PydanticModel):
264286
"""
@@ -302,6 +324,8 @@ class _Node(DbtInfoMixin, PydanticModel):
302324
_croniter: t.Optional[CroniterCache] = None
303325
__inferred_interval_unit: t.Optional[IntervalUnit] = None
304326

327+
_cron_tz_validator = cron_tz_validator
328+
305329
def __str__(self) -> str:
306330
path = f": {self._path.name}" if self._path else ""
307331
return f"{self.__class__.__name__}<{self.name}{path}>"
@@ -328,27 +352,6 @@ def _name_validator(cls, v: t.Any) -> t.Optional[str]:
328352
return v.meta["sql"]
329353
return str(v)
330354

331-
@field_validator("cron_tz", mode="before")
332-
def _cron_tz_validator(cls, v: t.Any) -> t.Optional[zoneinfo.ZoneInfo]:
333-
if not v or v == "UTC":
334-
return None
335-
336-
v = str_or_exp_to_str(v)
337-
338-
try:
339-
return zoneinfo.ZoneInfo(v)
340-
except Exception as e:
341-
available_timezones = zoneinfo.available_timezones()
342-
343-
if available_timezones:
344-
raise ConfigError(f"{e}. {v} must be in {available_timezones}.")
345-
else:
346-
raise ConfigError(
347-
f"{e}. IANA time zone data is not available on your system. `pip install tzdata` to leverage cron time zones or remove this field which will default to UTC."
348-
)
349-
350-
return None
351-
352355
@field_validator("start", "end", mode="before")
353356
@classmethod
354357
def _date_validator(cls, v: t.Any) -> t.Optional[TimeLike]:

tests/core/test_config.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -930,6 +930,65 @@ def test_gateway_model_defaults(tmp_path):
930930
assert ctx.config.model_defaults == expected
931931

932932

933+
def test_model_defaults_cron_tz(tmp_path):
934+
"""Test that cron_tz can be set in model_defaults."""
935+
import zoneinfo
936+
937+
config_path = tmp_path / "config_model_defaults_cron_tz.yaml"
938+
with open(config_path, "w", encoding="utf-8") as fd:
939+
fd.write(
940+
"""
941+
model_defaults:
942+
dialect: duckdb
943+
cron: '@daily'
944+
cron_tz: 'America/Los_Angeles'
945+
"""
946+
)
947+
948+
config = load_config_from_paths(
949+
Config,
950+
project_paths=[config_path],
951+
)
952+
953+
assert config.model_defaults.cron == "@daily"
954+
assert config.model_defaults.cron_tz == zoneinfo.ZoneInfo("America/Los_Angeles")
955+
assert config.model_defaults.cron_tz.key == "America/Los_Angeles"
956+
957+
958+
def test_gateway_model_defaults_cron_tz(tmp_path):
959+
"""Test that cron_tz can be set in gateway-specific model_defaults."""
960+
import zoneinfo
961+
962+
global_defaults = ModelDefaultsConfig(
963+
dialect="snowflake", owner="foo", cron="@daily", cron_tz="UTC"
964+
)
965+
gateway_defaults = ModelDefaultsConfig(
966+
dialect="duckdb", cron_tz="America/New_York"
967+
)
968+
969+
config = Config(
970+
gateways={
971+
"duckdb": GatewayConfig(
972+
connection=DuckDBConnectionConfig(database="db.db"),
973+
model_defaults=gateway_defaults,
974+
)
975+
},
976+
model_defaults=global_defaults,
977+
default_gateway="duckdb",
978+
)
979+
980+
ctx = Context(paths=tmp_path, config=config, gateway="duckdb")
981+
982+
expected = ModelDefaultsConfig(
983+
dialect="duckdb", owner="foo", cron="@daily", cron_tz="America/New_York"
984+
)
985+
986+
assert ctx.config.model_defaults == expected
987+
# Also verify the cron_tz is a ZoneInfo object
988+
assert isinstance(ctx.config.model_defaults.cron_tz, zoneinfo.ZoneInfo)
989+
assert ctx.config.model_defaults.cron_tz.key == "America/New_York"
990+
991+
933992
def test_redshift_merge_flag(tmp_path, mocker: MockerFixture):
934993
config_path = tmp_path / "config_redshift_merge.yaml"
935994
with open(config_path, "w", encoding="utf-8") as fd:

0 commit comments

Comments
 (0)