feat!: add cron_tz to allow scheduling according to time zones

Author: tobymao

sqlmesh/core/audit/definition.py

@@ -270,6 +270,8 @@ def metadata_hash(self) -> str:
                 *sorted(self.tags),
                 str(self.sorted_python_env),
                 self.stamp,
+                self.cron,
+                self.cron_tz.key if self.cron_tz else None,
             ]
 
             query = self.render_audit_query() or self.query

sqlmesh/core/model/definition.py

@@ -1086,6 +1086,7 @@ def metadata_hash(self) -> str:
                 self.description,
                 json.dumps(self.column_descriptions, sort_keys=True),
                 self.cron,
+                self.cron_tz.key if self.cron_tz else None,
                 str(self.start) if self.start else None,
                 str(self.end) if self.end else None,
                 str(self.retention) if self.retention else None,

sqlmesh/core/node.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing as t
+import zoneinfo
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
@@ -177,6 +178,7 @@ class _Node(PydanticModel):
             the date from the scheduler will be used
         cron: A cron string specifying how often the node should be run, leveraging the
             [croniter](https://github.com/kiorky/croniter) library.
+        cron_tz: Time zone for the cron, defaults to utc, [IANA time zones](https://docs.python.org/3/library/zoneinfo.html).
         interval_unit: The duration of an interval for the node. By default, it is computed from the cron expression.
         tags: A list of tags that can be used to filter nodes.
         stamp: An optional arbitrary string sequence used to create new node versions without making
@@ -190,6 +192,7 @@ class _Node(PydanticModel):
     start: t.Optional[TimeLike] = None
     end: t.Optional[TimeLike] = None
     cron: SQLGlotCron = "@daily"
+    cron_tz: t.Optional[zoneinfo.ZoneInfo] = None
     interval_unit_: t.Optional[IntervalUnit] = Field(alias="interval_unit", default=None)
     tags: t.List[str] = []
     stamp: t.Optional[str] = None
@@ -226,6 +229,22 @@ def _name_validator(cls, v: t.Any) -> t.Optional[str]:
             return v.meta["sql"]
         return str(v)
 
+    @field_validator("cron_tz", mode="before")
+    def _cron_tz_validator(cls, v: t.Any) -> t.Optional[zoneinfo.ZoneInfo]:
+        if not v or v == "UTC":
+            return None
+
+        v = str_or_exp_to_str(v)
+
+        try:
+            return zoneinfo.ZoneInfo(v)
+        except Exception as e:
+            raise ConfigError(
+                f"{e}. {v} must be in {zoneinfo.available_timezones()} or 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."
+            )
+
+        return None
+
     @field_validator("start", "end", mode="before")
     @classmethod
     def _date_validator(cls, v: t.Any) -> t.Optional[TimeLike]:
@@ -319,7 +338,7 @@ def croniter(self, value: TimeLike) -> CroniterCache:
         if self._croniter is None:
             self._croniter = CroniterCache(self.cron, value)
         else:
-            self._croniter.curr = to_datetime(value)
+            self._croniter.curr = to_datetime(value, tz=self.cron_tz)
         return self._croniter
 
     def cron_next(self, value: TimeLike, estimate: bool = False) -> datetime:

sqlmesh/migrations/v0071_add_dev_version_to_intervals.py

@@ -31,10 +31,10 @@ def migrate(state_sync, **kwargs):  # type: ignore
     )
     engine_adapter.execute(alter_table_exp)
 
-    used_dev_versions: t.Set[t.Tuple[str, str]] = set()
-    used_versions: t.Set[t.Tuple[str, str]] = set()
-    used_snapshot_ids: t.Set[t.Tuple[str, str]] = set()
-    snapshot_ids_to_dev_versions: t.Dict[t.Tuple[str, str], str] = {}
+    used_dev_versions = set()
+    used_versions = set()
+    used_snapshot_ids = set()
+    snapshot_ids_to_dev_versions = {}
 
     _migrate_snapshots(
         engine_adapter,

sqlmesh/utils/cron.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import typing as t
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, tzinfo
 from functools import lru_cache
 
 from croniter import croniter
@@ -34,21 +34,22 @@ def interval_seconds(cron: str) -> int:
 
 
 class CroniterCache:
-    def __init__(self, cron: str, time: t.Optional[TimeLike] = None):
+    def __init__(self, cron: str, time: t.Optional[TimeLike] = None, tz: t.Optional[tzinfo] = None):
         self.cron = cron
-        self.curr: datetime = to_datetime(now() if time is None else time)
+        self.tz = tz
+        self.curr: datetime = to_datetime(now() if time is None else time, tz=self.tz)
         self.interval_seconds = interval_seconds(self.cron)
 
     def get_next(self, estimate: bool = False) -> datetime:
         if estimate and self.interval_seconds:
             self.curr = self.curr + timedelta(seconds=self.interval_seconds)
         else:
-            self.curr = to_datetime(croniter(self.cron, self.curr).get_next() * 1000)
+            self.curr = to_datetime(croniter(self.cron, self.curr).get_next() * 1000, tz=self.tz)
         return self.curr
 
     def get_prev(self, estimate: bool = False) -> datetime:
         if estimate and self.interval_seconds:
             self.curr = self.curr - timedelta(seconds=self.interval_seconds)
         else:
-            self.curr = to_datetime(croniter(self.cron, self.curr).get_prev() * 1000)
+            self.curr = to_datetime(croniter(self.cron, self.curr).get_prev() * 1000, tz=self.tz)
         return self.curr

sqlmesh/utils/date.py

@@ -5,14 +5,13 @@
 import typing as t
 import warnings
 
-from pandas.api.types import is_datetime64_any_dtype  # type: ignore
-
-from datetime import date, datetime, timedelta, timezone
+from datetime import date, datetime, timedelta, timezone, tzinfo
 
 import dateparser
 import pandas as pd
 from dateparser import freshness_date_parser as freshness_date_parser_module
 from dateparser.freshness_date_parser import freshness_date_parser
+from pandas.api.types import is_datetime64_any_dtype  # type: ignore
 from sqlglot import exp
 
 from sqlmesh.utils import ttl_cache
@@ -149,19 +148,21 @@ def to_datetime(
     value: TimeLike,
     relative_base: t.Optional[datetime] = None,
     check_categorical_relative_expression: bool = True,
+    tz: t.Optional[tzinfo] = None,
 ) -> datetime:
     """Converts a value into a UTC datetime object.
 
     Args:
         value: A variety of date formats. If the value is number-like, it is assumed to be millisecond epochs.
         relative_base: The datetime to reference for time expressions that are using relative terms.
         check_categorical_relative_expression: If True, takes into account the relative expressions that are categorical.
+        tz: Timezone to convert datetime to, defaults to utc
 
     Raises:
         ValueError if value cannot be converted to a datetime.
 
     Returns:
-        A datetime object with tz utc.
+        A datetime object with tz (default UTC).
     """
     if isinstance(value, datetime):
         dt: t.Optional[datetime] = value
@@ -198,9 +199,11 @@ def to_datetime(
     if dt is None:
         raise ValueError(f"Could not convert `{value}` to datetime.")
 
+    tz = tz or UTC
+
     if dt.tzinfo:
-        return dt if dt.tzinfo == UTC else dt.astimezone(UTC)
-    return dt.replace(tzinfo=UTC)
+        return dt if dt.tzinfo == tz else dt.astimezone(tz)
+    return dt.replace(tzinfo=tz)
 
 
 def to_date(value: TimeLike, relative_base: t.Optional[datetime] = None) -> date:

sqlmesh/utils/pydantic.py

@@ -2,6 +2,7 @@
 
 import json
 import typing as t
+from datetime import tzinfo
 
 import pydantic
 from pydantic import ValidationInfo as ValidationInfo
@@ -72,6 +73,7 @@ class PydanticModel(pydantic.BaseModel):
             exp.Tuple: _expression_encoder,
             AuditQueryTypes: _expression_encoder,  # type: ignore
             ModelQueryTypes: _expression_encoder,  # type: ignore
+            tzinfo: lambda tz: tz.key,
         },
         arbitrary_types_allowed=True,
         extra="forbid",

tests/core/test_integration.py

@@ -4884,6 +4884,58 @@ def test_plan_production_environment_statements(tmp_path: Path):
     assert environment_statements[0].python_env["__sqlmesh__vars__"].payload == "{'var_5': 5}"
 
 
+@time_machine.travel("2025-03-08 00:00:00 UTC")
+def test_tz(init_and_plan_context):
+    context, _ = init_and_plan_context("examples/sushi")
+
+    model = context.get_model("sushi.waiter_revenue_by_day")
+    context.upsert_model(
+        SqlModel.parse_obj(
+            {**model.dict(), "cron_tz": "America/Los_Angeles", "start": "2025-03-07"}
+        )
+    )
+
+    def assert_intervals(plan, intervals):
+        assert (
+            next(
+                intervals.intervals
+                for intervals in plan.missing_intervals
+                if intervals.snapshot_id.name == model.fqn
+            )
+            == intervals
+        )
+
+    plan = context.plan_builder("prod", skip_tests=True).build()
+
+    assert_intervals(plan, [(to_timestamp("2025-03-07"), to_timestamp("2025-03-08"))])
+
+    with time_machine.travel("2025-03-09 07:00:00 UTC"):
+        plan = context.plan_builder("prod", skip_tests=True).build()
+
+        assert_intervals(
+            plan,
+            [
+                (to_timestamp("2025-03-07"), to_timestamp("2025-03-08")),
+            ],
+        )
+
+    with time_machine.travel("2025-03-09 08:00:00 UTC"):
+        plan = context.plan_builder("prod", skip_tests=True).build()
+
+        assert_intervals(
+            plan,
+            [
+                (to_timestamp("2025-03-07"), to_timestamp("2025-03-08")),
+                (to_timestamp("2025-03-08"), to_timestamp("2025-03-09")),
+            ],
+        )
+
+        context.apply(plan)
+
+        plan = context.plan_builder("prod", skip_tests=True).build()
+        assert not plan.missing_intervals
+
+
 def apply_to_environment(
     context: Context,
     environment: str,

tests/core/test_model.py

@@ -2113,7 +2113,7 @@ def python_model_prop(context, **kwargs):
         ),
     }
 
-    snapshot: Snapshot = make_snapshot(m)
+    snapshot = make_snapshot(m)
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
     # Rendering the properties will result to a TRANSIENT creatable_type and the removal of the conditional prop
@@ -3674,7 +3674,7 @@ def test_conditional_physical_properties(make_snapshot):
     )
 
     # substitution occurs at runtime
