Merge branch 'main' into feat-bigquery-reservations

Author: scottypate

.circleci/continue_config.yml

@@ -93,9 +93,9 @@ jobs:
       - run:
           name: Run linters and code style checks
           command: make py-style
-      - run:
-          name: Exercise the benchmarks
-          command: make benchmark-ci
+      # - run:
+      #     name: Exercise the benchmarks
+      #     command: make benchmark-ci
       - run:
           name: Run cicd tests
           command: make cicd-test

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "requests",
     "rich[jupyter]",
     "ruamel.yaml",
-    "sqlglot[rs]~=27.27.0",
+    "sqlglot[rs]~=27.28.0",
     "tenacity",
     "time-machine",
     "json-stream"

sqlmesh/core/context.py

@@ -154,6 +154,8 @@
     )
     from sqlmesh.core.snapshot import Node
 
+    from sqlmesh.core.snapshot.definition import Intervals
+
     ModelOrSnapshot = t.Union[str, Model, Snapshot]
     NodeOrSnapshot = t.Union[str, Model, StandaloneAudit, Snapshot]
 
@@ -276,6 +278,7 @@ def __init__(
         default_dialect: t.Optional[str] = None,
         default_catalog: t.Optional[str] = None,
         is_restatement: t.Optional[bool] = None,
+        parent_intervals: t.Optional[Intervals] = None,
         variables: t.Optional[t.Dict[str, t.Any]] = None,
         blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
     ):
