Merge branch 'main' into fix/limit-duckdb

Author: cmgoffena13

pnpm-lock.yaml

@@ -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

sqlmesh/core/plan/explainer.py

@@ -34,6 +34,7 @@
     import pandas as pd
 
     from sqlglot.dialects.dialect import DialectType
+    from sqlglot.generator import Generator
 
     Row = t.Dict[str, t.Any]
 
@@ -48,6 +49,54 @@
 }
 
 
+_FROZEN_TIME_GENERATORS: t.Dict[
+    t.Tuple[t.Type[Generator], str, t.Optional[str]], t.Type[Generator]
+] = {}
+_FROZEN_TIME_GENERATORS_LOCK = threading.Lock()
+
+
+def _frozen_time_generator_class(
+    generator_class: t.Type[Generator], execution_time: str, dialect: t.Optional[str]
+) -> t.Type[Generator]:
+    """Returns a subclass of `generator_class` whose CURRENT_* transforms render `execution_time`.
+
+    SQLGlot caches one dispatch table per generator class, so a subclass gets its own table and
+    the dialect's shared generator class is never modified. Subclasses are cached per
+    (generator, execution time, dialect) so tests that share an execution time share a table.
+    """
+    key = (generator_class, execution_time, dialect)
+    with _FROZEN_TIME_GENERATORS_LOCK:
+        klass = _FROZEN_TIME_GENERATORS.get(key)
+        if klass is None:
+            exec_time = exp.Literal.string(execution_time)
+            klass = t.cast(
+                t.Type["Generator"],
+                type(
+                    f"{generator_class.__name__}FrozenTime",
+                    (generator_class,),
+                    {
+                        "TRANSFORMS": {
+                            **generator_class.TRANSFORMS,
+                            exp.CurrentDate: lambda self, _: self.sql(
+                                exp.cast(exec_time, "date", dialect=dialect)
+                            ),
+                            exp.CurrentDatetime: lambda self, _: self.sql(
+                                exp.cast(exec_time, "datetime", dialect=dialect)
+                            ),
+                            exp.CurrentTime: lambda self, _: self.sql(
+                                exp.cast(exec_time, "time", dialect=dialect)
+                            ),
+                            exp.CurrentTimestamp: lambda self, _: self.sql(
+                                exp.cast(exec_time, "timestamp", dialect=dialect)
+                            ),
+                        }
+                    },
+                ),
+            )
+            _FROZEN_TIME_GENERATORS[key] = klass
+        return klass
+
+
 class ModelTest(unittest.TestCase):
     __test__ = False
 
@@ -116,31 +165,22 @@ def __init__(
         )
         self._qualified_fixture_schema = schema_(self._fixture_schema, self._fixture_catalog)
 
-        self._transforms = self._test_adapter_dialect.generator_class.TRANSFORMS
         self._execution_time = str(self.body.get("vars", {}).get("execution_time") or "")
 
         if self._execution_time:
             # Normalizes the execution time by converting it into UTC timezone
             self._execution_time = str(to_datetime(self._execution_time))
 
-        # When execution_time is set, we mock the CURRENT_* SQL expressions so they always return it
+        # When execution_time is set, the CURRENT_* SQL expressions must render as that time. The
+        # overrides live on a per-test generator subclass rather than on the dialect's shared
+        # generator class: SQLGlot caches one dispatch table per generator class, so patching the
+        # shared one is visible to every other thread that renders SQL (e.g. a concurrent test
+        # creating its fixture views) and races with the patch's restoration.
+        self._generator_class = self._test_adapter_dialect.generator_class
         if self._execution_time:
-            exec_time = exp.Literal.string(self._execution_time)
-            self._transforms = {
-                **self._transforms,
-                exp.CurrentDate: lambda self, _: self.sql(
-                    exp.cast(exec_time, "date", dialect=dialect)
-                ),
-                exp.CurrentDatetime: lambda self, _: self.sql(
-                    exp.cast(exec_time, "datetime", dialect=dialect)
-                ),
-                exp.CurrentTime: lambda self, _: self.sql(
-                    exp.cast(exec_time, "time", dialect=dialect)
-                ),
-                exp.CurrentTimestamp: lambda self, _: self.sql(
-                    exp.cast(exec_time, "timestamp", dialect=dialect)
-                ),
-            }
+            self._generator_class = _frozen_time_generator_class(
+                self._generator_class, self._execution_time, dialect
+            )
 
         super().__init__()
 
@@ -600,15 +640,18 @@ def _normalize_column_name(self, name: str) -> str:
         return normalized_name
 
     @contextmanager
