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

sqlmesh/core/context.py

@@ -93,7 +93,7 @@ from sqlmesh.core.plan.definition import UserProvidedFlags
 from sqlmesh.core.reference import ReferenceGraph
 from sqlmesh.core.scheduler import Scheduler, CompletionStatus
 from sqlmesh.core.schema_loader import create_external_models_file
-from sqlmesh.core.selector import Selector
+from sqlmesh.core.selector import Selector, NativeSelector
 from sqlmesh.core.snapshot import (
     DeployabilityIndex,
     Snapshot,
@@ -107,14 +107,15 @@ from sqlmesh.core.state_sync import (
     CachingStateSync,
     StateReader,
     StateSync,
-    cleanup_expired_views,
 )
+from sqlmesh.core.janitor import cleanup_expired_views, delete_expired_snapshots
 from sqlmesh.core.table_diff import TableDiff
 from sqlmesh.core.test import (
     ModelTextTestResult,
     ModelTestMetadata,
     generate_test,
     run_tests,
+    filter_tests_by_patterns,
 )
 from sqlmesh.core.user import User
 from sqlmesh.utils import UniqueKeyDict, Verbosity
@@ -139,6 +140,7 @@ from sqlmesh.utils.errors import (
 )
 from sqlmesh.utils.config import print_config
 from sqlmesh.utils.jinja import JinjaMacroRegistry
+from sqlmesh.utils.windows import IS_WINDOWS, fix_windows_path
 
 if t.TYPE_CHECKING:
     import pandas as pd
@@ -153,6 +155,8 @@ if t.TYPE_CHECKING:
     )
     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]
 
@@ -274,6 +278,8 @@ class ExecutionContext(BaseContext):
         deployability_index: t.Optional[DeployabilityIndex] = None,
         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,
     ):
@@ -284,6 +290,8 @@ class ExecutionContext(BaseContext):
         self._default_dialect = default_dialect
         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]:
@@ -308,6 +316,14 @@ class ExecutionContext(BaseContext):
         """Returns the gateway name."""
         return self.var(c.GATEWAY)
 
+    @property
+    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)
@@ -328,6 +344,7 @@ class ExecutionContext(BaseContext):
             self.deployability_index,
             self._default_dialect,
             self._default_catalog,
+            self._is_restatement,
             variables=variables,
             blueprint_variables=blueprint_variables,
         )
@@ -368,6 +385,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         load: bool = True,
         users: t.Optional[t.List[User]] = None,
         config_loader_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
+        selector: t.Optional[t.Type[Selector]] = None,
     ):
         self.configs = (
             config
@@ -381,6 +399,11 @@ class GenericContext(BaseContext, t.Generic[C]):
         self._standalone_audits: UniqueKeyDict[str, StandaloneAudit] = UniqueKeyDict(
             "standaloneaudits"
         )
+        self._model_test_metadata: t.List[ModelTestMetadata] = []
+        self._model_test_metadata_path_index: t.Dict[Path, t.List[ModelTestMetadata]] = {}
+        self._model_test_metadata_fully_qualified_name_index: t.Dict[str, ModelTestMetadata] = {}
+        self._models_with_tests: t.Set[str] = set()
+
         self._macros: UniqueKeyDict[str, ExecutableOrMacro] = UniqueKeyDict("macros")
         self._metrics: UniqueKeyDict[str, Metric] = UniqueKeyDict("metrics")
         self._jinja_macros = JinjaMacroRegistry()
@@ -390,6 +413,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         self._engine_adapter: t.Optional[EngineAdapter] = None
         self._linters: t.Dict[str, Linter] = {}
         self._loaded: bool = False
+        self._selector_cls = selector or NativeSelector
 
         self.path, self.config = t.cast(t.Tuple[Path, C], next(iter(self.configs.items())))
 
@@ -618,6 +642,10 @@ class GenericContext(BaseContext, t.Generic[C]):
         self._excluded_requirements.clear()
         self._linters.clear()
         self._environment_statements = []
+        self._model_test_metadata.clear()
+        self._model_test_metadata_path_index.clear()
+        self._model_test_metadata_fully_qualified_name_index.clear()
+        self._models_with_tests.clear()
 
         for loader, project in zip(self._loaders, loaded_projects):
             self._jinja_macros = self._jinja_macros.merge(project.jinja_macros)
@@ -630,6 +658,16 @@ class GenericContext(BaseContext, t.Generic[C]):
             self._excluded_requirements.update(project.excluded_requirements)
             self._environment_statements.extend(project.environment_statements)
 
+            self._model_test_metadata.extend(project.model_test_metadata)
+            for metadata in project.model_test_metadata:
+                if metadata.path not in self._model_test_metadata_path_index:
+                    self._model_test_metadata_path_index[metadata.path] = []
+                self._model_test_metadata_path_index[metadata.path].append(metadata)
+                self._model_test_metadata_fully_qualified_name_index[
+                    metadata.fully_qualified_test_name
+                ] = metadata
+                self._models_with_tests.add(metadata.model_name)
+
             config = loader.config
             self._linters[config.project] = Linter.from_rules(
                 BUILTIN_RULES.union(project.user_rules), config.linter
@@ -1031,6 +1069,11 @@ class GenericContext(BaseContext, t.Generic[C]):
         """Returns all registered standalone audits in this context."""
         return MappingProxyType(self._standalone_audits)
 
+    @property
+    def models_with_tests(self) -> t.Set[str]:
+        """Returns all models with tests in this context."""
+        return self._models_with_tests
+
     @property
     def snapshots(self) -> t.Dict[str, Snapshot]:
         """Generates and returns snapshots based on models registered in this context.
@@ -1429,6 +1472,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         explain: t.Optional[bool] = None,
         ignore_cron: t.Optional[bool] = None,
         min_intervals: t.Optional[int] = None,
+        always_include_local_changes: t.Optional[bool] = None,
     ) -> PlanBuilder:
         """Creates a plan builder.
 
@@ -1467,6 +1511,8 @@ class GenericContext(BaseContext, t.Generic[C]):
             diff_rendered: Whether the diff should compare raw vs rendered models
             min_intervals: Adjust the plan start date on a per-model basis in order to ensure at least this many intervals are covered
                 on every model when checking for missing intervals
+            always_include_local_changes: Usually when restatements are present, local changes in the filesystem are ignored.
+                However, it can be desirable to deploy changes + restatements in the same plan, so this flag overrides the default behaviour.
 
         Returns:
             The plan builder.
@@ -1583,13 +1629,20 @@ class GenericContext(BaseContext, t.Generic[C]):
                 "Selector did not return any models. Please check your model selection and try again."
             )
 
+        if always_include_local_changes is None:
+            # default behaviour - if restatements are detected; we operate entirely out of state and ignore local changes
+            force_no_diff = restate_models is not None or (
+                backfill_models is not None and not backfill_models
+            )
+        else:
+            force_no_diff = not always_include_local_changes
+
         snapshots = self._snapshots(models_override)
         context_diff = self._context_diff(
             environment or c.PROD,
             snapshots=snapshots,
             create_from=create_from,
-            force_no_diff=restate_models is not None
-            or (backfill_models is not None and not backfill_models),
+            force_no_diff=force_no_diff,
             ensure_finalized_snapshots=self.config.plan.use_finalized_state,
             diff_rendered=diff_rendered,
             always_recreate_environment=self.config.plan.always_recreate_environment,
@@ -1644,6 +1697,14 @@ class GenericContext(BaseContext, t.Generic[C]):
         elif forward_only is None:
             forward_only = self.config.plan.forward_only
 
+        # When handling prod restatements, only clear intervals from other model versions if we are using full virtual environments
+        # If we are not, then there is no point, because none of the data in dev environments can be promoted by definition
+        restate_all_snapshots = (
+            expanded_restate_models is not None
+            and not is_dev
+            and self.config.virtual_environment_mode.is_full
+        )
+
         return self.PLAN_BUILDER_TYPE(
             context_diff=context_diff,
             start=start,
@@ -1651,6 +1712,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             execution_time=execution_time,
             apply=self.apply,
             restate_models=expanded_restate_models,
+            restate_all_snapshots=restate_all_snapshots,
             backfill_models=backfill_models,
             no_gaps=no_gaps,
             skip_backfill=skip_backfill,
@@ -1678,9 +1740,9 @@ class GenericContext(BaseContext, t.Generic[C]):
             console=self.console,
             user_provided_flags=user_provided_flags,
             selected_models={
-                dbt_name
+                dbt_unique_id
                 for model in model_selector.expand_model_selections(select_models or "*")
-                if (dbt_name := snapshots[model].node.dbt_name)
+                if (dbt_unique_id := snapshots[model].node.dbt_unique_id)
             },
             explain=explain or False,
             ignore_cron=ignore_cron or False,
@@ -2183,7 +2245,7 @@ class GenericContext(BaseContext, t.Generic[C]):
 
             pd.set_option("display.max_columns", None)
 
-        test_meta = self.load_model_tests(tests=tests, patterns=match_patterns)
+        test_meta = self.select_tests(tests=tests, patterns=match_patterns)
 
         result = run_tests(
             model_test_metadata=test_meta,
@@ -2242,6 +2304,7 @@ class GenericContext(BaseContext, t.Generic[C]):
                 snapshot=snapshot,
                 start=start,
                 end=end,
+                execution_time=execution_time,
                 snapshots=self.snapshots,
             ):
                 audit_id = f"{audit_result.audit.name}"
@@ -2562,12 +2625,15 @@ class GenericContext(BaseContext, t.Generic[C]):
         )
 
     def clear_caches(self) -> None:
-        for path in self.configs:
-            cache_path = path / c.CACHE
-            if cache_path.exists():
-                rmtree(cache_path)
-        if self.cache_dir.exists():
-            rmtree(self.cache_dir)
+        paths_to_remove = [path / c.CACHE for path in self.configs]
+        paths_to_remove.append(self.cache_dir)
+
+        if IS_WINDOWS:
+            paths_to_remove = [fix_windows_path(path) for path in paths_to_remove]
+
+        for path in paths_to_remove:
+            if path.exists():
+                rmtree(path)
 
         if isinstance(self._state_sync, CachingStateSync):
             self._state_sync.clear_cache()
@@ -2824,19 +2890,14 @@ class GenericContext(BaseContext, t.Generic[C]):
         # Clean up expired environments by removing their views and schemas
         self._cleanup_environments(current_ts=current_ts)
 
-        cleanup_targets = self.state_sync.get_expired_snapshots(
-            ignore_ttl=ignore_ttl, current_ts=current_ts
-        )
-
-        # Remove the expired snapshots tables
-        self.snapshot_evaluator.cleanup(
-            target_snapshots=cleanup_targets,
-            on_complete=self.console.update_cleanup_progress,
+        delete_expired_snapshots(
+            self.state_sync,
+            self.snapshot_evaluator,
+            current_ts=current_ts,
+            ignore_ttl=ignore_ttl,
+            console=self.console,
+            batch_size=self.config.janitor.expired_snapshots_batch_size,
         )
-
-        # Delete the expired snapshot records from the state sync
-        self.state_sync.delete_expired_snapshots(ignore_ttl=ignore_ttl, current_ts=current_ts)
-
         self.state_sync.compact_intervals()
 
     def _cleanup_environments(self, current_ts: t.Optional[int] = None) -> None:
@@ -2874,7 +2935,7 @@ class GenericContext(BaseContext, t.Generic[C]):
     def _new_selector(
         self, models: t.Optional[UniqueKeyDict[str, Model]] = None, dag: t.Optional[DAG[str]] = None
     ) -> Selector:
-        return Selector(
+        return self._selector_cls(
             self.state_reader,
             models=models or self._models,
             context_path=self.path,
@@ -3157,18 +3218,34 @@ class GenericContext(BaseContext, t.Generic[C]):
 
         return all_violations
 
-    def load_model_tests(
-        self, tests: t.Optional[t.List[str]] = None, patterns: list[str] | None = None
+    def select_tests(
+        self,
+        tests: t.Optional[t.List[str]] = None,
+        patterns: t.Optional[t.List[str]] = None,
     ) -> t.List[ModelTestMetadata]:
-        # If a set of specific test path(s) are provided, we can use a single loader
-        # since it's not required to walk every tests/ folder in each repo
-        loaders = [self._loaders[0]] if tests else self._loaders
+        """Filter pre-loaded test metadata based on tests and patterns."""
+
+        test_meta = self._model_test_metadata
+
+        if tests:
+            filtered_tests = []
+            for test in tests:
+                if "::" in test:
+                    if test in self._model_test_metadata_fully_qualified_name_index:
+                        filtered_tests.append(
+                            self._model_test_metadata_fully_qualified_name_index[test]
+                        )
+                else:
+                    test_path = Path(test)
+                    if test_path in self._model_test_metadata_path_index:
+                        filtered_tests.extend(self._model_test_metadata_path_index[test_path])
+
+            test_meta = filtered_tests
 
-        model_tests = []
-        for loader in loaders:
-            model_tests.extend(loader.load_model_tests(tests=tests, patterns=patterns))
+        if patterns:
+            test_meta = filter_tests_by_patterns(test_meta, patterns)
 
-        return model_tests
+        return test_meta
 
 
 class Context(GenericContext[Config]):

sqlmesh/core/dialect.py

@@ -174,6 +174,7 @@ def _parse_id_var(
 
     while (
         identifier
+        and not identifier.args.get("quoted")
         and self._is_connected()
         and (
             self._match_texts(("{", SQLMESH_MACRO_PREFIX))
@@ -349,6 +350,7 @@ def _parse_select(
     parse_subquery_alias: bool = True,
     parse_set_operation: bool = True,
     consume_pipe: bool = True,
+    from_: t.Optional[exp.From] = None,
 ) -> t.Optional[exp.Expression]:
     select = self.__parse_select(  # type: ignore
         nested=nested,
@@ -356,6 +358,7 @@ def _parse_select(
         parse_subquery_alias=parse_subquery_alias,
         parse_set_operation=parse_set_operation,
         consume_pipe=consume_pipe,
+        from_=from_,
     )
 
     if (

sqlmesh/core/engine_adapter/_typing.py

@@ -30,3 +30,5 @@ if t.TYPE_CHECKING:
     ]
 
     QueryOrDF = t.Union[Query, DF]
+    GrantsConfig = t.Dict[str, t.List[str]]
+    DCL = t.TypeVar("DCL", exp.Grant, exp.Revoke)