Fix: Make sure start_at / end_at are strings when storing an environment record in state sync (#870)

Author: izeigerman

sqlmesh/core/state_sync/engine_adapter.py

@@ -43,7 +43,7 @@
 from sqlmesh.core.snapshot.definition import _parents_from_model
 from sqlmesh.core.state_sync.base import SCHEMA_VERSION, StateSync, Versions
 from sqlmesh.core.state_sync.common import CommonStateSyncMixin, transactional
-from sqlmesh.utils.date import now_timestamp
+from sqlmesh.utils.date import now_timestamp, time_like_to_str
 from sqlmesh.utils.errors import SQLMeshError
 
 logger = logging.getLogger(__name__)
@@ -617,8 +617,8 @@ def _environment_to_df(environment: Environment) -> pd.DataFrame:
             {
                 "name": environment.name,
                 "snapshots": json.dumps([snapshot.dict() for snapshot in environment.snapshots]),
-                "start_at": environment.start_at,
-                "end_at": environment.end_at,
+                "start_at": time_like_to_str(environment.start_at),
+                "end_at": time_like_to_str(environment.end_at) if environment.end_at else None,
                 "plan_id": environment.plan_id,
                 "previous_plan_id": environment.previous_plan_id,
                 "expiration_ts": environment.expiration_ts,

sqlmesh/utils/date.py

@@ -257,3 +257,11 @@ def validate_date_range(
         raise ValueError(
             f"Start date / time ({start}) can't be greater than end date / time ({end})"
         )
+
+
+def time_like_to_str(time_like: TimeLike) -> str:
+    if isinstance(time_like, str):
+        return time_like
+    if is_date(time_like):
+        return to_ds(time_like)
+    return to_datetime(time_like).isoformat()

tests/core/test_state_sync.py

@@ -551,6 +551,37 @@ def test_delete_expired_environments(state_sync: EngineAdapterStateSync, make_sn
     assert state_sync.get_environment(env_b.name) == env_b
 
 
+def test_environment_start_as_timestamp(
+    state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
+):
+    snapshot = make_snapshot(
+        SqlModel(
+            name="a",
+            query=parse_one("select a, ds"),
+        ),
+    )
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    state_sync.push_snapshots([snapshot])
+
+    now_ts = now_timestamp()
+
+    env = Environment(
+        name="test_environment_a",
+        snapshots=[snapshot.table_info],
+        start_at=now_ts,
+        end_at=None,
+        plan_id="test_plan_id",
+        previous_plan_id="test_plan_id",
+        expiration_ts=now_ts - 1000,
+    )
+    state_sync.promote(env)
+
+    stored_env = state_sync.get_environment(env.name)
+    assert stored_env
+    assert stored_env.start_at == to_datetime(now_ts).isoformat()
+
+
 def test_missing_intervals(sushi_context_pre_scheduling: Context) -> None:
     sushi_context = sushi_context_pre_scheduling
     state_sync = sushi_context.state_reader