Merge branch 'main' into fix/audit-condition-z-score

Author: mday-io

pnpm-lock.yaml

@@ -9,7 +9,7 @@ requires-python = ">= 3.9"
 dependencies = [
     "click",
     "croniter",
-    "duckdb>=0.10.0,!=0.10.3",
+    "duckdb>=0.10.0,!=0.10.3,<2",
     "dateparser<=1.2.1",
     "humanize",
     "hyperscript>=0.1.0",
@@ -121,7 +121,7 @@ fabric = ["pyodbc>=5.0.0"]
 fabric-mssql-python = ["mssql-python>=1.1.0;python_version>=\"3.10\""]
 gcppostgres = ["cloud-sql-python-connector[pg8000]>=1.8.0"]
 github = ["PyGithub>=2.6.0"]
-motherduck = ["duckdb>=1.3.2"]
+motherduck = ["duckdb>=1.3.2,<2"]
 mssql = ["pymssql"]
 mssql-odbc = ["pyodbc>=5.0.0"]
 mssql-python = ["mssql-python>=1.1.0;python_version>=\"3.10\""]

pyproject.toml

@@ -748,7 +748,25 @@ def load(
             prod = self.state_reader.get_environment(c.PROD)
 
             if prod:
-                for snapshot in self.state_reader.get_snapshots(prod.snapshots).values():
+                # Environment snapshot infos already contain the names we need to distinguish
+                # local nodes from nodes owned by another project. Hydrating local snapshots
+                # here is wasteful: their payloads are not used, and a large remote state can
+                # spend most of Context.load() decoding those model object graphs.
+                #
+                # Only hydrate names which are absent locally. These can be deleted nodes from
+                # this project or nodes from another project which must be merged into the
+                # context. The project field on the hydrated node disambiguates the two cases.
+                remote_snapshot_infos = []
+                for snapshot_info in prod.snapshots:
+                    local_store = (
+                        self._standalone_audits if snapshot_info.is_audit else self._models
+                    )
+                    if snapshot_info.name in local_store:
+                        uncached.add(snapshot_info.name)
+                    else:
+                        remote_snapshot_infos.append(snapshot_info)
+
+                for snapshot in self.state_reader.get_snapshots(remote_snapshot_infos).values():
                     if snapshot.node.project in self._projects:
                         uncached.add(snapshot.name)
                     else:

sqlmesh/core/context.py

@@ -112,22 +112,48 @@ def with_optimized_query(self, model: Model, name: t.Optional[str] = None) -> bo
         cache_entry = self._file_cache.get(name)
         if cache_entry:
             try:
-                # If the optimized rendered query is None, then there are likely adapter calls in the query
-                # that prevent us from rendering it at load time. This means that we can safely set the
-                # unoptimized cache to None as well to prevent attempts to render it downstream.
-                optimized = cache_entry.optimized_rendered_query is not None
-                model._query_renderer.update_cache(
-                    cache_entry.optimized_rendered_query,
-                    cache_entry.renderer_violations,
-                    optimized=optimized,
-                )
+                self.with_optimized_query_entry(model, cache_entry)
                 return True
             except Exception as ex:
                 logger.warning("Failed to load a cache entry '%s': %s", name, ex)
 
         self._put(name, model)
         return False
 
+    @staticmethod
+    def with_optimized_query_entry(model: Model, cache_entry: OptimizedQueryCacheEntry) -> None:
+        """Adds an already-decoded optimized query cache entry to a model."""
+        if not isinstance(model, SqlModel):
+            return
+
+        # If the optimized rendered query is None, then there are likely adapter calls in the
+        # query that prevent us from rendering it at load time. This means that we can safely
+        # set the unoptimized cache to None as well to prevent attempts to render it downstream.
+        optimized = cache_entry.optimized_rendered_query is not None
+        model._query_renderer.update_cache(
+            cache_entry.optimized_rendered_query,
+            cache_entry.renderer_violations,
+            optimized=optimized,
+        )
+
+    def get_or_create_entry(
+        self, model: Model, name: t.Optional[str] = None
+    ) -> t.Optional[OptimizedQueryCacheEntry]:
+        """Returns one decoded cache entry, creating it when necessary."""
+        if not isinstance(model, SqlModel):
+            return None
+
+        name = self._entry_name(model) if name is None else name
+        cache_entry = self._file_cache.get(name)
+        if cache_entry:
+            try:
+                self.with_optimized_query_entry(model, cache_entry)
+                return cache_entry
+            except Exception as ex:
+                logger.warning("Failed to load a cache entry '%s': %s", name, ex)
+
+        return self._put(name, model)
+
     def put(self, model: Model) -> t.Optional[str]:
         if not isinstance(model, SqlModel):
             return None
@@ -140,14 +166,15 @@ def put(self, model: Model) -> t.Optional[str]:
         self._put(name, model)
         return name
 
-    def _put(self, name: str, model: SqlModel) -> None:
+    def _put(self, name: str, model: SqlModel) -> OptimizedQueryCacheEntry:
         optimized_query = model.render_query()
 
         new_entry = OptimizedQueryCacheEntry(
             optimized_rendered_query=optimized_query,
             renderer_violations=model.violated_rules_for_query,
         )
         self._file_cache.put(name, value=new_entry)
+        return new_entry
 
     @staticmethod
     def _entry_name(model: SqlModel) -> str:
@@ -197,7 +224,13 @@ def load_optimized_query(
 
 def load_optimized_query_and_mapping(
     model: Model, mapping: t.Dict
-) -> t.Tuple[str, t.Optional[str], str, str, t.Dict]:
+) -> t.Tuple[
+    str,
+    t.Optional[OptimizedQueryCacheEntry],
+    str,
+    str,
+    t.Dict,
+]:
     assert _optimized_query_cache
 
     schema = MappingSchema(normalize=False)
@@ -207,13 +240,13 @@ def load_optimized_query_and_mapping(
 
     if isinstance(model, SqlModel):
         entry_name = _optimized_query_cache._entry_name(model)
-        _optimized_query_cache.with_optimized_query(model, entry_name)
+        cache_entry = _optimized_query_cache.get_or_create_entry(model, entry_name)
     else:
-        entry_name = None
+        cache_entry = None
 
     return (
         model.fqn,
-        entry_name,
+        cache_entry,
         model.data_hash,
         model.metadata_hash,
         model.mapping_schema,

sqlmesh/core/model/cache.py

@@ -87,13 +87,16 @@ def process_models(completed_model: t.Optional[Model] = None) -> None:
             for future in as_completed(futures):
                 try:
                     futures.remove(future)
-                    fqn, entry_name, data_hash, metadata_hash, mapping_schema = future.result()
+                    fqn, cache_entry, data_hash, metadata_hash, mapping_schema = future.result()
                     model = models[fqn]
                     model._data_hash = data_hash
                     model._metadata_hash = metadata_hash
                     if model.mapping_schema != mapping_schema:
                         model.set_mapping_schema(mapping_schema)
-                    optimized_query_cache.with_optimized_query(model, entry_name)
+                    if cache_entry is not None:
+                        # The worker has already decompressed and decoded this entry. Passing
+                        # it back avoids a second gzip/pickle read in the parent process.
+                        optimized_query_cache.with_optimized_query_entry(model, cache_entry)
                     _update_schema_with_model(schema, model)
                     process_models(completed_model=model)
                 except Exception as ex:

sqlmesh/core/model/schema.py

@@ -8,9 +8,11 @@
 
 from rich.console import Console as RichConsole
 from rich.tree import Tree
+from sqlglot import exp
 from sqlglot.dialects.dialect import DialectType
 from sqlmesh.core import constants as c
 from sqlmesh.core.console import Console, TerminalConsole, get_console
+from sqlmesh.core.dialect import schema_
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.plan.common import (
     SnapshotIntervalClearRequest,
@@ -145,6 +147,29 @@ def visit_before_all_stage(self, stage: stages.BeforeAllStage) -> Tree:
     def visit_after_all_stage(self, stage: stages.AfterAllStage) -> Tree:
         return Tree("[bold]Execute after all statements[/bold]")
 
+    def visit_physical_layer_schema_creation_stage(
+        self, stage: stages.PhysicalLayerSchemaCreationStage
+    ) -> t.Optional[Tree]:
+        tables = [
+            exp.to_table(
+                snapshot.table_name(is_deployable=stage.deployability_index.is_deployable(snapshot))
+            )
+            for snapshot in stage.snapshots
+            if snapshot.is_model and not snapshot.is_symbolic
+        ]
+        schemas = {
+            schema_(table.args["db"], table.args.get("catalog")).sql(dialect=self.dialect)
+            for table in tables
+            if table.db
+        }
+        if not schemas:
+            return None
+
+        tree = Tree("[bold]Create physical schemas if they do not exist[/bold]")
+        for schema in sorted(schemas):
+            tree.add(schema)
+        return tree
+
     def visit_physical_layer_update_stage(self, stage: stages.PhysicalLayerUpdateStage) -> Tree:
         snapshots = [
             s for s in stage.snapshots if s.snapshot_id in stage.snapshots_with_missing_intervals