Skip to content

Commit 40f64db

Browse files
Merge branch 'main' into feat/local_tz
2 parents 1233d63 + 87e179b commit 40f64db

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: 26 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -260,6 +260,30 @@ def dbt_fqn(self) -> t.Optional[str]:
260260
}
261261

262262

263+
def _cron_tz_validator(cls: t.Type, v: t.Any) -> t.Optional[zoneinfo.ZoneInfo]:
264+
if not v or v == "UTC":
265+
return None
266+
267+
v = str_or_exp_to_str(v)
268+
269+
try:
270+
return zoneinfo.ZoneInfo(v)
271+
except Exception as e:
272+
available_timezones = zoneinfo.available_timezones()
273+
274+
if available_timezones:
275+
raise ConfigError(f"{e}. {v} must be in {available_timezones}.")
276+
else:
277+
raise ConfigError(
278+
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."
279+
)
280+
281+
return None
282+
283+
284+
cron_tz_validator = field_validator("cron_tz", mode="before")(_cron_tz_validator)
285+
286+
263287
class _Node(DbtInfoMixin, PydanticModel):
264288
"""
265289
Node is the core abstraction for entity that can be executed within the scheduler.
@@ -302,6 +326,8 @@ class _Node(DbtInfoMixin, PydanticModel):
302326
_croniter: t.Optional[CroniterCache] = None
303327
__inferred_interval_unit: t.Optional[IntervalUnit] = None
304328

329+
_cron_tz_validator = cron_tz_validator
330+
305331
def __str__(self) -> str:
306332
path = f": {self._path.name}" if self._path else ""
307333
return f"{self.__class__.__name__}<{self.name}{path}>"
@@ -328,27 +354,6 @@ def _name_validator(cls, v: t.Any) -> t.Optional[str]:
328354
return v.meta["sql"]
329355
return str(v)
330356

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-
352357
@field_validator("start", "end", mode="before")
353358
@classmethod
354359
def _date_validator(cls, v: t.Any) -> t.Optional[TimeLike]:

tests/core/test_config.py

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -975,6 +975,63 @@ def test_gateway_model_defaults(tmp_path):
975975
assert ctx.config.model_defaults == expected
976976

977977

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

0 commit comments

Comments
 (0)