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

sqlmesh/core/state_sync/db/snapshot.py

@@ -6,7 +6,6 @@ import logging
 from pathlib import Path
 from collections import defaultdict
 from sqlglot import exp
-from pydantic import Field
 
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.state_sync.db.utils import (
@@ -15,9 +14,7 @@ from sqlmesh.core.state_sync.db.utils import (
     snapshot_id_filter,
     fetchone,
     fetchall,
-    create_batches,
 )
-from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.model import SeedModel, ModelKindName
 from sqlmesh.core.snapshot.cache import SnapshotCache
@@ -28,13 +25,18 @@ from sqlmesh.core.snapshot import (
     SnapshotNameVersion,
     SnapshotInfoLike,
     Snapshot,
+    SnapshotIdAndVersion,
     SnapshotId,
     SnapshotFingerprint,
-    SnapshotChangeCategory,
+)
+from sqlmesh.core.state_sync.common import (
+    RowBoundary,
+    ExpiredSnapshotBatch,
+    ExpiredBatchRange,
+    LimitBoundary,
 )
 from sqlmesh.utils.migration import index_text_type, blob_text_type
 from sqlmesh.utils.date import now_timestamp, TimeLike, to_timestamp
-from sqlmesh.utils.pydantic import PydanticModel
 from sqlmesh.utils import unique
 
 if t.TYPE_CHECKING:
@@ -63,6 +65,7 @@ class SnapshotState:
             "name": exp.DataType.build(index_type),
             "identifier": exp.DataType.build(index_type),
             "version": exp.DataType.build(index_type),
+            "dev_version": exp.DataType.build(index_type),
             "snapshot": exp.DataType.build(blob_type),
             "kind_name": exp.DataType.build("text"),
             "updated_ts": exp.DataType.build("bigint"),
@@ -70,6 +73,7 @@ class SnapshotState:
             "ttl_ms": exp.DataType.build("bigint"),
             "unrestorable": exp.DataType.build("boolean"),
             "forward_only": exp.DataType.build("boolean"),
+            "fingerprint": exp.DataType.build(blob_type),
         }
 
         self._auto_restatement_columns_to_types = {
@@ -164,94 +168,136 @@ class SnapshotState:
         self,
         environments: t.Iterable[Environment],
         current_ts: int,
-        ignore_ttl: bool = False,
-    ) -> t.List[SnapshotTableCleanupTask]:
-        """Aggregates the id's of the expired snapshots and creates a list of table cleanup tasks.
-
-        Expired snapshots are snapshots that have exceeded their time-to-live
-        and are no longer in use within an environment.
-
-        Returns:
-            The set of expired snapshot ids.
-            The list of table cleanup tasks.
-        """
-        _, cleanup_targets = self._get_expired_snapshots(
-            environments=environments,
-            current_ts=current_ts,
-            ignore_ttl=ignore_ttl,
+        ignore_ttl: bool,
+        batch_range: ExpiredBatchRange,
+    ) -> t.Optional[ExpiredSnapshotBatch]:
+        expired_query = exp.select("name", "identifier", "version", "updated_ts").from_(
+            self.snapshots_table
         )
-        return cleanup_targets
-
-    def _get_expired_snapshots(
-        self,
-        environments: t.Iterable[Environment],
-        current_ts: int,
-        ignore_ttl: bool = False,
-    ) -> t.Tuple[t.Set[SnapshotId], t.List[SnapshotTableCleanupTask]]:
-        expired_query = exp.select("name", "identifier", "version").from_(self.snapshots_table)
 
         if not ignore_ttl:
             expired_query = expired_query.where(
                 (exp.column("updated_ts") + exp.column("ttl_ms")) <= current_ts
             )
 
+        expired_query = expired_query.where(batch_range.where_filter)
+
+        promoted_snapshot_ids = {
+            snapshot.snapshot_id
+            for environment in environments
+            for snapshot in (
+                environment.snapshots
+                if environment.finalized_ts is not None
+                # If the environment is not finalized, check both the current snapshots and the previous finalized snapshots
+                else [*environment.snapshots, *(environment.previous_finalized_snapshots or [])]
+            )
+        }
+
+        if promoted_snapshot_ids:
+            not_in_conditions = [
+                exp.not_(condition)
+                for condition in snapshot_id_filter(
+                    self.engine_adapter,
+                    promoted_snapshot_ids,
+                    batch_size=self.SNAPSHOT_BATCH_SIZE,
+                )
+            ]
+            expired_query = expired_query.where(exp.and_(*not_in_conditions))
+
+        expired_query = expired_query.order_by(
+            exp.column("updated_ts"), exp.column("name"), exp.column("identifier")
+        )
+
+        if isinstance(batch_range.end, LimitBoundary):
+            expired_query = expired_query.limit(batch_range.end.batch_size)
+
+        rows = fetchall(self.engine_adapter, expired_query)
+
+        if not rows:
+            return None
+
         expired_candidates = {
             SnapshotId(name=name, identifier=identifier): SnapshotNameVersion(
                 name=name, version=version
             )
-            for name, identifier, version in fetchall(self.engine_adapter, expired_query)
+            for name, identifier, version, _ in rows
         }
         if not expired_candidates:
-            return set(), []
+            return None
 
-        promoted_snapshot_ids = {
-            snapshot.snapshot_id
-            for environment in environments
-            for snapshot in environment.snapshots
-        }
-
-        def _is_snapshot_used(snapshot: SharedVersionSnapshot) -> bool:
+        def _is_snapshot_used(snapshot: SnapshotIdAndVersion) -> bool:
             return (
                 snapshot.snapshot_id in promoted_snapshot_ids
                 or snapshot.snapshot_id not in expired_candidates
             )
 
