sqlmesh 0.217.1.dev1__py3-none-any.whl → 0.227.2.dev4__py3-none-any.whl

sqlmesh/core/node.py

@@ -153,6 +153,101 @@ class IntervalUnit(str, Enum):
         return self.seconds * 1000
 
 
+class DbtNodeInfo(PydanticModel):
+    """
+    Represents dbt-specific model information set by the dbt loader and intended to be made available at the Snapshot level
+    (as opposed to hidden within the individual model jinja macro registries).
+
+    This allows for things like injecting implementations of variables / functions into the Jinja context that are compatible with
+    their dbt equivalents but are backed by the sqlmesh snapshots in any given plan / environment
+    """
+
+    unique_id: str
+    """This is the node/resource name/unique_id that's used as the node key in the dbt manifest.
+    It's prefixed by the resource type and is exposed in context variables like {{ selected_resources }}.
+
+    Examples:
+        - test.jaffle_shop.unique_stg_orders_order_id.e3b841c71a
+        - seed.jaffle_shop.raw_payments
+        - model.jaffle_shop.stg_orders
+    """
+
+    name: str
+    """Name of this object in the dbt global namespace, used by things like {{ ref() }} calls.    
+    
+    Examples:
+        - unique_stg_orders_order_id
+        - raw_payments
+        - stg_orders
+    """
+
+    fqn: str
+    """Used for selectors in --select/--exclude.
+    Takes the filesystem into account so may be structured differently to :unique_id.
+    
+    Examples:
+        - jaffle_shop.staging.unique_stg_orders_order_id
+        - jaffle_shop.raw_payments
+        - jaffle_shop.staging.stg_orders
+    """
+
+    alias: t.Optional[str] = None
+    """This is dbt's way of overriding the _physical table_ a model is written to.
+
+    It's used in the following situation:
+     - Say you have two models, "stg_customers" and "customers"
+     - You want "stg_customers" to be written to the "staging" schema as eg "staging.customers" - NOT "staging.stg_customers"
+     - But you cant rename the file to "customers" because it will conflict with your other model file "customers"
+     - Even if you put it in a different folder, eg "staging/customers.sql" - dbt still has a global namespace so it will conflict
+        when you try to do something like "{{ ref('customers') }}"
+     - So dbt's solution to this problem is to keep calling it "stg_customers" at the dbt project/model level,
+        but allow overriding the physical table to "customers" via something like "{{ config(alias='customers', schema='staging') }}"
+
+    Note that if :alias is set, it does *not* replace :name at the model level and cannot be used interchangably with :name.
+    It also does not affect the :fqn or :unique_id. It's just used to override :name when it comes time to generate the physical table name.
+    """
+
+    @model_validator(mode="after")
+    def post_init(self) -> Self:
+        # by default, dbt sets alias to the same as :name
+        # however, we only want to include :alias if it is actually different / actually providing an override
+        if self.alias == self.name:
+            self.alias = None
+        return self
+
+    def to_expression(self) -> exp.Expression:
+        """Produce a SQLGlot expression representing this object, for use in things like the model/audit definition renderers"""
+        return exp.tuple_(
+            *(
+                exp.PropertyEQ(this=exp.var(k), expression=exp.Literal.string(v))
+                for k, v in sorted(self.model_dump(exclude_none=True).items())
+            )
+        )
+
+
+class DbtInfoMixin:
+    """This mixin encapsulates properties that only exist for dbt compatibility and are otherwise not required
+    for native projects"""
+
+    @property
+    def dbt_node_info(self) -> t.Optional[DbtNodeInfo]:
+        raise NotImplementedError()
+
+    @property
+    def dbt_unique_id(self) -> t.Optional[str]:
+        """Used for compatibility with jinja context variables such as {{ selected_resources }}"""
+        if self.dbt_node_info:
+            return self.dbt_node_info.unique_id
+        return None
+
+    @property
+    def dbt_fqn(self) -> t.Optional[str]:
+        """Used in the selector engine for compatibility with selectors that select models by dbt fqn"""
+        if self.dbt_node_info:
+            return self.dbt_node_info.fqn
+        return None
+
+
 # this must be sorted in descending order
 INTERVAL_SECONDS = {
     IntervalUnit.YEAR: 60 * 60 * 24 * 365,
@@ -165,7 +260,7 @@ INTERVAL_SECONDS = {
 }
 
 
-class _Node(PydanticModel):
+class _Node(DbtInfoMixin, PydanticModel):
     """
     Node is the core abstraction for entity that can be executed within the scheduler.
 
@@ -199,7 +294,7 @@ class _Node(PydanticModel):
     interval_unit_: t.Optional[IntervalUnit] = Field(alias="interval_unit", default=None)
     tags: t.List[str] = []
     stamp: t.Optional[str] = None
-    dbt_name: t.Optional[str] = None  # dbt node name
+    dbt_node_info_: t.Optional[DbtNodeInfo] = Field(alias="dbt_node_info", default=None)
     _path: t.Optional[Path] = None
     _data_hash: t.Optional[str] = None
     _metadata_hash: t.Optional[str] = None
@@ -446,6 +541,10 @@ class _Node(PydanticModel):
         """Return True if this is an audit node"""
         return False
 
+    @property
+    def dbt_node_info(self) -> t.Optional[DbtNodeInfo]:
+        return self.dbt_node_info_
+
 
 class NodeType(str, Enum):
     MODEL = "model"

sqlmesh/core/plan/builder.py

@@ -65,6 +65,9 @@ class PlanBuilder:
         restate_models: A list of models for which the data should be restated for the time range
             specified in this plan. Note: models defined outside SQLMesh (external) won't be a part
             of the restatement.
+        restate_all_snapshots: If restatements are present, this flag indicates whether or not the intervals
+            being restated should be cleared from state for other versions of this model (typically, versions that are present in other environments).
+            If set to None, the default behaviour is to not clear anything unless the target environment is prod.
         backfill_models: A list of fully qualified model names for which the data should be backfilled as part of this plan.
         no_gaps:  Whether to ensure that new snapshots for nodes that are already a
             part of the target environment have no data gaps when compared against previous
@@ -103,6 +106,7 @@ class PlanBuilder:
         execution_time: t.Optional[TimeLike] = None,
         apply: t.Optional[t.Callable[[Plan], None]] = None,
         restate_models: t.Optional[t.Iterable[str]] = None,
+        restate_all_snapshots: bool = False,
         backfill_models: t.Optional[t.Iterable[str]] = None,
         no_gaps: bool = False,
         skip_backfill: bool = False,
@@ -154,6 +158,7 @@ class PlanBuilder:
         self._auto_categorization_enabled = auto_categorization_enabled
         self._include_unmodified = include_unmodified
         self._restate_models = set(restate_models) if restate_models is not None else None
+        self._restate_all_snapshots = restate_all_snapshots
         self._effective_from = effective_from
 
         # note: this deliberately doesnt default to now() here.
@@ -277,7 +282,6 @@ class PlanBuilder:
         if self._latest_plan:
             return self._latest_plan
 
-        self._ensure_no_new_snapshots_with_restatements()
         self._ensure_new_env_with_changes()
         self._ensure_valid_date_range()
         self._ensure_no_broken_references()
@@ -338,7 +342,9 @@ class PlanBuilder:
             directly_modified=directly_modified,
             indirectly_modified=indirectly_modified,
             deployability_index=deployability_index,
+            selected_models_to_restate=self._restate_models,
             restatements=restatements,
+            restate_all_snapshots=self._restate_all_snapshots,
             start_override_per_model=self._start_override_per_model,
             end_override_per_model=end_override_per_model,
             selected_models_to_backfill=self._backfill_models,
@@ -674,6 +680,14 @@ class PlanBuilder:
                     if mode == AutoCategorizationMode.FULL:
                         snapshot.categorize_as(SnapshotChangeCategory.BREAKING, forward_only)
         elif self._context_diff.indirectly_modified(snapshot.name):
+            if snapshot.is_materialized_view and not forward_only:
+                # We categorize changes as breaking to allow for instantaneous switches in a virtual layer.
+                # Otherwise, there might be a potentially long downtime during MVs recreation.
+                # In the case of forward-only changes this optimization is not applicable because we want to continue
+                # using the same (existing) table version.
+                snapshot.categorize_as(SnapshotChangeCategory.INDIRECT_BREAKING, forward_only)
+                return
+
             all_upstream_forward_only = set()
             all_upstream_categories = set()
             direct_parent_categories = set()
@@ -858,15 +872,6 @@ class PlanBuilder:
                     f"""Removed {broken_references_msg} are referenced in '{snapshot.name}'. Please remove broken references before proceeding."""
                 )
 
-    def _ensure_no_new_snapshots_with_restatements(self) -> None:
-        if self._restate_models is not None and (
-            self._context_diff.new_snapshots or self._context_diff.modified_snapshots
-        ):
-            raise PlanError(
-                "Model changes and restatements can't be a part of the same plan. "
-                "Revert or apply changes before proceeding with restatements."
-            )
-
     def _ensure_new_env_with_changes(self) -> None:
         if (
             self._is_dev

sqlmesh/core/plan/common.py

@@ -1,19 +1,40 @@
 from __future__ import annotations
+import typing as t
+import logging
+from dataclasses import dataclass, field
 
-from sqlmesh.core.snapshot import Snapshot
+from sqlmesh.core.state_sync import StateReader
+from sqlmesh.core.snapshot import Snapshot, SnapshotId, SnapshotIdAndVersion, SnapshotNameVersion
+from sqlmesh.core.snapshot.definition import Interval
+from sqlmesh.utils.dag import DAG
+from sqlmesh.utils.date import now_timestamp
+
+logger = logging.getLogger(__name__)
 
 
 def should_force_rebuild(old: Snapshot, new: Snapshot) -> bool:
     if new.is_view and new.is_indirect_non_breaking and not new.is_forward_only:
         # View models always need to be rebuilt to reflect updated upstream dependencies
         return True
-    if new.is_seed:
+    if new.is_seed and not (
+        new.is_metadata
+        and new.previous_version
+        and new.previous_version.snapshot_id(new.name) == old.snapshot_id
+    ):
         # Seed models always need to be rebuilt to reflect changes in the seed file
+        # Unless only their metadata has been updated (eg description added) and the seed file has not been touched
         return True
     return is_breaking_kind_change(old, new)
 
 
 def is_breaking_kind_change(old: Snapshot, new: Snapshot) -> bool:
+    if new.is_model != old.is_model:
+        # If one is a model and the other isn't, then we need to rebuild
+        return True
+    if not new.is_model or not old.is_model:
+        # If neither are models, then we don't need to rebuild
+        # Note that the remaining checks only apply to model snapshots
+        return False
     if old.virtual_environment_mode != new.virtual_environment_mode:
         # If the virtual environment mode has changed, then we need to rebuild
         return True
@@ -27,3 +48,176 @@ def is_breaking_kind_change(old: Snapshot, new: Snapshot) -> bool:
         # If the partitioning hasn't changed, then we don't need to rebuild
         return False
     return True
+
+
+@dataclass
+class SnapshotIntervalClearRequest:
+    # affected snapshot
+    snapshot: SnapshotIdAndVersion
+
+    # which interval to clear
+    interval: Interval
+
+    # which environments this snapshot is currently promoted
+    # note that this can be empty if the snapshot exists because its ttl has not expired
+    # but it is not part of any particular environment
+    environment_names: t.Set[str] = field(default_factory=set)
+
+    @property
+    def snapshot_id(self) -> SnapshotId:
+        return self.snapshot.snapshot_id
+
+    @property
+    def sorted_environment_names(self) -> t.List[str]:
+        return list(sorted(self.environment_names))
+
+
+def identify_restatement_intervals_across_snapshot_versions(
+    state_reader: StateReader,
+    prod_restatements: t.Dict[str, Interval],
+    disable_restatement_models: t.Set[str],
+    loaded_snapshots: t.Dict[SnapshotId, Snapshot],
+    current_ts: t.Optional[int] = None,
+) -> t.Dict[SnapshotId, SnapshotIntervalClearRequest]:
+    """
+    Given a map of snapshot names + intervals to restate in prod:
+        - Look up matching snapshots (match based on name - regardless of version, to get all versions)
+        - For each match, also match downstream snapshots in each dev environment while filtering out models that have restatement disabled
+        - Return a list of all snapshots that are affected + the interval that needs to be cleared for each
+
+    The goal here is to produce a list of intervals to invalidate across all dev snapshots so that a subsequent plan or
+    cadence run in those environments causes the intervals to be repopulated.
+    """
+    if not prod_restatements:
+        return {}
+
+    # Although :loaded_snapshots is sourced from RestatementStage.all_snapshots, since the only time we ever need
+    # to clear intervals across all environments is for prod, the :loaded_snapshots here are always from prod
+    prod_name_versions: t.Set[SnapshotNameVersion] = {
+        s.name_version for s in loaded_snapshots.values()
+    }
+
+    snapshot_intervals_to_clear: t.Dict[SnapshotId, SnapshotIntervalClearRequest] = {}
+
+    for env_summary in state_reader.get_environments_summary():
+        # Fetch the full environment object one at a time to avoid loading all environments into memory at once
+        env = state_reader.get_environment(env_summary.name)
+        if not env:
+            logger.warning("Environment %s not found", env_summary.name)
+            continue
+
+        snapshots_by_name = {s.name: s.table_info for s in env.snapshots}
+
+        # We dont just restate matching snapshots, we also have to restate anything downstream of them
+        # so that if A gets restated in prod and dev has A <- B <- C, B and C get restated in dev
+        env_dag = DAG({s.name: {p.name for p in s.parents} for s in env.snapshots})
+
+        for restate_snapshot_name, interval in prod_restatements.items():
+            if restate_snapshot_name not in snapshots_by_name:
+                # snapshot is not promoted in this environment
+                continue
+
+            affected_snapshot_names = [
+                x
+                for x in ([restate_snapshot_name] + env_dag.downstream(restate_snapshot_name))
+                if x not in disable_restatement_models
+            ]
+
+            for affected_snapshot_name in affected_snapshot_names:
+                affected_snapshot = snapshots_by_name[affected_snapshot_name]
+
+                # Don't clear intervals for a dev snapshot if it shares the same physical version with prod.
+                # Otherwise, prod will be affected by what should be a dev operation
+                if affected_snapshot.name_version in prod_name_versions:
+                    continue
+
+                clear_request = snapshot_intervals_to_clear.get(affected_snapshot.snapshot_id)
+                if not clear_request:
+                    clear_request = SnapshotIntervalClearRequest(
+                        snapshot=affected_snapshot.id_and_version, interval=interval
+                    )
+                    snapshot_intervals_to_clear[affected_snapshot.snapshot_id] = clear_request
+
+                clear_request.environment_names |= set([env.name])
+
+    # snapshot_intervals_to_clear now contains the entire hierarchy of affected snapshots based
+    # on building the DAG for each environment and including downstream snapshots
+    # but, what if there are affected snapshots that arent part of any environment?
+    unique_snapshot_names = set(snapshot_id.name for snapshot_id in snapshot_intervals_to_clear)
+
+    current_ts = current_ts or now_timestamp()
+    all_matching_non_prod_snapshots = {
+        s.snapshot_id: s
+        for s in state_reader.get_snapshots_by_names(
+            snapshot_names=unique_snapshot_names, current_ts=current_ts, exclude_expired=True
+        )
+        # Don't clear intervals for a snapshot if it shares the same physical version with prod.
+        # Otherwise, prod will be affected by what should be a dev operation
+        if s.name_version not in prod_name_versions
+    }
+
+    # identify the ones that we havent picked up yet, which are the ones that dont exist in any environment
+    if remaining_snapshot_ids := set(all_matching_non_prod_snapshots).difference(
+        snapshot_intervals_to_clear
+    ):
+        # these snapshot id's exist in isolation and may be related to a downstream dependency of the :prod_restatements,
+        # rather than directly related, so we can't simply look up the interval to clear based on :prod_restatements.
+        # To figure out the interval that should be cleared, we can match to the existing list based on name
+        # and conservatively take the widest interval that shows up
+        snapshot_name_to_widest_interval: t.Dict[str, Interval] = {}
+        for s_id, clear_request in snapshot_intervals_to_clear.items():
+            current_start, current_end = snapshot_name_to_widest_interval.get(
+                s_id.name, clear_request.interval
+            )
+            next_start, next_end = clear_request.interval
+
+            next_start = min(current_start, next_start)
+            next_end = max(current_end, next_end)
+
+            snapshot_name_to_widest_interval[s_id.name] = (next_start, next_end)
+
+        for remaining_snapshot_id in remaining_snapshot_ids:
+            remaining_snapshot = all_matching_non_prod_snapshots[remaining_snapshot_id]
+            snapshot_intervals_to_clear[remaining_snapshot_id] = SnapshotIntervalClearRequest(
+                snapshot=remaining_snapshot,
+                interval=snapshot_name_to_widest_interval[remaining_snapshot_id.name],
+            )
+
+    # for any affected full_history_restatement_only snapshots, we need to widen the intervals being restated to
+    # include the whole time range for that snapshot. This requires a call to state to load the full snapshot record,
+    # so we only do it if necessary
+    full_history_restatement_snapshot_ids = [
+        # FIXME: full_history_restatement_only is just one indicator that the snapshot can only be fully refreshed, the other one is Model.depends_on_self
+        # however, to figure out depends_on_self, we have to render all the model queries which, alongside having to fetch full snapshots from state,
+        # is problematic in secure environments that are deliberately isolated from arbitrary user code (since rendering a query may require user macros to be present)
+        # So for now, these are not considered
+        s_id
+        for s_id, s in snapshot_intervals_to_clear.items()
+        if s.snapshot.full_history_restatement_only
+    ]
+    if full_history_restatement_snapshot_ids:
+        # only load full snapshot records that we havent already loaded
+        additional_snapshots = state_reader.get_snapshots(
+            [
+                s.snapshot_id
+                for s in full_history_restatement_snapshot_ids
+                if s.snapshot_id not in loaded_snapshots
+            ]
+        )
+
+        all_snapshots = loaded_snapshots | additional_snapshots
+
+        for full_snapshot_id in full_history_restatement_snapshot_ids:
+            full_snapshot = all_snapshots[full_snapshot_id]
+            intervals_to_clear = snapshot_intervals_to_clear[full_snapshot_id]
+
+            original_start, original_end = intervals_to_clear.interval
+
+            # get_removal_interval() widens intervals if necessary
+            new_interval = full_snapshot.get_removal_interval(
+                start=original_start, end=original_end
+            )
+
+            intervals_to_clear.interval = new_interval
+
+    return snapshot_intervals_to_clear

sqlmesh/core/plan/definition.py

@@ -58,7 +58,18 @@ class Plan(PydanticModel, frozen=True):
     indirectly_modified: t.Dict[SnapshotId, t.Set[SnapshotId]]
 
     deployability_index: DeployabilityIndex
+    selected_models_to_restate: t.Optional[t.Set[str]] = None
+    """Models that have been explicitly selected for restatement by a user"""
     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
+    """
+    restate_all_snapshots: bool
+    """Whether or not to clear intervals from state for other versions of the models listed in :restatements"""
+
     start_override_per_model: t.Optional[t.Dict[str, datetime]]
     end_override_per_model: t.Optional[t.Dict[str, datetime]]
 
@@ -202,8 +213,8 @@ class Plan(PydanticModel, frozen=True):
 
         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 = {
@@ -214,12 +225,14 @@ class Plan(PydanticModel, frozen=True):
                         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
@@ -259,6 +272,7 @@ class Plan(PydanticModel, frozen=True):
             skip_backfill=self.skip_backfill,
             empty_backfill=self.empty_backfill,
             restatements={s.name: i for s, i in self.restatements.items()},
+            restate_all_snapshots=self.restate_all_snapshots,
             is_dev=self.is_dev,
             allow_destructive_models=self.allow_destructive_models,
             allow_additive_models=self.allow_additive_models,
@@ -303,6 +317,7 @@ class EvaluatablePlan(PydanticModel):
     skip_backfill: bool
     empty_backfill: bool
     restatements: t.Dict[str, Interval]
+    restate_all_snapshots: bool
     is_dev: bool
     allow_destructive_models: t.Set[str]
     allow_additive_models: t.Set[str]