-    snapshot: Snapshot = make_snapshot(full_model)
+    snapshot = make_snapshot(full_model)
     snapshot_view: Snapshot = make_snapshot(view_model)
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
@@ -3775,7 +3775,7 @@ def test_model_defaults_macros(make_snapshot):
         variables={"gateway": "dev", "create_type": "SECURE"},
     )
 
-    snapshot: Snapshot = make_snapshot(model)
+    snapshot = make_snapshot(model)
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
     # Validate rendering of model defaults
@@ -3881,7 +3881,7 @@ def python_model_prop_macro(context, **kwargs):
     # Validate disabling attribute dynamically
     assert not m.storage_format
 
-    snapshot: Snapshot = make_snapshot(m)
+    snapshot = make_snapshot(m)
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
     # Ensure properties are not rendered at load time

tests/core/test_snapshot_evaluator.py

@@ -3192,8 +3192,8 @@ def insert(
 def test_custom_materialization_strategy_with_custom_properties(adapter_mock, make_snapshot):
     custom_insert_kind = None
 
-    class TestCustomKind(CustomKind):  # type: ignore[no-untyped-def]
-        _primary_key: t.List[exp.Expression]
+    class TestCustomKind(CustomKind):
+        _primary_key: t.List[exp.Expression]  # type: ignore[no-untyped-def]
 
         @model_validator(mode="after")
         def _validate_model(self) -> Self: