Skip to content

Commit f22cc38

Browse files
committed
feat: add cron_tz to model_defaults
Signed-off-by: lafirm <136463254+lafirm@users.noreply.github.com>
1 parent ee960d1 commit f22cc38

4 files changed

Lines changed: 88 additions & 22 deletions

File tree

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
@@ -964,6 +964,65 @@ def test_gateway_model_defaults(tmp_path):
964964
assert ctx.config.model_defaults == expected
965965

966966

967+
def test_model_defaults_cron_tz(tmp_path):
968+
"""Test that cron_tz can be set in model_defaults."""
969+
import zoneinfo
970+
971+
config_path = tmp_path / "config_model_defaults_cron_tz.yaml"
972+
with open(config_path, "w", encoding="utf-8") as fd:
973+
fd.write(
974+
"""
975+
model_defaults:
976+
dialect: duckdb
977+
cron: '@daily'
978+
cron_tz: 'America/Los_Angeles'
979+
"""
980+
)
981+
982+
config = load_config_from_paths(
983+
Config,
984+
project_paths=[config_path],
985+
)
986+
987+
assert config.model_defaults.cron == "@daily"
988+
assert config.model_defaults.cron_tz == zoneinfo.ZoneInfo("America/Los_Angeles")
989+
assert config.model_defaults.cron_tz.key == "America/Los_Angeles"
990+
991+
992+
def test_gateway_model_defaults_cron_tz(tmp_path):
993+
"""Test that cron_tz can be set in gateway-specific model_defaults."""
994+
import zoneinfo
995+
996+
global_defaults = ModelDefaultsConfig(
997+
dialect="snowflake", owner="foo", cron="@daily", cron_tz="UTC"
998+
)
999+
gateway_defaults = ModelDefaultsConfig(
1000+
dialect="duckdb", cron_tz="America/New_York"
1001+
)
1002+
1003+
config = Config(
1004+
gateways={
1005+
"duckdb": GatewayConfig(
1006+
connection=DuckDBConnectionConfig(database="db.db"),
1007+
model_defaults=gateway_defaults,
1008+
)
1009+
},
1010+
model_defaults=global_defaults,
1011+
default_gateway="duckdb",
1012+
)
1013+
1014+
ctx = Context(paths=tmp_path, config=config, gateway="duckdb")
1015+
1016+
expected = ModelDefaultsConfig(
1017+
dialect="duckdb", owner="foo", cron="@daily", cron_tz="America/New_York"
1018+
)
1019+
1020+
assert ctx.config.model_defaults == expected
1021+
# Also verify the cron_tz is a ZoneInfo object
1022+
assert isinstance(ctx.config.model_defaults.cron_tz, zoneinfo.ZoneInfo)
1023+
assert ctx.config.model_defaults.cron_tz.key == "America/New_York"
1024+
1025+
9671026
def test_redshift_merge_flag(tmp_path, mocker: MockerFixture):
9681027
config_path = tmp_path / "config_redshift_merge.yaml"
9691028
with open(config_path, "w", encoding="utf-8") as fd:

0 commit comments

Comments
 (0)