-    def _concurrent_render_context(self) -> t.Iterator[None]:
+    def _concurrent_render_context(self, patch_shared_dialect: bool = False) -> t.Iterator[None]:
         """
         Context manager that ensures that the tests are executed safely in a concurrent environment.
-        This is needed in case `execution_time` is set, as we'd then have to:
-        - Freeze time through `time_machine` (not thread safe)
-        - Globally patch the SQLGlot dialect so that any date/time nodes are evaluated at the `execution_time` during generation
+        This is needed in case `execution_time` is set, as we'd then have to freeze time through
+        `time_machine`, which is not thread safe.
+
+        SQL model tests render through `self._generator_class`, so the shared dialect is never
+        modified. Python model tests may run arbitrary SQL through the engine adapter, whose
+        generator cannot be swapped per test, so they additionally patch the shared generator's
+        transforms (`patch_shared_dialect=True`) while holding the lock.
         """
         import time_machine
-        from sqlglot.generator import _DISPATCH_CACHE
 
         lock_ctx: AbstractContextManager = (
             self.CONCURRENT_RENDER_LOCK if self.concurrency else nullcontext()
@@ -618,19 +661,30 @@ def _concurrent_render_context(self) -> t.Iterator[None]:
         dispatch_patch_ctx: AbstractContextManager = nullcontext()
 
         if self._execution_time:
-            generator_class = self._test_adapter_dialect.generator_class
             time_ctx = time_machine.travel(self._execution_time, tick=False)
-            dialect_patch_ctx = patch.dict(generator_class.TRANSFORMS, self._transforms)
+
+        if self._execution_time and patch_shared_dialect:
+            from sqlglot.generator import _DISPATCH_CACHE
+
+            generator_class = self._test_adapter_dialect.generator_class
+            transforms = self._generator_class.TRANSFORMS
+            dialect_patch_ctx = patch.dict(generator_class.TRANSFORMS, transforms)
 
             # sqlglot caches a dispatch table per generator class, so we need to patch
             # it as well to ensure the overridden transforms are actually used
             dispatch = _DISPATCH_CACHE.get(generator_class)
             if dispatch is not None:
-                dispatch_patch_ctx = patch.dict(dispatch, self._transforms)
+                dispatch_patch_ctx = patch.dict(dispatch, transforms)
 
         with lock_ctx, time_ctx, dialect_patch_ctx, dispatch_patch_ctx:
             yield
 
+    def _generate_sql(self, expression: exp.Expr) -> str:
+        """Generates SQL for the testing engine, rendering CURRENT_* at `execution_time` when set."""
+        return self._generator_class(
+            dialect=self._test_adapter_dialect, pretty=self.engine_adapter._pretty_sql
+        ).generate(expression)
+
     def _execute(self, query: exp.Query | str) -> pd.DataFrame:
         """Executes the given query using the testing engine adapter and returns a DataFrame."""
         return self.engine_adapter.fetchdf(query)
@@ -698,9 +752,7 @@ def test_ctes(self, ctes: t.Dict[str, exp.Expr], recursive: bool = False) -> Non
                 with self._concurrent_render_context():
                     # Similar to the model's query, we render the CTE query under the locked context
                     # so that the execution (fetchdf) can continue concurrently between the threads
-                    sql = cte_query.sql(
-                        self._test_adapter_dialect, pretty=self.engine_adapter._pretty_sql
-                    )
+                    sql = self._generate_sql(cte_query)
 
                 actual = self._execute(sql)
                 expected = self._create_df(values, columns=cte_query.named_selects, partial=partial)
@@ -712,7 +764,7 @@ def runTest(self) -> None:
             # Render the model's query and generate the SQL under the locked context so that
             # execution (fetchdf) can continue concurrently between the threads
             query = self._render_model_query()
-            sql = query.sql(self._test_adapter_dialect, pretty=self.engine_adapter._pretty_sql)
+            sql = self._generate_sql(query)
 
         with_clause = query.args.get("with_")
 
@@ -817,7 +869,7 @@ def _execute_model(self) -> pd.DataFrame:
         """Executes the python model and returns a DataFrame."""
         import pandas as pd
 
-        with self._concurrent_render_context():
+        with self._concurrent_render_context(patch_shared_dialect=True):
             variables = self.body.get("vars", {}).copy()
             time_kwargs = {key: variables.pop(key) for key in TIME_KWARG_KEYS if key in variables}
             df = next(self.model.render(context=self.context, variables=variables, **time_kwargs))

sqlmesh/core/test/definition.py

@@ -106,7 +106,7 @@ def _migrate_intervals(
             continue
 
         dev_version = snapshot_ids_to_dev_versions.get((name, identifier))
-        if dev_version not in used_dev_versions and is_dev:
+        if (name, dev_version) not in used_dev_versions and is_dev:
             # If the interval's dev version is no longer used and this is a dev interval, we can safely delete it
             continue