chore: add interval inclusive/exclusive method (#2543)

* chore: add interval inclusive/exclusive method

* add add_snapshot_intervals method

Author: eakmanrq

sqlmesh/core/state_sync/base.py

@@ -18,7 +18,7 @@
     SnapshotTableCleanupTask,
     SnapshotTableInfo,
 )
-from sqlmesh.core.snapshot.definition import Interval
+from sqlmesh.core.snapshot.definition import Interval, SnapshotIntervals
 from sqlmesh.utils import major_minor
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import SQLMeshError
@@ -303,26 +303,6 @@ def invalidate_environment(self, name: str) -> None:
             name: The name of the environment to invalidate.
         """
 
-    @abc.abstractmethod
-    def add_interval(
-        self,
-        snapshot: Snapshot,
-        start: TimeLike,
-        end: TimeLike,
-        is_dev: bool = False,
-    ) -> None:
-        """Add an interval to a snapshot and sync it to the store.
-
-        Snapshots must be pushed before adding intervals to them.
-
-        Args:
-            snapshot: The snapshot like object to add an interval to.
-            start: The start of the interval to add.
-            end: The end of the interval to add.
-            is_dev: Indicates whether the given interval is being added while in
-                development mode.is_dev.
-        """
-
     @abc.abstractmethod
     def remove_interval(
         self,
@@ -429,6 +409,42 @@ def migrate(
     def rollback(self) -> None:
         """Rollback to previous backed up state."""
 
+    @abc.abstractmethod
+    def _add_snapshot_intervals(self, snapshot_intervals: SnapshotIntervals) -> None:
+        """Add snapshot intervals to state
+
+        Args:
+            snapshot_intervals: The snapshot intervals to add.
+        """
+
+    def add_interval(
+        self,
+        snapshot: Snapshot,
+        start: TimeLike,
+        end: TimeLike,
+        is_dev: bool = False,
+    ) -> None:
+        """Add an interval to a snapshot and sync it to the store.
+
+        Args:
+            snapshot: The snapshot like object to add an interval to.
+            start: The start of the interval to add.
+            end: The end of the interval to add.
+            is_dev: Indicates whether the given interval is being added while in development mode
+        """
+        start_ts, end_ts = snapshot.inclusive_exclusive(start, end, strict=False)
+        if not snapshot.version:
+            raise SQLMeshError("Snapshot version must be set to add an interval.")
+        intervals = [(start_ts, end_ts)]
+        snapshot_intervals = SnapshotIntervals(
+            name=snapshot.name,
+            identifier=snapshot.identifier,
+            version=snapshot.version,
+            intervals=intervals if not is_dev else [],
+            dev_intervals=intervals if is_dev else [],
+        )
+        self._add_snapshot_intervals(snapshot_intervals)
+
 
 class DelegatingStateSync(StateSync):
     def __init__(self, state_sync: StateSync) -> None:

sqlmesh/core/state_sync/cache.py

@@ -11,7 +11,7 @@
     SnapshotInfoLike,
     SnapshotTableCleanupTask,
 )
-from sqlmesh.core.snapshot.definition import Interval
+from sqlmesh.core.snapshot.definition import Interval, SnapshotIntervals
 from sqlmesh.core.state_sync.base import DelegatingStateSync, StateSync
 from sqlmesh.utils.date import TimeLike, now_timestamp
 
@@ -129,15 +129,9 @@ def delete_expired_snapshots(self) -> t.List[SnapshotTableCleanupTask]:
         self.snapshot_cache.clear()
         return self.state_sync.delete_expired_snapshots()
 
-    def add_interval(
-        self,
-        snapshot: Snapshot,
-        start: TimeLike,
-        end: TimeLike,
-        is_dev: bool = False,
-    ) -> None:
-        self.snapshot_cache.pop(snapshot.snapshot_id, None)
-        self.state_sync.add_interval(snapshot, start, end, is_dev)
+    def _add_snapshot_intervals(self, snapshot_intervals: SnapshotIntervals) -> None:
+        self.snapshot_cache.pop(snapshot_intervals.snapshot_id, None)
+        self.state_sync._add_snapshot_intervals(snapshot_intervals)
 
     def remove_interval(
         self,

sqlmesh/core/state_sync/engine_adapter.py

@@ -624,27 +624,48 @@ def add_interval(
         end: TimeLike,
         is_dev: bool = False,
     ) -> None:
-        start_ts, end_ts = snapshot.inclusive_exclusive(start, end, strict=False)
-        if start_ts >= end_ts:
-            logger.info(
-                "Skipping partial interval (%s, %s) for snapshot %s",
-                start,
-                end,
-                snapshot.snapshot_id,
-            )
-            return
+        super().add_interval(snapshot, start, end, is_dev)
 
-        logger.info(
-            "Adding %s (%s, %s) for snapshot %s",
-            "dev interval" if is_dev else "interval",
-            start_ts,
-            end_ts,
-            snapshot.snapshot_id,
-        )
+    @transactional()
+    def _add_snapshot_intervals(self, snapshot_intervals: SnapshotIntervals) -> None:
+        def remove_partial_intervals(
+            intervals: t.List[Interval], snapshot_id: SnapshotId, *, is_dev: bool
+        ) -> t.List[Interval]:
+            results = []
+            for start_ts, end_ts in intervals:
+                if start_ts < end_ts:
+                    logger.info(
+                        "Adding %s (%s, %s) for snapshot %s",
+                        "dev interval" if is_dev else "interval",
+                        start_ts,
+                        end_ts,
+                        snapshot_id,
+                    )
+                    results.append((start_ts, end_ts))
+                else:
+                    logger.info(
+                        "Skipping partial interval (%s, %s) for snapshot %s",
+                        start_ts,
+                        end_ts,
+                        snapshot_intervals.snapshot_id,
+                    )
+            return results
 
+        snapshot_intervals = snapshot_intervals.copy(
+            update={
+                "intervals": remove_partial_intervals(
+                    snapshot_intervals.intervals, snapshot_intervals.snapshot_id, is_dev=False
+                ),
+                "dev_intervals": remove_partial_intervals(
+                    snapshot_intervals.dev_intervals, snapshot_intervals.snapshot_id, is_dev=True
+                ),
+            }
+        )
+        if not snapshot_intervals.intervals and not snapshot_intervals.dev_intervals:
+            return
         self.engine_adapter.insert_append(
             self.intervals_table,
-            _intervals_to_df([(snapshot, (start_ts, end_ts))], is_dev, False),
+            _snapshot_interval_to_df(snapshot_intervals, is_removed=False),
             columns_to_types=self._interval_columns_to_types,
         )
 
@@ -1312,6 +1333,27 @@ def _intervals_to_df(
     )
 
 
+def _snapshot_interval_to_df(
+    snapshot_intervals: SnapshotIntervals,
+    is_removed: bool = False,
+) -> pd.DataFrame:
+    return pd.DataFrame(
+        [
+            _interval_to_df(
+                snapshot_intervals,
+                start_ts,
+                end_ts,
+                is_dev=is_dev,
+                is_removed=is_removed,
+            )
+            for is_dev in (False, True)
+            for start_ts, end_ts in getattr(
+                snapshot_intervals, "dev_intervals" if is_dev else "intervals"
+            )
+        ]
+    )
+
+
 def _interval_to_df(
     snapshot: t.Union[SnapshotInfoLike, SnapshotIntervals],
     start_ts: int,

sqlmesh/schedulers/airflow/state_sync.py

@@ -12,7 +12,7 @@
     SnapshotInfoLike,
     SnapshotTableCleanupTask,
 )
-from sqlmesh.core.snapshot.definition import Interval
+from sqlmesh.core.snapshot.definition import Interval, SnapshotIntervals
 from sqlmesh.core.state_sync import StateSync, Versions
 from sqlmesh.core.state_sync.base import PromotionResult
 from sqlmesh.schedulers.airflow.client import AirflowClient
@@ -219,6 +219,9 @@ def add_interval(
         """
         raise NotImplementedError("Adding intervals is not supported by the Airflow state sync.")
 
+    def _add_snapshot_intervals(self, snapshot_intervals: SnapshotIntervals) -> None:
+        raise NotImplementedError("Adding intervals is not supported by the Airflow state sync.")
+
     def remove_interval(
         self,
         snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],