diff --git a/docs/examples/sqlmesh_cli_crash_course.md b/docs/examples/sqlmesh_cli_crash_course.md index 0bf5780f12..e2a240f22d 100644 --- a/docs/examples/sqlmesh_cli_crash_course.md +++ b/docs/examples/sqlmesh_cli_crash_course.md @@ -738,6 +738,11 @@ You'll use these commands as needed to validate that your changes are behaving a This is a great way to verify that your model's SQL is looking as expected before applying the changes. It is especially important if you're migrating from one query engine to another (ex: postgres to databricks). +In large projects, add `--use-project-index` to load only the model being rendered and its upstream +dependencies. To enable this behavior by default, set +[`render.use_project_index`](../reference/configuration.md#render) to `true` in the project +configuration. + === "SQLMesh" ```bash @@ -1254,4 +1259,4 @@ If you notice you have a lot of old development schemas/data, you can clean them ```bash tcloud sqlmesh janitor - ``` \ No newline at end of file + ``` diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index 41f4b05594..b120efa8a0 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -90,6 +90,7 @@ The `config` sub-module API documentation describes the individual classes used - [Connection configuration](https://sqlmesh.readthedocs.io/en/latest/_readthedocs/html/sqlmesh/core/config/connection.html) (separate classes for each supported database/engine) - [Scheduler configuration](https://sqlmesh.readthedocs.io/en/latest/_readthedocs/html/sqlmesh/core/config/scheduler.html) (separate classes for each supported scheduler) - [Plan change categorization configuration](https://sqlmesh.readthedocs.io/en/latest/_readthedocs/html/sqlmesh/core/config/categorizer.html#CategorizerConfig): `CategorizerConfig()` +- [Render configuration](https://sqlmesh.readthedocs.io/en/latest/_readthedocs/html/sqlmesh/core/config/render.html): `RenderConfig()` - [User configuration](https://sqlmesh.readthedocs.io/en/latest/_readthedocs/html/sqlmesh/core/user.html#User): `User()` - [Notification configuration](https://sqlmesh.readthedocs.io/en/latest/_readthedocs/html/sqlmesh/core/notification_target.html) (separate classes for each notification target) @@ -331,7 +332,7 @@ The cache directory is automatically created if it doesn't exist. You can clear #### Project index -The `--use-project-index` option on supported commands maintains a persistent model dependency index in the cache directory. Each project writes a file named `__model_index.json`. +The `--use-project-index` option on the `lint`, `plan`, and `render` commands maintains a persistent model dependency index in the cache directory. Each project writes a file named `__model_index.json`. The option can be enabled by default for each command with `linter.use_project_index`, `plan.use_project_index`, or `render.use_project_index`, respectively. A full project load with the option enabled creates or refreshes the index. SQLMesh invalidates it when relevant configuration, gateway, macro, audit, or signal metadata changes, or when the set of model files changes. If the index is missing, invalid, or stale, SQLMesh safely falls back to a full project load and rebuilds it. @@ -339,6 +340,33 @@ For operations targeting selected models, the index allows SQLMesh to load only In multi-repository projects, dependencies that cross project boundaries may not be represented by an individual project's index. SQLMesh detects incomplete scoped loads and falls back to loading the full configured project set. +#### Indexed planning + +Indexed planning also reuses snapshot state already loaded while building the plan and scopes graph +work to changed or selected model lineage. It does not change the resulting plan. Enable it by +default for the CLI and Python API with `plan.use_project_index`: + +=== "YAML" + + ```yaml linenums="1" + plan: + use_project_index: true + ``` + +=== "Python" + + ```python linenums="1" + from sqlmesh.core.config import Config, ModelDefaultsConfig, PlanConfig + + config = Config( + model_defaults=ModelDefaultsConfig(dialect="duckdb"), + plan=PlanConfig(use_project_index=True), + ) + ``` + +`Context.plan` and `Context.plan_builder` use this configuration value when `use_project_index` is +omitted. Passing `use_project_index=False` explicitly disables indexed planning for that API call. + ### Table/view storage locations SQLMesh creates schemas, physical tables, and views in the data warehouse/engine. Learn more about why and how SQLMesh creates schema in the ["Why does SQLMesh create schemas?" FAQ](../faq/faq.md#schema-question). @@ -1503,6 +1531,35 @@ SQLMesh provides a linter that checks for potential issues in your models' code. Learn more about linting configuration in the [linting guide](./linter.md). +### Rendering + +By default, `sqlmesh render` loads every model in the project. In large projects, you can use the +persistent project index to load only the model being rendered and its transitive upstream +dependencies. Enable indexed rendering for an individual command with `--use-project-index`, or +make it the project default with the `render.use_project_index` configuration option. + +=== "YAML" + + ```yaml linenums="1" + render: + use_project_index: true + ``` + +=== "Python" + + ```python linenums="1" + from sqlmesh.core.config import Config, ModelDefaultsConfig, RenderConfig + + config = Config( + model_defaults=ModelDefaultsConfig(dialect="duckdb"), + render=RenderConfig(use_project_index=True), + ) + ``` + +`Context.render` uses the configured value when `use_project_index` is omitted. Passing +`use_project_index=False` explicitly disables indexed rendering for that API call. See the +[`render` CLI reference](../reference/cli.md#render) for the other rendering options. + ### Debug mode To enable debug mode set the `SQLMESH_DEBUG` environment variable to one of the following values: "1", "true", "t", "yes" or "y". diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 82b4161277..9787d4252e 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -405,6 +405,17 @@ Options: versions of the models and standalone audits. --explain Explain the plan instead of applying it. + --ignore-cron Run all missing intervals, ignoring + individual cron schedules. Only applies if + --run is set. + --min-intervals TEXT For every model, ensure at least this many + intervals are covered by a missing intervals + check regardless of the plan start date + --use-project-index Refresh the persistent project index, reuse + loaded snapshot state, and scope plan graph + work to changed or selected model lineage + without changing the plan result. Can also + be enabled with plan.use_project_index. -v, --verbose Verbose output. Use -vv for very verbose output. --help Show this message and exit. @@ -447,6 +458,10 @@ Options: only they will be expanded as raw queries. --dialect TEXT The SQL dialect to render the query as. --no-format Disable fancy formatting of the query. + --use-project-index Use the persistent project index to load and + render only the target model and its upstream + dependencies. Can also be enabled with + render.use_project_index. --max-text-width INTEGER The max number of characters in a segment before creating new lines in pretty mode. --leading-comma Determines whether or not the comma is leading diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index a1bf400c32..f89c361b86 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -67,6 +67,12 @@ See all the keys allowed in `model_defaults` at the [model configuration referen | `linter.enabled` | Whether linting is enabled (Default: `False`) | boolean | N | | `linter.use_project_index` | Whether to use the persistent project index for linting. Targeted linting loads selected models and their upstream dependencies. (Default: `False`) | boolean | N | +### Render + +| Option | Description | Type | Required | +|----------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------|----------| +| `render.use_project_index` | Whether to use the persistent project index when rendering. Only the target model and its upstream dependencies are loaded. (Default: `False`) | boolean | N | + ### Variables The `variables` key can be used to provide values for user-defined variables, accessed using the [`@VAR` macro function](../concepts/macros/sqlmesh_macros.md#global-variables) in SQL model definitions, [`context.var` method](../concepts/models/python_models.md#global-variables) in Python model definitions, and [`evaluator.var` method](../concepts/macros/sqlmesh_macros.md#accessing-global-variable-values) in Python macro functions. @@ -102,6 +108,7 @@ Configuration for the `sqlmesh plan` command. | `no_diff` | Don't show diffs for changed models (Default: False) | boolean | N | | `no_prompts` | Disables interactive prompts in CLI (Default: True) | boolean | N | | `always_recreate_environment` | Always recreates the target environment from the environment specified in `create_from` (by default `prod`) (Default: False) | boolean | N | +| `use_project_index` | Whether to refresh the persistent project index, reuse loaded snapshot state, and scope plan graph work to changed or selected model lineage without changing the plan result. (Default: `False`) | boolean | N | ## Run diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py index c2dc1e3dbb..182cf6480e 100644 --- a/sqlmesh/cli/main.py +++ b/sqlmesh/cli/main.py @@ -64,6 +64,13 @@ def _sqlmesh_version() -> str: return "0.0.0" +def _raise_if_no_models(context: Context, path: t.Any) -> None: + if not context.models: + raise click.ClickException( + f"`{path}` doesn't seem to have any models... cd into the proper directory or specify the path(s) with -p." + ) + + @click.group(cls=_SQLMeshGroup, no_args_is_help=True) @click.version_option(version=_sqlmesh_version(), message="%(version)s") @opt.paths @@ -141,8 +148,8 @@ def cli( if ctx.invoked_subcommand in SKIP_LOAD_COMMANDS: load = False - # Unlike the other commands above, lint can scope its own load for multi-project contexts. - if ctx.invoked_subcommand == "lint": + # These commands can scope their own load for multi-project contexts. + if ctx.invoked_subcommand in ("lint", "plan", "render"): load = False configs = load_configs(config, Context.CONFIG_TYPE, paths, dotenv_path=dotenv) @@ -163,10 +170,8 @@ def cli( logger.exception("Failed to initialize SQLMesh context") raise - if load and not context.models: - raise click.ClickException( - f"`{paths}` doesn't seem to have any models... cd into the proper directory or specify the path(s) with -p." - ) + if load: + _raise_if_no_models(context, paths) ctx.obj = context @@ -284,6 +289,12 @@ def init( help="The SQL dialect to render the query as.", ) @click.option("--no-format", is_flag=True, help="Disable fancy formatting of the query.") +@click.option( + "--use-project-index", + is_flag=True, + default=None, + help="Use the persistent project index to load and render only the target model and its upstream dependencies. Can also be enabled with render.use_project_index.", +) @opt.format_options @click.pass_context @error_handler @@ -297,19 +308,20 @@ def render( expand: t.Optional[t.Union[bool, t.Iterable[str]]] = None, dialect: t.Optional[str] = None, no_format: bool = False, + use_project_index: t.Optional[bool] = None, **format_kwargs: t.Any, ) -> None: """Render a model's query, optionally expanding referenced models.""" - model = ctx.obj.get_model(model, raise_if_missing=True) - rendered = ctx.obj.render( model, start=start, end=end, execution_time=execution_time, expand=expand, + use_project_index=use_project_index, ) + model = ctx.obj.get_model(model, raise_if_missing=True) format_config = ctx.obj.config_for_node(model).format format_kwargs = { **format_config.generator_options, @@ -561,6 +573,12 @@ def diff(ctx: click.Context, environment: t.Optional[str] = None) -> None: default=None, help="For every model, ensure at least this many intervals are covered by a missing intervals check regardless of the plan start date", ) +@click.option( + "--use-project-index", + is_flag=True, + default=None, + help="Refresh the persistent project index, reuse loaded snapshot state, and scope plan graph work to changed or selected model lineage without changing the plan result. Can also be enabled with plan.use_project_index.", +) @opt.verbose @click.pass_context @error_handler @@ -579,8 +597,18 @@ def plan( allow_additive_models = kwargs.pop("allow_additive_model") or None backfill_models = kwargs.pop("backfill_model") or None ignore_cron = kwargs.pop("ignore_cron") or None + use_project_index = kwargs.pop("use_project_index") setattr(get_console(), "verbosity", Verbosity(verbose)) + context.load( + use_project_index=( + context.config.plan.use_project_index + if use_project_index is None + else use_project_index + ) + ) + _raise_if_no_models(context, context.path) + context.plan( environment, restate_models=restate_models, @@ -589,6 +617,7 @@ def plan( allow_additive_models=allow_additive_models, backfill_models=backfill_models, ignore_cron=ignore_cron, + use_project_index=use_project_index, **kwargs, ) @@ -1239,10 +1268,7 @@ def lint( use_project_index=use_project_index, ) - if not obj.models: - raise click.ClickException( - f"`{obj.path}` doesn't seem to have any models... cd into the proper directory or specify the path(s) with -p." - ) + _raise_if_no_models(obj, obj.path) @cli.group(no_args_is_help=True) diff --git a/sqlmesh/core/config/__init__.py b/sqlmesh/core/config/__init__.py index 50d2d9a5a2..83d958eb7e 100644 --- a/sqlmesh/core/config/__init__.py +++ b/sqlmesh/core/config/__init__.py @@ -37,6 +37,7 @@ from sqlmesh.core.config.naming import NameInferenceConfig as NameInferenceConfig from sqlmesh.core.config.linter import LinterConfig as LinterConfig from sqlmesh.core.config.plan import PlanConfig as PlanConfig +from sqlmesh.core.config.render import RenderConfig as RenderConfig from sqlmesh.core.config.root import Config as Config, DbtConfig as DbtConfig from sqlmesh.core.config.run import RunConfig as RunConfig from sqlmesh.core.config.scheduler import BuiltInSchedulerConfig as BuiltInSchedulerConfig diff --git a/sqlmesh/core/config/plan.py b/sqlmesh/core/config/plan.py index df1ca44873..93b43f25c0 100644 --- a/sqlmesh/core/config/plan.py +++ b/sqlmesh/core/config/plan.py @@ -21,6 +21,8 @@ class PlanConfig(BaseConfig): use_finalized_state: Whether to compare against the latest finalized environment state, or to use whatever state the target environment is currently in. always_recreate_environment: Whether to always recreate the target environment from the `create_from` environment. + use_project_index: Whether to use the persistent project index and related planning + optimizations. """ forward_only: bool = False @@ -32,3 +34,4 @@ class PlanConfig(BaseConfig): auto_apply: bool = False use_finalized_state: bool = False always_recreate_environment: bool = False + use_project_index: bool = False diff --git a/sqlmesh/core/config/render.py b/sqlmesh/core/config/render.py new file mode 100644 index 0000000000..0b5ef84474 --- /dev/null +++ b/sqlmesh/core/config/render.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from sqlmesh.core.config.base import BaseConfig + + +class RenderConfig(BaseConfig): + """Configuration for rendering model queries. + + Args: + use_project_index: Whether to use the persistent project index when rendering. + """ + + use_project_index: bool = False diff --git a/sqlmesh/core/config/root.py b/sqlmesh/core/config/root.py index b36b7dadc1..9d1971b93a 100644 --- a/sqlmesh/core/config/root.py +++ b/sqlmesh/core/config/root.py @@ -35,6 +35,7 @@ from sqlmesh.core.config.naming import NameInferenceConfig as NameInferenceConfig from sqlmesh.core.config.linter import LinterConfig as LinterConfig from sqlmesh.core.config.plan import PlanConfig +from sqlmesh.core.config.render import RenderConfig from sqlmesh.core.config.run import RunConfig from sqlmesh.core.config.dbt import DbtConfig from sqlmesh.core.config.scheduler import ( @@ -141,6 +142,7 @@ class Config(BaseConfig): format: The formatting options for SQL code. ui: The UI configuration for SQLMesh. plan: The plan configuration. + render: The render configuration. migration: The migration configuration. variables: A dictionary of variables that can be used in models / macros. disable_anonymized_analytics: Whether to disable the anonymized analytics collection. @@ -183,6 +185,7 @@ class Config(BaseConfig): format: FormatConfig = FormatConfig() ui: UIConfig = UIConfig() plan: PlanConfig = PlanConfig() + render: RenderConfig = RenderConfig() migration: MigrationConfig = MigrationConfig() model_naming: NameInferenceConfig = NameInferenceConfig() variables: t.Dict[str, t.Any] = {} @@ -208,6 +211,7 @@ class Config(BaseConfig): "ui": UpdateStrategy.NESTED_UPDATE, "loader_kwargs": UpdateStrategy.KEY_UPDATE, "plan": UpdateStrategy.NESTED_UPDATE, + "render": UpdateStrategy.NESTED_UPDATE, "before_all": UpdateStrategy.EXTEND, "after_all": UpdateStrategy.EXTEND, "linter": UpdateStrategy.NESTED_UPDATE, diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py index e335253016..2c9259ecd8 100644 --- a/sqlmesh/core/context.py +++ b/sqlmesh/core/context.py @@ -99,6 +99,7 @@ Snapshot, SnapshotEvaluator, SnapshotFingerprint, + SnapshotId, missing_intervals, to_table_mapping, ) @@ -730,21 +731,18 @@ def load( self._update_model_schemas_and_validate(model_fqns) return self - # Load environment statements from state for projects not in current load + uncached = set() + if self._load_state and any(self._projects): prod = self.state_reader.get_environment(c.PROD) + if prod: + # Load environment statements from state for projects not in current load existing_statements = self.state_reader.get_environment_statements(c.PROD) for stmt in existing_statements: if stmt.project and stmt.project not in self._projects: self._environment_statements.append(stmt) - uncached = set() - - if self._load_state and any(self._projects): - prod = self.state_reader.get_environment(c.PROD) - - if prod: for snapshot in self.state_reader.get_snapshots(prod.snapshots).values(): if snapshot.node.project in self._projects: uncached.add(snapshot.name) @@ -1199,6 +1197,7 @@ def render( end: t.Optional[TimeLike] = None, execution_time: t.Optional[TimeLike] = None, expand: t.Union[bool, t.Iterable[str]] = False, + use_project_index: t.Optional[bool] = None, **kwargs: t.Any, ) -> exp.Expr: """Renders a model's query, expanding macros with provided kwargs, and optionally expanding referenced models. @@ -1211,11 +1210,23 @@ def render( expand: Whether or not to use expand materialized models, defaults to False. If True, all referenced models are expanded as raw queries. If a list, only referenced models are expanded as raw queries. + use_project_index: Whether to use the persistent project index to load and + render only the target model and its transitive upstream dependencies. If + omitted, the value of ``render.use_project_index`` is used. Returns: The rendered expression. """ execution_time = execution_time or now() + use_project_index = ( + self.config.render.use_project_index if use_project_index is None else use_project_index + ) + + if not self._loaded: + target_fqns = ( + {self._node_or_snapshot_to_fqn(model_or_snapshot)} if use_project_index else None + ) + self.load(model_fqns=target_fqns, use_project_index=use_project_index) model = self.get_model(model_or_snapshot, raise_if_missing=True) @@ -1245,7 +1256,19 @@ def render( ) return next(pandas_to_sql(t.cast(pd.DataFrame, df), model.columns_to_types)) - snapshots = self.snapshots + if use_project_index: + # Only the target model and its transitive upstream dependencies can be referenced + # by the rendered query, so there is no need to create snapshots for the rest. + upstream_fqns = {model.fqn, *self.dag.upstream(model.fqn)} + upstream_models: UniqueKeyDict[str, Model] = UniqueKeyDict( + "models", {fqn: m for fqn, m in self._models.items() if fqn in upstream_fqns} + ) + snapshots = self._snapshots( + upstream_models, + include_standalone_audits=False, + ) + else: + snapshots = self.snapshots deployability_index = DeployabilityIndex.create(snapshots.values(), start=start) return model.render_query_or_raise( @@ -1438,6 +1461,7 @@ def plan( explain: t.Optional[bool] = None, ignore_cron: t.Optional[bool] = None, min_intervals: t.Optional[int] = None, + use_project_index: t.Optional[bool] = None, ) -> Plan: """Interactively creates a plan. @@ -1487,6 +1511,10 @@ def plan( explain: Whether to explain the plan instead of applying it. 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 + use_project_index: Whether to refresh the persistent project index, reuse loaded + snapshot state, and scope plan graph work to changed or selected model lineage. + If omitted, the value of ``plan.use_project_index`` is used. This optimization + does not change the resulting plan. Returns: The populated Plan object. @@ -1518,6 +1546,7 @@ def plan( explain=explain, ignore_cron=ignore_cron, min_intervals=min_intervals, + use_project_index=use_project_index, ) plan = plan_builder.build() @@ -1575,6 +1604,7 @@ def plan_builder( ignore_cron: t.Optional[bool] = None, min_intervals: t.Optional[int] = None, always_include_local_changes: t.Optional[bool] = None, + use_project_index: t.Optional[bool] = None, ) -> PlanBuilder: """Creates a plan builder. @@ -1617,6 +1647,10 @@ def plan_builder( 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. + use_project_index: Whether to refresh the persistent project index, reuse loaded + snapshot state, and scope plan graph work to changed or selected model lineage. + If omitted, the value of ``plan.use_project_index`` is used. This optimization + does not change the resulting plan. Returns: The plan builder. @@ -1655,6 +1689,13 @@ def plan_builder( k: v for k, v in kwargs.items() if v is not None } + use_project_index = ( + self.config.plan.use_project_index if use_project_index is None else use_project_index + ) + + if not self._loaded: + self.load(use_project_index=use_project_index) + skip_tests = explain or skip_tests or False no_gaps = no_gaps or False skip_backfill = skip_backfill or False @@ -1751,7 +1792,7 @@ def plan_builder( else: force_no_diff = not always_include_local_changes - snapshots = self._snapshots(models_override) + snapshots, stored_snapshot_ids = self._snapshots_and_stored_ids(models_override) context_diff = self._context_diff( environment or c.PROD, snapshots=snapshots, @@ -1760,6 +1801,7 @@ def plan_builder( ensure_finalized_snapshots=self.config.plan.use_finalized_state, diff_rendered=diff_rendered, always_recreate_environment=self.config.plan.always_recreate_environment, + stored_snapshot_ids=stored_snapshot_ids if use_project_index else None, ) modified_model_names = { *context_diff.modified_snapshots, @@ -1881,6 +1923,7 @@ def plan_builder( }, explain=explain or False, ignore_cron=ignore_cron or False, + scope_to_changed_lineage=use_project_index, ) def apply( @@ -3009,9 +3052,24 @@ def _get_engine_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter: return self.engine_adapter def _snapshots( - self, models_override: t.Optional[UniqueKeyDict[str, Model]] = None + self, + models_override: t.Optional[UniqueKeyDict[str, Model]] = None, + include_standalone_audits: bool = True, ) -> t.Dict[str, Snapshot]: - nodes = {**(models_override or self._models), **self._standalone_audits} + return self._snapshots_and_stored_ids( + models_override, + include_standalone_audits=include_standalone_audits, + )[0] + + def _snapshots_and_stored_ids( + self, + models_override: t.Optional[UniqueKeyDict[str, Model]] = None, + include_standalone_audits: bool = True, + ) -> t.Tuple[t.Dict[str, Snapshot], t.Set[SnapshotId]]: + """Returns the snapshots along with the IDs of those that exist in the state.""" + nodes: t.Dict[str, Node] = dict(models_override or self._models) + if include_standalone_audits: + nodes.update(self._standalone_audits) snapshots = self._nodes_to_snapshots(nodes) stored_snapshots = self.state_reader.get_snapshots(snapshots.values()) @@ -3036,7 +3094,10 @@ def _snapshots( # Keep the original model instance to preserve the query cache. snapshot.node = snapshots[snapshot.name].node - return {name: stored_snapshots.get(s.snapshot_id, s) for name, s in snapshots.items()} + merged_snapshots = { + name: stored_snapshots.get(s.snapshot_id, s) for name, s in snapshots.items() + } + return merged_snapshots, set(stored_snapshots) def _context_diff( self, @@ -3047,6 +3108,7 @@ def _context_diff( ensure_finalized_snapshots: bool = False, diff_rendered: bool = False, always_recreate_environment: bool = False, + stored_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None, ) -> ContextDiff: environment = Environment.sanitize_name(environment) if force_no_diff: @@ -3065,6 +3127,7 @@ def _context_diff( gateway_managed_virtual_layer=self.config.gateway_managed_virtual_layer, infer_python_dependencies=self.config.infer_python_dependencies, always_recreate_environment=always_recreate_environment, + stored_snapshot_ids=stored_snapshot_ids, ) def _destroy(self) -> bool: diff --git a/sqlmesh/core/context_diff.py b/sqlmesh/core/context_diff.py index 047e58609a..0b6aab146a 100644 --- a/sqlmesh/core/context_diff.py +++ b/sqlmesh/core/context_diff.py @@ -104,6 +104,7 @@ def create( gateway_managed_virtual_layer: bool = False, infer_python_dependencies: bool = True, always_recreate_environment: bool = False, + stored_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None, ) -> ContextDiff: """Create a ContextDiff object. @@ -124,6 +125,9 @@ def create( model-specific gateway rather than the default gateway. infer_python_dependencies: Whether to statically analyze Python code to automatically infer Python package requirements. + stored_snapshot_ids: The IDs of the provided snapshots that are known to exist in the state. + When provided, the given snapshots are assumed to have already been hydrated from the state, + so only the previous versions of the modified snapshots are fetched. Returns: The ContextDiff object. @@ -177,9 +181,21 @@ def create( and snapshot.fingerprint != remote_snapshot_name_to_info[snapshot.name].fingerprint } - stored = state_reader.get_snapshots( - [*snapshots.values(), *modified_snapshot_name_to_snapshot_info.values()] - ) + if stored_snapshot_ids is None: + stored = state_reader.get_snapshots( + [*snapshots.values(), *modified_snapshot_name_to_snapshot_info.values()] + ) + else: + # The provided snapshots have already been hydrated from the state, so only the + # previous versions of the modified snapshots need to be fetched. + stored = { + snapshot.snapshot_id: snapshot.copy() + for snapshot in snapshots.values() + if snapshot.snapshot_id in stored_snapshot_ids + } + stored.update( + state_reader.get_snapshots(modified_snapshot_name_to_snapshot_info.values()) + ) merged_snapshots = {} modified_snapshots = {} diff --git a/sqlmesh/core/plan/builder.py b/sqlmesh/core/plan/builder.py index a6307a9ffd..8f39f4eafd 100644 --- a/sqlmesh/core/plan/builder.py +++ b/sqlmesh/core/plan/builder.py @@ -98,6 +98,8 @@ class PlanBuilder: end_override_per_model: A mapping of model FQNs to target end dates. ignore_cron: Whether to ignore the node's cron schedule when computing missing intervals. explain: Whether to explain the plan instead of applying it. + scope_to_changed_lineage: Whether to scope plan graph work to changed or selected + model lineage without changing the resulting plan. """ def __init__( @@ -138,6 +140,7 @@ def __init__( console: t.Optional[PlanBuilderConsole] = None, user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = None, selected_models: t.Optional[t.Set[str]] = None, + scope_to_changed_lineage: bool = False, ): self._context_diff = context_diff self._no_gaps = no_gaps @@ -183,6 +186,7 @@ def __init__( self._choices: t.Dict[SnapshotId, SnapshotChangeCategory] = {} self._user_provided_flags = user_provided_flags self._selected_models = selected_models + self._scope_to_changed_lineage = scope_to_changed_lineage self._explain = explain self._start = start @@ -380,9 +384,65 @@ def build(self) -> Plan: return plan def _build_dag(self) -> DAG[SnapshotId]: + snapshots = self._context_diff.snapshots + if not self._scope_to_changed_lineage or self._restate_all_snapshots: + relevant_snapshot_ids = set(snapshots) + else: + model_fqn_to_snapshot_id = { + snapshot.name: snapshot_id for snapshot_id, snapshot in snapshots.items() + } + relevant_snapshot_ids = { + *self._context_diff.added, + *self._context_diff.new_snapshots, + *(new.snapshot_id for new, _ in self._context_diff.modified_snapshots.values()), + } + + selected_model_names = { + *(self._restate_models or set()), + *(self._backfill_models or set()), + *(self._selected_models or set()), + } + relevant_snapshot_ids.update( + model_fqn_to_snapshot_id[name] + for name in selected_model_names + if name in model_fqn_to_snapshot_id + ) + + children: t.Dict[SnapshotId, t.Set[SnapshotId]] = defaultdict(set) + for s_id, snapshot in snapshots.items(): + for parent_id in snapshot.parents: + if parent_id in snapshots: + children[parent_id].add(s_id) + + root_snapshot_ids = relevant_snapshot_ids.copy() + upstream_stack = list(relevant_snapshot_ids) + while upstream_stack: + s_id = upstream_stack.pop() + for parent_id in snapshots[s_id].parents: + if parent_id in snapshots and parent_id not in relevant_snapshot_ids: + relevant_snapshot_ids.add(parent_id) + upstream_stack.append(parent_id) + + downstream_snapshot_ids = root_snapshot_ids.copy() + downstream_stack = list(root_snapshot_ids) + while downstream_stack: + s_id = downstream_stack.pop() + for child_id in children.get(s_id, set()) - downstream_snapshot_ids: + downstream_snapshot_ids.add(child_id) + relevant_snapshot_ids.add(child_id) + downstream_stack.append(child_id) + dag: DAG[SnapshotId] = DAG() - for s_id, context_snapshot in self._context_diff.snapshots.items(): - dag.add(s_id, context_snapshot.parents) + for s_id in relevant_snapshot_ids: + context_snapshot = snapshots[s_id] + dag.add( + s_id, + ( + parent_id + for parent_id in context_snapshot.parents + if parent_id in relevant_snapshot_ids + ), + ) return dag def _build_restatements( diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index a8ee1aaa39..608c543241 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -179,6 +179,34 @@ def test_plan(runner, tmp_path): assert "sqlmesh_example.incremental_model [insert 2020-01-01 - 2022-12-31]" in result.output +@time_machine.travel(FREEZE_TIME) +def test_plan_use_project_index(runner, tmp_path): + create_example_project(tmp_path) + config_path = tmp_path / "config.yaml" + config_path.write_text( + config_path.read_text(encoding="utf-8").replace( + "plan:\n no_prompts: false", + "plan:\n no_prompts: false\n use_project_index: true", + ), + encoding="utf-8", + ) + + result = runner.invoke( + cli, + [ + "--log-file-dir", + tmp_path, + "--paths", + tmp_path, + "plan", + ], + input="y\n", + ) + + assert_plan_success(result) + assert list((tmp_path / ".cache").glob("*_model_index.json")) + + def test_plan_skip_tests(runner, tmp_path): create_example_project(tmp_path) @@ -2111,6 +2139,19 @@ def test_render(runner: CliRunner, tmp_path: Path): assert expected in cleaned_output + indexed_result = runner.invoke( + cli, + [ + "--paths", + str(tmp_path), + "render", + "sqlmesh_example.full_model", + "--use-project-index", + "--no-format", + ], + ) + assert indexed_result.exit_code == 0 + @time_machine.travel(FREEZE_TIME) def test_signals(runner: CliRunner, tmp_path: Path): diff --git a/tests/core/test_config.py b/tests/core/test_config.py index 0da5b6e22f..28cf2c5df0 100644 --- a/tests/core/test_config.py +++ b/tests/core/test_config.py @@ -16,6 +16,8 @@ ModelDefaultsConfig, BigQueryConnectionConfig, MotherDuckConnectionConfig, + PlanConfig, + RenderConfig, BuiltInSchedulerConfig, EnvironmentSuffixTarget, TableNamingConvention, @@ -70,6 +72,20 @@ def python_config_path(tmp_path_factory) -> Path: return config_path +def test_render_config() -> None: + config = Config.parse_obj({"render": {"use_project_index": True}}) + + assert config.render == RenderConfig(use_project_index=True) + assert Config().update_with(config).render.use_project_index is True + + +def test_plan_project_index_config() -> None: + config = Config.parse_obj({"plan": {"use_project_index": True}}) + + assert config.plan == PlanConfig(use_project_index=True) + assert Config().update_with(config).plan.use_project_index is True + + def test_update_with_gateways(): gateway0_config = GatewayConfig(connection=DuckDBConnectionConfig()) gateway1_config = GatewayConfig(connection=DuckDBConnectionConfig(database="test")) diff --git a/tests/core/test_context.py b/tests/core/test_context.py index 75737f1edb..a2068f0a71 100644 --- a/tests/core/test_context.py +++ b/tests/core/test_context.py @@ -30,6 +30,7 @@ LinterConfig, ModelDefaultsConfig, PlanConfig, + RenderConfig, SnowflakeConnectionConfig, ) from sqlmesh.core.context import Context @@ -276,6 +277,196 @@ def test_render_seed_model(sushi_context, assert_exp_eq): ) +@pytest.mark.slow +def test_render_only_creates_snapshots_for_upstream_models(sushi_context: Context): + model = sushi_context.get_model("sushi.top_waiters", raise_if_missing=True) + upstream_fqns = {model.fqn, *sushi_context.dag.upstream(model.fqn)} + + # Sanity check that the project contains models outside of the target model's subgraph. + assert set(sushi_context.models) - upstream_fqns + + with patch.object( + sushi_context.state_reader, + "get_snapshots", + wraps=sushi_context.state_reader.get_snapshots, + ) as default_get_snapshots_mock: + sushi_context.render("sushi.top_waiters") + + default_requested_names = { + snapshot.name + for call_args in default_get_snapshots_mock.call_args_list + for snapshot in call_args.args[0] + } + assert set(sushi_context.models) <= default_requested_names + + with patch.object( + sushi_context.state_reader, + "get_snapshots", + wraps=sushi_context.state_reader.get_snapshots, + ) as get_snapshots_mock: + sushi_context.render("sushi.top_waiters", use_project_index=True) + + requested_names = { + snapshot.name + for call_args in get_snapshots_mock.call_args_list + for snapshot in call_args.args[0] + } + assert model.fqn in requested_names + assert requested_names == upstream_fqns + + +def test_render_only_loads_upstream_model_files(tmp_path: pathlib.Path) -> None: + create_temp_file( + tmp_path, + pathlib.Path("models", "a.sql"), + "MODEL(name a, kind FULL); SELECT 1 AS col;", + ) + create_temp_file( + tmp_path, + pathlib.Path("models", "b.sql"), + "MODEL(name b, kind FULL); SELECT col FROM a;", + ) + create_temp_file( + tmp_path, + pathlib.Path("models", "c.sql"), + "MODEL(name c, kind FULL); SELECT col FROM b;", + ) + config = Config( + model_defaults=ModelDefaultsConfig(dialect="duckdb"), + render=RenderConfig(use_project_index=True), + ) + + # Populate the persistent model path/dependency index. + Context(config=config, paths=tmp_path, load=False).load(use_project_index=True) + + ctx = Context(config=config, paths=tmp_path, load=False) + loader = t.cast(SqlMeshLoader, ctx._loaders[0]) + with patch.object( + loader, + "_load_sql_models", + wraps=loader._load_sql_models, + ) as load_sql_models_mock: + ctx.render("b") + + selected_paths = load_sql_models_mock.call_args.kwargs["selected_paths"] + assert {path.name for path in selected_paths} == {"a.sql", "b.sql"} + assert set(ctx.models) == { + ctx.get_model("a", raise_if_missing=True).fqn, + ctx.get_model("b", raise_if_missing=True).fqn, + } + + non_indexed_ctx = Context(config=config, paths=tmp_path, load=False) + non_indexed_loader = t.cast(SqlMeshLoader, non_indexed_ctx._loaders[0]) + with patch.object( + non_indexed_loader, + "_load_sql_models", + wraps=non_indexed_loader._load_sql_models, + ) as load_sql_models_mock: + non_indexed_ctx.render("b", use_project_index=False) + + assert load_sql_models_mock.call_args.kwargs["selected_paths"] is None + assert len(non_indexed_ctx.models) == 3 + + +@pytest.mark.slow +def test_plan_builder_fetches_stored_snapshots_once(sushi_context: Context): + sushi_context.upsert_model("sushi.customers", stamp="force a new snapshot version") + snapshot_count = len(sushi_context.snapshots) + + with patch.object( + sushi_context.state_reader, + "get_snapshots", + wraps=sushi_context.state_reader.get_snapshots, + ) as default_get_snapshots_mock: + sushi_context.plan_builder("dev", skip_tests=True, skip_linter=True) + + default_full_fetches = [ + call_args + for call_args in default_get_snapshots_mock.call_args_list + if len(list(call_args.args[0])) >= snapshot_count + ] + + sushi_context.config.plan = PlanConfig(use_project_index=True) + + with patch.object( + sushi_context.state_reader, + "get_snapshots", + wraps=sushi_context.state_reader.get_snapshots, + ) as get_snapshots_mock: + plan_builder = sushi_context.plan_builder( + "dev", + skip_tests=True, + skip_linter=True, + ) + + full_fetches = [ + call_args + for call_args in get_snapshots_mock.call_args_list + if len(list(call_args.args[0])) >= snapshot_count + ] + assert len(full_fetches) < len(default_full_fetches) + + # The diff is unaffected and still reflects the modified model. + context_diff = plan_builder._context_diff + assert '"memory"."sushi"."customers"' in context_diff.modified_snapshots + + +@pytest.mark.slow +def test_project_index_plan_matches_default_plan(sushi_context: Context) -> None: + sushi_context.upsert_model("sushi.customers", stamp="force a new snapshot version") + + sushi_context.config.plan = PlanConfig(use_project_index=True) + + def build_plan(use_project_index: t.Optional[bool] = None) -> Plan: + if use_project_index is None: + plan_builder = sushi_context.plan_builder( + "dev", + start="2023-01-01", + end="2023-01-07", + execution_time="2023-01-08", + skip_tests=True, + skip_linter=True, + ) + else: + plan_builder = sushi_context.plan_builder( + "dev", + start="2023-01-01", + end="2023-01-07", + execution_time="2023-01-08", + skip_tests=True, + skip_linter=True, + use_project_index=use_project_index, + ) + return plan_builder.build() + + default_plan = build_plan(use_project_index=False) + indexed_plan = build_plan() + + assert indexed_plan.directly_modified == default_plan.directly_modified + assert indexed_plan.indirectly_modified == default_plan.indirectly_modified + assert indexed_plan.restatements == default_plan.restatements + assert indexed_plan.models_to_backfill == default_plan.models_to_backfill + assert indexed_plan.missing_intervals == default_plan.missing_intervals + + default_snapshots = { + snapshot.name: snapshot for snapshot in default_plan.context_diff.snapshots.values() + } + indexed_snapshots = { + snapshot.name: snapshot for snapshot in indexed_plan.context_diff.snapshots.values() + } + assert set(indexed_snapshots) == set(default_snapshots) + assert {name: snapshot.change_category for name, snapshot in indexed_snapshots.items()} == { + name: snapshot.change_category for name, snapshot in default_snapshots.items() + } + assert { + name: indexed_plan.deployability_index.is_deployable(snapshot) + for name, snapshot in indexed_snapshots.items() + } == { + name: default_plan.deployability_index.is_deployable(snapshot) + for name, snapshot in default_snapshots.items() + } + + @pytest.mark.slow def test_diff(sushi_context: Context, mocker: MockerFixture): mock_console = mocker.Mock() diff --git a/tests/core/test_plan.py b/tests/core/test_plan.py index add85d2eec..572b69cca9 100644 --- a/tests/core/test_plan.py +++ b/tests/core/test_plan.py @@ -4510,3 +4510,86 @@ def test_forward_only_indirect_change_to_materialized_view(make_snapshot): # Forward-only indirect changes to MVs should not always be classified as indirect breaking. # Instead, we want to preserve the standard categorization. assert snapshot_b_new.change_category == SnapshotChangeCategory.INDIRECT_NON_BREAKING + + +def test_plan_builder_scopes_dag_to_changed_model_lineage(make_snapshot): + snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 1 as id"))) + snapshot_b_old = make_snapshot( + SqlModel(name="b", query=parse_one("select id from a")), + nodes={snapshot_a.name: snapshot_a.model}, + ) + snapshot_b_new = make_snapshot( + SqlModel(name="b", query=parse_one("select id + 1 as id from a")), + nodes={snapshot_a.name: snapshot_a.model}, + ) + snapshot_c = make_snapshot( + SqlModel(name="c", query=parse_one("select id from b")), + nodes={ + snapshot_a.name: snapshot_a.model, + snapshot_b_new.name: snapshot_b_new.model, + }, + ) + snapshot_unrelated = make_snapshot( + SqlModel(name="unrelated", query=parse_one("select 2 as id")) + ) + snapshot_sibling = make_snapshot( + SqlModel(name="sibling", query=parse_one("select id from a")), + nodes={snapshot_a.name: snapshot_a.model}, + ) + + context_diff = ContextDiff( + environment="prod", + is_new_environment=False, + is_unfinalized_environment=False, + normalize_environment_name=True, + create_from="prod", + create_from_env_exists=True, + added=set(), + removed_snapshots={}, + modified_snapshots={ + snapshot_b_new.name: (snapshot_b_new, snapshot_b_old), + }, + snapshots={ + snapshot_a.snapshot_id: snapshot_a, + snapshot_b_new.snapshot_id: snapshot_b_new, + snapshot_c.snapshot_id: snapshot_c, + snapshot_sibling.snapshot_id: snapshot_sibling, + snapshot_unrelated.snapshot_id: snapshot_unrelated, + }, + new_snapshots={snapshot_b_new.snapshot_id: snapshot_b_new}, + previous_plan_id=None, + previously_promoted_snapshot_ids=set(), + previous_finalized_snapshots=None, + previous_gateway_managed_virtual_layer=False, + gateway_managed_virtual_layer=False, + environment_statements=[], + ) + + assert set(PlanBuilder(context_diff)._build_dag()) == set(context_diff.snapshots) + + dag = PlanBuilder( + context_diff, + scope_to_changed_lineage=True, + )._build_dag() + + assert set(dag) == { + snapshot_a.snapshot_id, + snapshot_b_new.snapshot_id, + snapshot_c.snapshot_id, + } + assert snapshot_sibling.snapshot_id not in dag + assert snapshot_unrelated.snapshot_id not in dag + + restatement_dag = PlanBuilder( + context_diff, + restate_models={snapshot_unrelated.name}, + scope_to_changed_lineage=True, + )._build_dag() + assert snapshot_unrelated.snapshot_id in restatement_dag + + backfill_dag = PlanBuilder( + context_diff, + backfill_models={snapshot_unrelated.name}, + scope_to_changed_lineage=True, + )._build_dag() + assert snapshot_unrelated.snapshot_id in backfill_dag