@@ -287,6 +290,7 @@ def __init__(
         self._variables = variables or {}
         self._blueprint_variables = blueprint_variables or {}
         self._is_restatement = is_restatement
+        self._parent_intervals = parent_intervals
 
     @property
     def default_dialect(self) -> t.Optional[str]:
@@ -315,6 +319,10 @@ def gateway(self) -> t.Optional[str]:
     def is_restatement(self) -> t.Optional[bool]:
         return self._is_restatement
 
+    @property
+    def parent_intervals(self) -> t.Optional[Intervals]:
+        return self._parent_intervals
+
     def var(self, var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
         """Returns a variable value."""
         return self._variables.get(var_name.lower(), default)
@@ -2271,6 +2279,7 @@ def audit(
                 snapshot=snapshot,
                 start=start,
                 end=end,
+                execution_time=execution_time,
                 snapshots=self.snapshots,
             ):
                 audit_id = f"{audit_result.audit.name}"

sqlmesh/core/engine_adapter/base.py

@@ -551,11 +551,13 @@ def replace_query(
                     target_table,
                     source_queries,
                     target_columns_to_types,
+                    **kwargs,
                 )
         return self._insert_overwrite_by_condition(
             target_table,
             source_queries,
             target_columns_to_types,
+            **kwargs,
         )
 
     def create_index(
@@ -1614,7 +1616,7 @@ def _insert_overwrite_by_time_partition(
         **kwargs: t.Any,
     ) -> None:
         return self._insert_overwrite_by_condition(
-            table_name, source_queries, target_columns_to_types, where
+            table_name, source_queries, target_columns_to_types, where, **kwargs
         )
 
     def _values_to_sql(

sqlmesh/core/engine_adapter/mssql.py

@@ -423,7 +423,9 @@ def _insert_overwrite_by_condition(
         insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None,
         **kwargs: t.Any,
     ) -> None:
-        if not where or where == exp.true():
+        # note that this is passed as table_properties here rather than physical_properties
+        use_merge_strategy = kwargs.get("table_properties", {}).get("mssql_merge_exists")
+        if (not where or where == exp.true()) and not use_merge_strategy:
             # this is a full table replacement, call the base strategy to do DELETE+INSERT
             # which will result in TRUNCATE+INSERT due to how we have overridden self.delete_from()
             return EngineAdapter._insert_overwrite_by_condition(
@@ -436,7 +438,7 @@ def _insert_overwrite_by_condition(
                 **kwargs,
             )
 
-        # For actual conditional overwrites, use MERGE from InsertOverwriteWithMergeMixin
+        # For conditional overwrites or when mssql_merge_exists is set use MERGE
         return super()._insert_overwrite_by_condition(
             table_name=table_name,
             source_queries=source_queries,

sqlmesh/core/lineage.py

@@ -66,6 +66,7 @@ def lineage(
         scope=scope,
         trim_selects=trim_selects,
         dialect=model.dialect,
+        copy=False,
     )
 
 

sqlmesh/core/plan/definition.py

@@ -63,7 +63,7 @@ class Plan(PydanticModel, frozen=True):
     restatements: t.Dict[SnapshotId, Interval]
     """
     All models being restated, which are typically the explicitly selected ones + their downstream dependencies.
-    
+
     Note that dev previews are also considered restatements, so :selected_models_to_restate can be empty
     while :restatements is still populated with dev previews
     """
@@ -213,8 +213,8 @@ def environment(self) -> Environment:
 
         snapshots_by_name = self.context_diff.snapshots_by_name
         snapshots = [s.table_info for s in self.snapshots.values()]
-        promoted_snapshot_ids = None
-        if self.is_dev and not self.include_unmodified:
+        promotable_snapshot_ids = None
+        if self.is_dev:
             if self.selected_models_to_backfill is not None:
                 # Only promote models that have been explicitly selected for backfill.
                 promotable_snapshot_ids = {
@@ -225,12 +225,14 @@ def environment(self) -> Environment:
                         if m in snapshots_by_name
                     ],
                 }
-            else:
+            elif not self.include_unmodified:
                 promotable_snapshot_ids = self.context_diff.promotable_snapshot_ids.copy()
 
-            promoted_snapshot_ids = [
-                s.snapshot_id for s in snapshots if s.snapshot_id in promotable_snapshot_ids
-            ]
+        promoted_snapshot_ids = (
+            [s.snapshot_id for s in snapshots if s.snapshot_id in promotable_snapshot_ids]
+            if promotable_snapshot_ids is not None
+            else None
+        )
 
         previous_finalized_snapshots = (
             self.context_diff.environment_snapshots

sqlmesh/core/renderer.py

@@ -196,7 +196,14 @@ def _resolve_table(table: str | exp.Table) -> str:
             **kwargs,
         }
 
+        if this_model:
+            render_kwargs["this_model"] = this_model
+
+        macro_evaluator.locals.update(render_kwargs)
+
         variables = kwargs.pop("variables", {})
+        if variables:
+            macro_evaluator.locals.setdefault(c.SQLMESH_VARS, {}).update(variables)
 
         expressions = [self._expression]
         if isinstance(self._expression, d.Jinja):
@@ -268,14 +275,6 @@ def _resolve_table(table: str | exp.Table) -> str:
                             f"Could not parse the rendered jinja at '{self._path}'.\n{ex}"
                         ) from ex
 
-        if this_model:
-            render_kwargs["this_model"] = this_model
-
-        macro_evaluator.locals.update(render_kwargs)
-
-        if variables:
-            macro_evaluator.locals.setdefault(c.SQLMESH_VARS, {}).update(variables)
-
         for definition in self._macro_definitions:
             try:
                 macro_evaluator.evaluate(definition)

sqlmesh/core/scheduler.py

@@ -352,7 +352,7 @@ def batch_intervals(
             )
             for snapshot, intervals in merged_intervals.items()
         }
-        snapshot_batches = {}
+        snapshot_batches: t.Dict[Snapshot, Intervals] = {}
         all_unready_intervals: t.Dict[str, set[Interval]] = {}
         for snapshot_id in dag:
             if snapshot_id not in snapshot_intervals:
@@ -364,13 +364,22 @@ def batch_intervals(
 
             adapter = self.snapshot_evaluator.get_adapter(snapshot.model_gateway)
 
+            parent_intervals: Intervals = []
+            for parent_id in snapshot.parents:
+                parent_snapshot, _ = snapshot_intervals.get(parent_id, (None, []))
+                if not parent_snapshot or parent_snapshot.is_external:
+                    continue
+
+                parent_intervals.extend(snapshot_batches[parent_snapshot])
+
             context = ExecutionContext(
                 adapter,
                 self.snapshots_by_name,
                 deployability_index,
                 default_dialect=adapter.dialect,
                 default_catalog=self.default_catalog,
                 is_restatement=is_restatement,
+                parent_intervals=parent_intervals,
             )
 
             intervals = self._check_ready_intervals(
@@ -538,6 +547,10 @@ def run_node(node: SchedulingUnit) -> None:
                             execution_time=execution_time,
                         )
                     else:
+                        # If batch_index > 0, then the target table must exist since the first batch would have created it
+                        target_table_exists = (
+                            snapshot.snapshot_id not in snapshots_to_create or node.batch_index > 0
+                        )
                         audit_results = self.evaluate(
                             snapshot=snapshot,
                             environment_naming_info=environment_naming_info,
@@ -548,7 +561,7 @@ def run_node(node: SchedulingUnit) -> None:
                             batch_index=node.batch_index,
                             allow_destructive_snapshots=allow_destructive_snapshots,
                             allow_additive_snapshots=allow_additive_snapshots,
-                            target_table_exists=snapshot.snapshot_id not in snapshots_to_create,
+                            target_table_exists=target_table_exists,
                             selected_models=selected_models,
                         )
 
@@ -646,6 +659,7 @@ def _dag(
         }
         snapshots_to_create = snapshots_to_create or set()
         original_snapshots_to_create = snapshots_to_create.copy()
+        upstream_dependencies_cache: t.Dict[SnapshotId, t.Set[SchedulingUnit]] = {}
 
         snapshot_dag = snapshot_dag or snapshots_to_dag(batches)
         dag = DAG[SchedulingUnit]()
@@ -657,12 +671,15 @@ def _dag(
             snapshot = self.snapshots_by_name[snapshot_id.name]
             intervals = intervals_per_snapshot.get(snapshot.name, [])
 
-            upstream_dependencies: t.List[SchedulingUnit] = []
+            upstream_dependencies: t.Set[SchedulingUnit] = set()
 
             for p_sid in snapshot.parents:
-                upstream_dependencies.extend(
+                upstream_dependencies.update(
                     self._find_upstream_dependencies(
-                        p_sid, intervals_per_snapshot, original_snapshots_to_create
+                        p_sid,
+                        intervals_per_snapshot,
+                        original_snapshots_to_create,
+                        upstream_dependencies_cache,
                     )
                 )
 
@@ -713,29 +730,42 @@ def _find_upstream_dependencies(
         parent_sid: SnapshotId,
         intervals_per_snapshot: t.Dict[str, Intervals],
         snapshots_to_create: t.Set[SnapshotId],
-    ) -> t.List[SchedulingUnit]:
+        cache: t.Dict[SnapshotId, t.Set[SchedulingUnit]],
+    ) -> t.Set[SchedulingUnit]:
         if parent_sid not in self.snapshots:
-            return []
+            return set()
+        if parent_sid in cache:
+            return cache[parent_sid]
 
         p_intervals = intervals_per_snapshot.get(parent_sid.name, [])
 
+        parent_node: t.Optional[SchedulingUnit] = None
         if p_intervals:
             if len(p_intervals) > 1:
-                return [DummyNode(snapshot_name=parent_sid.name)]
-            interval = p_intervals[0]
-            return [EvaluateNode(snapshot_name=parent_sid.name, interval=interval, batch_index=0)]
-        if parent_sid in snapshots_to_create:
-            return [CreateNode(snapshot_name=parent_sid.name)]
+                parent_node = DummyNode(snapshot_name=parent_sid.name)
+            else:
+                interval = p_intervals[0]
+                parent_node = EvaluateNode(
+                    snapshot_name=parent_sid.name, interval=interval, batch_index=0
+                )
+        elif parent_sid in snapshots_to_create:
+            parent_node = CreateNode(snapshot_name=parent_sid.name)
+
+        if parent_node is not None:
+            cache[parent_sid] = {parent_node}
+            return {parent_node}
+
         # This snapshot has no intervals and doesn't need creation which means
         # that it can be a transitive dependency
-        transitive_deps: t.List[SchedulingUnit] = []
+        transitive_deps: t.Set[SchedulingUnit] = set()
         parent_snapshot = self.snapshots[parent_sid]
         for grandparent_sid in parent_snapshot.parents:
-            transitive_deps.extend(
+            transitive_deps.update(
                 self._find_upstream_dependencies(
-                    grandparent_sid, intervals_per_snapshot, snapshots_to_create
+                    grandparent_sid, intervals_per_snapshot, snapshots_to_create, cache
                 )
             )
+        cache[parent_sid] = transitive_deps
         return transitive_deps
 
     def _run_or_audit(

sqlmesh/core/signal.py

@@ -2,6 +2,7 @@
 
 import typing as t
 from sqlmesh.utils import UniqueKeyDict, registry_decorator
+from sqlmesh.utils.errors import MissingSourceError
 
 if t.TYPE_CHECKING:
     from sqlmesh.core.context import ExecutionContext
@@ -42,7 +43,16 @@ class signal(registry_decorator):
 
 
 @signal()
-def freshness(batch: DatetimeRanges, snapshot: Snapshot, context: ExecutionContext) -> bool:
+def freshness(
+    batch: DatetimeRanges,
+    snapshot: Snapshot,
+    context: ExecutionContext,
+) -> bool:
+    """
+    Implements model freshness as a signal, i.e it considers this model to be fresh if:
+    - Any upstream SQLMesh model has available intervals to compute i.e is fresh
+    - Any upstream external model has been altered since the last time the model was evaluated
+    """
     adapter = context.engine_adapter
     if context.is_restatement or not adapter.SUPPORTS_METADATA_TABLE_LAST_MODIFIED_TS:
         return True
@@ -54,24 +64,35 @@ def freshness(batch: DatetimeRanges, snapshot: Snapshot, context: ExecutionConte
         if deployability_index.is_deployable(snapshot)
         else snapshot.dev_last_altered_ts
     )
+
     if not last_altered_ts:
         return True
 
     parent_snapshots = {context.snapshots[p.name] for p in snapshot.parents}
-    if len(parent_snapshots) != len(snapshot.node.depends_on) or not all(
-        p.is_external for p in parent_snapshots
-    ):
-        # The mismatch can happen if e.g an external model is not registered in the project
+
+    upstream_parent_snapshots = {p for p in parent_snapshots if not p.is_external}
+    external_parents = snapshot.node.depends_on - {p.name for p in upstream_parent_snapshots}
+
+    if context.parent_intervals:
+        # At least one upstream sqlmesh model has intervals to compute (i.e is fresh),
+        # so the current model is considered fresh too
         return True
 
-    # Finding new data means that the upstream depedencies have been altered
-    # since the last time the model was evaluated
-    upstream_dep_has_new_data = any(
-        upstream_last_altered_ts > last_altered_ts
-        for upstream_last_altered_ts in adapter.get_table_last_modified_ts(
-            [p.name for p in parent_snapshots]
+    if external_parents:
+        external_last_altered_timestamps = adapter.get_table_last_modified_ts(
+            list(external_parents)
+        )
+
+        if len(external_last_altered_timestamps) != len(external_parents):
+            raise MissingSourceError(
+                f"Expected {len(external_parents)} sources to be present, but got {len(external_last_altered_timestamps)}."
+            )
+
+        # Finding new data means that the upstream depedencies have been altered
+        # since the last time the model was evaluated
+        return any(
+            external_last_altered_ts > last_altered_ts
+            for external_last_altered_ts in external_last_altered_timestamps
         )
-    )
 
-    # Returning true is a no-op, returning False nullifies the batch so the model will not be evaluated.
-    return upstream_dep_has_new_data
+    return False