-        unique_expired_versions = unique(expired_candidates.values())
-        version_batches = create_batches(
-            unique_expired_versions, batch_size=self.SNAPSHOT_BATCH_SIZE
+        # Extract cursor values from last row for pagination
+        last_row = rows[-1]
+        last_row_boundary = RowBoundary(
+            updated_ts=last_row[3],
+            name=last_row[0],
+            identifier=last_row[1],
         )
-        cleanup_targets = []
-        expired_snapshot_ids = set()
-        for versions_batch in version_batches:
-            snapshots = self._get_snapshots_with_same_version(versions_batch)
-
-            snapshots_by_version = defaultdict(set)
-            snapshots_by_dev_version = defaultdict(set)
-            for s in snapshots:
-                snapshots_by_version[(s.name, s.version)].add(s.snapshot_id)
-                snapshots_by_dev_version[(s.name, s.dev_version)].add(s.snapshot_id)
-
-            expired_snapshots = [s for s in snapshots if not _is_snapshot_used(s)]
-            expired_snapshot_ids.update([s.snapshot_id for s in expired_snapshots])
-
-            for snapshot in expired_snapshots:
-                shared_version_snapshots = snapshots_by_version[(snapshot.name, snapshot.version)]
-                shared_version_snapshots.discard(snapshot.snapshot_id)
-
-                shared_dev_version_snapshots = snapshots_by_dev_version[
-                    (snapshot.name, snapshot.dev_version)
-                ]
-                shared_dev_version_snapshots.discard(snapshot.snapshot_id)
-
-                if not shared_dev_version_snapshots:
-                    cleanup_targets.append(
-                        SnapshotTableCleanupTask(
-                            snapshot=snapshot.full_snapshot.table_info,
-                            dev_table_only=bool(shared_version_snapshots),
-                        )
+        # The returned batch_range represents the actual range of rows in this batch
+        result_batch_range = ExpiredBatchRange(
+            start=batch_range.start,
+            end=last_row_boundary,
+        )
+
+        unique_expired_versions = unique(expired_candidates.values())
+        expired_snapshot_ids: t.Set[SnapshotId] = set()
+        cleanup_tasks: t.List[SnapshotTableCleanupTask] = []
+
+        snapshots = self._get_snapshots_with_same_version(unique_expired_versions)
+
+        snapshots_by_version = defaultdict(set)
+        snapshots_by_dev_version = defaultdict(set)
+        for s in snapshots:
+            snapshots_by_version[(s.name, s.version)].add(s.snapshot_id)
+            snapshots_by_dev_version[(s.name, s.dev_version)].add(s.snapshot_id)
+
+        expired_snapshots = [s for s in snapshots if not _is_snapshot_used(s)]
+        all_expired_snapshot_ids = {s.snapshot_id for s in expired_snapshots}
+
+        cleanup_targets: t.List[t.Tuple[SnapshotId, bool]] = []
+        for snapshot in expired_snapshots:
+            shared_version_snapshots = snapshots_by_version[(snapshot.name, snapshot.version)]
+            shared_version_snapshots.discard(snapshot.snapshot_id)
+
+            shared_dev_version_snapshots = snapshots_by_dev_version[
+                (snapshot.name, snapshot.dev_version)
+            ]
+            shared_dev_version_snapshots.discard(snapshot.snapshot_id)
+
+            if not shared_dev_version_snapshots:
+                dev_table_only = bool(shared_version_snapshots)
+                cleanup_targets.append((snapshot.snapshot_id, dev_table_only))
+
+        snapshot_ids_to_cleanup = [snapshot_id for snapshot_id, _ in cleanup_targets]
+        full_snapshots = self._get_snapshots(snapshot_ids_to_cleanup)
+        for snapshot_id, dev_table_only in cleanup_targets:
+            if snapshot_id in full_snapshots:
+                cleanup_tasks.append(
+                    SnapshotTableCleanupTask(
+                        snapshot=full_snapshots[snapshot_id].table_info,
+                        dev_table_only=dev_table_only,
                     )
+                )
+                expired_snapshot_ids.add(snapshot_id)
+                all_expired_snapshot_ids.discard(snapshot_id)
+
+        # Add any remaining expired snapshots that don't require cleanup
+        if all_expired_snapshot_ids:
+            expired_snapshot_ids.update(all_expired_snapshot_ids)
+
+        if expired_snapshot_ids or cleanup_tasks:
+            return ExpiredSnapshotBatch(
+                expired_snapshot_ids=expired_snapshot_ids,
+                cleanup_tasks=cleanup_tasks,
+                batch_range=result_batch_range,
+            )
 
-        return expired_snapshot_ids, cleanup_targets
+        return None
 
     def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
         """Deletes snapshots.
@@ -288,6 +334,55 @@ class SnapshotState:
         """
         return self._get_snapshots(snapshot_ids)
 
+    def get_snapshots_by_names(
+        self,
+        snapshot_names: t.Iterable[str],
+        current_ts: t.Optional[int] = None,
+        exclude_expired: bool = True,
+    ) -> t.Set[SnapshotIdAndVersion]:
+        """Return the snapshot records for all versions of the specified snapshot names.
+
+        Args:
+            snapshot_names: Iterable of snapshot names to fetch all snapshot records for
+            current_ts: Sets the current time for identifying which snapshots have expired so they can be excluded (only relevant if :exclude_expired=True)
+            exclude_expired: Whether or not to return the snapshot id's of expired snapshots in the result
+
+        Returns:
+            A set containing all the matched snapshot records. To fetch full snapshots, pass it into StateSync.get_snapshots()
+        """
+        if not snapshot_names:
+            return set()
+
+        if exclude_expired:
+            current_ts = current_ts or now_timestamp()
+            unexpired_expr = (exp.column("updated_ts") + exp.column("ttl_ms")) > current_ts
+        else:
+            unexpired_expr = None
+
+        return {
+            SnapshotIdAndVersion(
+                name=name,
+                identifier=identifier,
+                version=version,
+                kind_name=kind_name or None,
+                dev_version=dev_version,
+                fingerprint=fingerprint,
+            )
+            for where in snapshot_name_filter(
+                snapshot_names=snapshot_names,
+                batch_size=self.SNAPSHOT_BATCH_SIZE,
+            )
+            for name, identifier, version, kind_name, dev_version, fingerprint in fetchall(
+                self.engine_adapter,
+                exp.select(
+                    "name", "identifier", "version", "kind_name", "dev_version", "fingerprint"
+                )
+                .from_(self.snapshots_table)
+                .where(where)
+                .and_(unexpired_expr),
+            )
+        }
+
     def snapshots_exist(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> t.Set[SnapshotId]:
         """Checks if snapshots exist.
 
@@ -571,7 +666,7 @@ class SnapshotState:
         self,
         snapshots: t.Collection[SnapshotNameVersionLike],
         lock_for_update: bool = False,
-    ) -> t.List[SharedVersionSnapshot]:
+    ) -> t.List[SnapshotIdAndVersion]:
         """Fetches all snapshots that share the same version as the snapshots.
 
         The output includes the snapshots with the specified identifiers.
@@ -593,14 +688,12 @@ class SnapshotState:
         ):
             query = (
                 exp.select(
-                    "snapshot",
                     "name",
                     "identifier",
                     "version",
-                    "updated_ts",
-                    "unpaused_ts",
-                    "unrestorable",
-                    "forward_only",
+                    "kind_name",
+                    "dev_version",
+                    "fingerprint",
                 )
                 .from_(exp.to_table(self.snapshots_table).as_("snapshots"))
                 .where(where)
@@ -611,17 +704,15 @@ class SnapshotState:
             snapshot_rows.extend(fetchall(self.engine_adapter, query))
 
         return [
-            SharedVersionSnapshot.from_snapshot_record(
+            SnapshotIdAndVersion(
                 name=name,
                 identifier=identifier,
                 version=version,
-                updated_ts=updated_ts,
-                unpaused_ts=unpaused_ts,
-                unrestorable=unrestorable,
-                forward_only=forward_only,
-                snapshot=snapshot,
+                kind_name=kind_name or None,
+                dev_version=dev_version,
+                fingerprint=SnapshotFingerprint.parse_raw(fingerprint),
             )
-            for snapshot, name, identifier, version, updated_ts, unpaused_ts, unrestorable, forward_only in snapshot_rows
+            for name, identifier, version, kind_name, dev_version, fingerprint in snapshot_rows
         ]
 
 
@@ -676,6 +767,8 @@ def _snapshots_to_df(snapshots: t.Iterable[Snapshot]) -> pd.DataFrame:
                 "ttl_ms": snapshot.ttl_ms,
                 "unrestorable": snapshot.unrestorable,
                 "forward_only": snapshot.forward_only,
+                "dev_version": snapshot.dev_version,
+                "fingerprint": snapshot.fingerprint.json(),
             }
             for snapshot in snapshots
         ]
@@ -695,88 +788,3 @@ def _auto_restatements_to_df(auto_restatements: t.Dict[SnapshotNameVersion, int]
             for name_version, ts in auto_restatements.items()
         ]
     )
-
-
-class SharedVersionSnapshot(PydanticModel):
-    """A stripped down version of a snapshot that is used for fetching snapshots that share the same version
-    with a significantly reduced parsing overhead.
-    """
-
-    name: str
-    version: str
-    dev_version_: t.Optional[str] = Field(alias="dev_version")
-    identifier: str
-    fingerprint: SnapshotFingerprint
-    interval_unit: IntervalUnit
-    change_category: SnapshotChangeCategory
-    updated_ts: int
-    unpaused_ts: t.Optional[int]
-    unrestorable: bool
-    disable_restatement: bool
-    effective_from: t.Optional[TimeLike]
-    raw_snapshot: t.Dict[str, t.Any]
-    forward_only: bool
-
-    @property
-    def snapshot_id(self) -> SnapshotId:
-        return SnapshotId(name=self.name, identifier=self.identifier)
-
-    @property
-    def is_forward_only(self) -> bool:
-        return self.forward_only or self.change_category == SnapshotChangeCategory.FORWARD_ONLY
-
-    @property
-    def normalized_effective_from_ts(self) -> t.Optional[int]:
-        return (
-            to_timestamp(self.interval_unit.cron_floor(self.effective_from))
-            if self.effective_from
-            else None
-        )
-
-    @property
-    def dev_version(self) -> str:
-        return self.dev_version_ or self.fingerprint.to_version()
-
-    @property
-    def full_snapshot(self) -> Snapshot:
-        return Snapshot(
-            **{
-                **self.raw_snapshot,
-                "updated_ts": self.updated_ts,
-                "unpaused_ts": self.unpaused_ts,
-                "unrestorable": self.unrestorable,
-                "forward_only": self.forward_only,
-            }
-        )
-
-    @classmethod
-    def from_snapshot_record(
-        cls,
-        *,
-        name: str,
-        identifier: str,
-        version: str,
-        updated_ts: int,
-        unpaused_ts: t.Optional[int],
-        unrestorable: bool,
-        forward_only: bool,
-        snapshot: str,
-    ) -> SharedVersionSnapshot:
-        raw_snapshot = json.loads(snapshot)
-        raw_node = raw_snapshot["node"]
-        return SharedVersionSnapshot(
-            name=name,
-            version=version,
-            dev_version=raw_snapshot.get("dev_version"),
-            identifier=identifier,
-            fingerprint=raw_snapshot["fingerprint"],
-            interval_unit=raw_node.get("interval_unit", IntervalUnit.from_cron(raw_node["cron"])),
-            change_category=raw_snapshot["change_category"],
-            updated_ts=updated_ts,
-            unpaused_ts=unpaused_ts,
-            unrestorable=unrestorable,
-            disable_restatement=raw_node.get("kind", {}).get("disable_restatement", False),
-            effective_from=raw_snapshot.get("effective_from"),
-            raw_snapshot=raw_snapshot,
-            forward_only=forward_only,
-        )

sqlmesh/core/table_diff.py

@@ -367,8 +367,8 @@ class TableDiff:
                 column_type = matched_columns[name]
                 qualified_column = exp.column(name, table)
 
-                if column_type.is_type(*exp.DataType.FLOAT_TYPES):
-                    return exp.func("ROUND", qualified_column, exp.Literal.number(self.decimals))
+                if column_type.is_type(*exp.DataType.REAL_TYPES):
+                    return self.adapter._normalize_decimal_value(qualified_column, self.decimals)
                 if column_type.is_type(*exp.DataType.NESTED_TYPES):
                     return self.adapter._normalize_nested_value(qualified_column)
 

sqlmesh/core/test/context.py

@@ -18,6 +18,8 @@ class TestExecutionContext(ExecutionContext):
         models: All upstream models to use for expansion and mapping of physical locations.
     """
 
+    __test__ = False  # prevent pytest trying to collect this as a test class
+
     def __init__(
         self,
         engine_adapter: EngineAdapter,

sqlmesh/core/test/definition.py

@@ -100,8 +100,11 @@ class ModelTest(unittest.TestCase):
         self._validate_and_normalize_test()
 
         if self.engine_adapter.default_catalog:
-            self._fixture_catalog: t.Optional[exp.Identifier] = exp.parse_identifier(
-                self.engine_adapter.default_catalog, dialect=self._test_adapter_dialect
+            self._fixture_catalog: t.Optional[exp.Identifier] = normalize_identifiers(
+                exp.parse_identifier(
+                    self.engine_adapter.default_catalog, dialect=self._test_adapter_dialect
+                ),
+                dialect=self._test_adapter_dialect,
             )
         else:
             self._fixture_catalog = None
@@ -451,6 +454,9 @@ class ModelTest(unittest.TestCase):
         query = outputs.get("query")
         partial = outputs.pop("partial", None)
 
+        if ctes is None and query is None:
+            _raise_error("Incomplete test, outputs must contain 'query' or 'ctes'", self.path)
+
         def _normalize_rows(
             values: t.List[Row] | t.Dict,
             name: str,
@@ -641,16 +647,16 @@ class ModelTest(unittest.TestCase):
             return self._execute(query)
 
         rows = values["rows"]
+        columns_str: t.Optional[t.List[str]] = None
         if columns:
+            columns_str = [str(c) for c in columns]
             referenced_columns = list(dict.fromkeys(col for row in rows for col in row))
             _raise_if_unexpected_columns(columns, referenced_columns)
 
             if partial:
-                columns = referenced_columns
+                columns_str = [c for c in columns_str if c in referenced_columns]
 
-        return pd.DataFrame.from_records(
-            rows, columns=[str(c) for c in columns] if columns else None
-        )
+        return pd.DataFrame.from_records(rows, columns=columns_str)
 
     def _add_missing_columns(
         self, query: exp.Query, all_columns: t.Optional[t.Collection[str]] = None
@@ -801,7 +807,7 @@ class PythonModelTest(ModelTest):
             actual_df.reset_index(drop=True, inplace=True)
             expected = self._create_df(values, columns=self.model.columns_to_types, partial=partial)
 
-            self.assert_equal(expected, actual_df, sort=False, partial=partial)
+            self.assert_equal(expected, actual_df, sort=True, partial=partial)
 
     def _execute_model(self) -> pd.DataFrame:
         """Executes the python model and returns a DataFrame."""
@@ -919,8 +925,7 @@ def generate_test(
                 cte_output = test._execute(cte_query)
                 ctes[cte.alias] = (
                     pandas_timestamp_to_pydatetime(
-                        cte_output.apply(lambda col: col.map(_normalize_df_value)),
-                        cte_query.named_selects,
+                        df=cte_output.apply(lambda col: col.map(_normalize_df_value)),
                     )
                     .replace({np.nan: None})
                     .to_dict(orient="records")

sqlmesh/dbt/adapter.py

@@ -115,30 +115,39 @@ class BaseAdapter(abc.ABC):
         """Returns the value quoted according to the quote policy."""
         return self.quote(value) if getattr(self.quote_policy, component_type, False) else value
 
-    def dispatch(self, name: str, package: t.Optional[str] = None) -> t.Callable:
+    def dispatch(
+        self,
+        macro_name: str,
+        macro_namespace: t.Optional[str] = None,
+    ) -> t.Callable:
         """Returns a dialect-specific version of a macro with the given name."""
         target_type = self.jinja_globals["target"]["type"]
-        macro_suffix = f"__{name}"
+        macro_suffix = f"__{macro_name}"
 
         def _relevance(package_name_pair: t.Tuple[t.Optional[str], str]) -> t.Tuple[int, int]:
             """Lower scores more relevant."""
-            macro_package, macro_name = package_name_pair
+            macro_package, name = package_name_pair
 
-            package_score = 0 if macro_package == package else 1
+            package_score = 0 if macro_package == macro_namespace else 1
             name_score = 1
 
-            if macro_name.startswith("default"):
+            if name.startswith("default"):
                 name_score = 2
-            elif macro_name.startswith(target_type):
+            elif name.startswith(target_type):
                 name_score = 0
 
             return name_score, package_score
 
         jinja_env = self.jinja_macros.build_environment(**self.jinja_globals).globals
-        packages_to_check: t.List[t.Optional[str]] = [
-            package,
-            *(k for k in jinja_env if k.startswith("dbt")),
-        ]
+
+        packages_to_check: t.List[t.Optional[str]] = [None]
+        if macro_namespace is not None:
+            if macro_namespace in jinja_env:
+                packages_to_check = [self.jinja_macros.root_package_name, macro_namespace]
+
+        # Add dbt packages as fallback
+        packages_to_check.extend(k for k in jinja_env if k.startswith("dbt"))
+
         candidates = {}
         for macro_package in packages_to_check:
             macros = jinja_env.get(macro_package, {}) if macro_package else jinja_env
@@ -156,7 +165,7 @@ class BaseAdapter(abc.ABC):
             sorted_candidates = sorted(candidates, key=_relevance)
             return candidates[sorted_candidates[0]]
 
-        raise ConfigError(f"Macro '{name}', package '{package}' was not found.")
+        raise ConfigError(f"Macro '{macro_name}', package '{macro_namespace}' was not found.")
 
     def type(self) -> str:
         return self.project_dialect or ""
@@ -168,7 +177,8 @@ class BaseAdapter(abc.ABC):
 
     @property
     def graph(self) -> t.Any:
-        return AttributeDict(
+        flat_graph = self.jinja_globals.get("flat_graph", None)
+        return flat_graph or AttributeDict(
             {
                 "exposures": {},
                 "groups": {},
@@ -276,10 +286,6 @@ class RuntimeAdapter(BaseAdapter):
             **table_mapping,
         }
 
-    @property
-    def graph(self) -> t.Any:
-        return self.jinja_globals.get("flat_graph", super().graph)
-
     def get_relation(
         self, database: t.Optional[str], schema: str, identifier: str
     ) -> t.Optional[BaseRelation]: