From 5ddff6619da4d1571162c17278c1dc6125519b63 Mon Sep 17 00:00:00 2001 From: Christoph Meier Date: Sat, 4 Jul 2026 23:34:54 -0400 Subject: [PATCH] feat: add --mode flag to create_external_models command Adds a --mode option to create_external_models with three values: - overwrite (default): current behavior, replaces same-gateway entries - sync: preserves hand-edited metadata, only syncs columns from DB, warns on stale entries but keeps them - sync_prune: same as sync but removes stale entries Introduces CreateExternalModelsMode enum following the codebase's (str, Enum) convention with is_* properties and a default classproperty. Fixes a bug where external models loaded from the existing YAML file were always re-fetched from the database, preventing stale detection. The new logic only fetches FQNs referenced by non-external model depends_on, excluding known internal models to prevent conversion. Signed-off-by: Christoph Meier --- docs/concepts/models/external_models.md | 6 +- docs/reference/cli.md | 10 +- docs/reference/notebook.md | 2 +- sqlmesh/cli/main.py | 8 + sqlmesh/core/context.py | 17 +- sqlmesh/core/schema_loader.py | 121 +++++++++++-- sqlmesh/magics.py | 10 +- tests/core/test_schema_loader.py | 230 +++++++++++++++++++++++- 8 files changed, 381 insertions(+), 23 deletions(-) diff --git a/docs/concepts/models/external_models.md b/docs/concepts/models/external_models.md index ef2b39a10c..1e154969b3 100644 --- a/docs/concepts/models/external_models.md +++ b/docs/concepts/models/external_models.md @@ -52,6 +52,8 @@ If SQLMesh does not have access to an external table's metadata, the table will `create_external_models` solely queries SQL engine metadata and does not query external tables themselves. +By default, the command overwrites same-gateway entries with freshly fetched columns and types. Use `--mode sync` to preserve any hand-edited metadata (such as descriptions, audits, or start dates) while only updating columns from the database. Use `--mode sync_prune` to additionally remove entries that are no longer referenced by any model. + ### Gateway-specific external models In some use-cases such as [isolated systems with multiple gateways](../../guides/isolated_systems.md#multiple-gateways), there are external models that only exist on a certain gateway. @@ -116,9 +118,9 @@ The file can be constructed by hand using a standard text editor or IDE. Sometimes, SQLMesh cannot infer the structure of a model and you need to add it manually. -However, since `sqlmesh create_external_models` replaces the `external_models.yaml` file, any manual changes you made to that file will be overwritten. +By default, running `sqlmesh create_external_models` replaces the entries in `external_models.yaml` — any manual changes to that file may be lost. To preserve hand-edited metadata such as descriptions and audits, pass `--mode sync`, which only updates columns from the database while keeping all other fields. -The solution is to create the manual model definition files in the `external_models/` directory, like so: +Alternatively, you can place manual model definitions in the `external_models/` directory, like so: ``` external_models.yaml diff --git a/docs/reference/cli.md b/docs/reference/cli.md index b65f8256ac..fbfddfc435 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -106,7 +106,15 @@ Usage: sqlmesh create_external_models [OPTIONS] Create a schema file containing external model schemas. Options: - --help Show this message and exit. + --strict Raise an error if the external model is + missing in the database + --mode [overwrite|sync|sync_prune] + The mode for updating external models. + overwrite replaces all entries (default), sync + syncs columns while preserving metadata and + warns on stale entries, sync_prune also + removes stale entries. + --help Show this message and exit. ``` ## create_test diff --git a/docs/reference/notebook.md b/docs/reference/notebook.md index 6cac4e1078..fe107ba747 100644 --- a/docs/reference/notebook.md +++ b/docs/reference/notebook.md @@ -304,7 +304,7 @@ Migrate SQLMesh to the current running version #### create_external_models ``` -%create_external_models +%create_external_models [--strict] [--mode {overwrite,sync,sync_prune}] Create a schema file containing external model schemas. ``` diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py index b3c7a7027b..7d8ff96328 100644 --- a/sqlmesh/cli/main.py +++ b/sqlmesh/cli/main.py @@ -977,6 +977,14 @@ def rollback(obj: Context) -> None: is_flag=True, help="Raise an error if the external model is missing in the database", ) +@click.option( + "--mode", + type=click.Choice(["overwrite", "sync", "sync_prune"], case_sensitive=False), + default="overwrite", + help="The mode for updating external models. overwrite replaces all entries (default), " + "sync syncs columns while preserving metadata and warns on stale entries, " + "sync_prune also removes stale entries.", +) @click.pass_obj @error_handler @cli_analytics diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py index 5902977331..97691be2f6 100644 --- a/sqlmesh/core/context.py +++ b/sqlmesh/core/context.py @@ -92,7 +92,10 @@ 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.schema_loader import ( + CreateExternalModelsMode, + create_external_models_file, +) from sqlmesh.core.selector import Selector, NativeSelector from sqlmesh.core.snapshot import ( DeployabilityIndex, @@ -2486,7 +2489,11 @@ def rollback(self) -> None: self._new_state_sync().rollback() @python_api_analytics - def create_external_models(self, strict: bool = False) -> None: + def create_external_models( + self, + strict: bool = False, + mode: t.Union[CreateExternalModelsMode, str] = CreateExternalModelsMode.default, + ) -> None: """Create a file to document the schema of external models. The external models file contains all columns and types of external models, allowing for more @@ -2494,7 +2501,12 @@ def create_external_models(self, strict: bool = False) -> None: Args: strict: If True, raise an error if the external model is missing in the database. + mode: The mode for updating external models. "overwrite" replaces all entries (default), + "sync" syncs columns while preserving metadata and warns on stale entries, + "sync_prune" also removes stale entries. """ + if isinstance(mode, str): + mode = CreateExternalModelsMode(mode) if not self._models: self.load(update_schemas=False) @@ -2529,6 +2541,7 @@ def create_external_models(self, strict: bool = False) -> None: max_workers=self.concurrent_tasks, strict=strict, all_models=self._models, + mode=mode, ) @python_api_analytics diff --git a/sqlmesh/core/schema_loader.py b/sqlmesh/core/schema_loader.py index 3fed8c5c09..6ff55e92de 100644 --- a/sqlmesh/core/schema_loader.py +++ b/sqlmesh/core/schema_loader.py @@ -2,6 +2,7 @@ import typing as t from concurrent.futures import ThreadPoolExecutor +from enum import Enum from pathlib import Path from sqlglot import exp @@ -11,10 +12,34 @@ from sqlmesh.core.engine_adapter import EngineAdapter from sqlmesh.core.model.definition import Model from sqlmesh.core.state_sync import StateReader -from sqlmesh.utils import UniqueKeyDict, yaml +from sqlmesh.utils import UniqueKeyDict, classproperty, yaml from sqlmesh.utils.errors import SQLMeshError +class CreateExternalModelsMode(str, Enum): + """Determines how existing entries are handled when creating external models.""" + + OVERWRITE = "overwrite" + SYNC = "sync" + SYNC_PRUNE = "sync_prune" + + @classproperty + def default(cls) -> CreateExternalModelsMode: + return CreateExternalModelsMode.OVERWRITE + + @property + def is_overwrite(self) -> bool: + return self == CreateExternalModelsMode.OVERWRITE + + @property + def is_sync(self) -> bool: + return self == CreateExternalModelsMode.SYNC + + @property + def is_sync_prune(self) -> bool: + return self == CreateExternalModelsMode.SYNC_PRUNE + + def create_external_models_file( path: Path, models: UniqueKeyDict[str, Model], @@ -25,6 +50,7 @@ def create_external_models_file( max_workers: int = 1, strict: bool = False, all_models: t.Optional[t.Dict[str, Model]] = None, + mode: CreateExternalModelsMode = CreateExternalModelsMode.default, ) -> None: """Create or replace a YAML file with column and types of all columns in all external models. @@ -40,16 +66,18 @@ def create_external_models_file( all_models: FQN to model across all loaded repos. When provided, a dependency is only classified as external if it is absent from this full set. This prevents cross-repo internal models from being misclassified as external in multi-repo setups. + mode: The mode for updating external models. "overwrite" replaces all entries (default), + "sync" syncs columns while preserving metadata and warns on stale entries, + "sync_prune" also removes stale entries. """ known_models: t.Dict[str, Model] = all_models if all_models is not None else models - external_model_fqns = set() - for fqn, model in models.items(): - if model.kind.is_external: - external_model_fqns.add(fqn) - for dep in model.depends_on: - if dep not in known_models: - external_model_fqns.add(dep) + external_model_fqns = { + dep + for model in models.values() + for dep in model.depends_on + if dep not in known_models or known_models[dep].kind.is_external + } # Make sure we don't convert internal models into external ones. existing_model_fqns = state_reader.nodes_exist(external_model_fqns, exclude_external=True) @@ -79,15 +107,78 @@ def create_external_models_file( if columns ] - # dont clobber existing entries from other gateways - entries_to_keep = ( - [e for e in yaml.load(path) if e.get("gateway", None) != gateway] - if path.exists() - else [] + if mode.is_overwrite: + # dont clobber existing entries from other gateways + entries_to_keep = ( + [e for e in yaml.load(path) if e.get("gateway", None) != gateway] + if path.exists() + else [] + ) + + with open(path, "w", encoding="utf-8") as file: + yaml.dump(entries_to_keep + schemas, file) + else: + _write_external_models_file_update(path, schemas, gateway, mode=mode) + + +def _write_external_models_file_update( + path: Path, + schemas: t.List[t.Dict[str, t.Any]], + gateway: t.Optional[str], + mode: CreateExternalModelsMode, +) -> None: + """Write external models file for sync and sync_prune modes. + + Preserves hand-edited metadata on existing entries, only syncing columns from the DB. + sync mode warns about stale entries but keeps them. + sync_prune mode removes stale entries. + """ + existing_entries = yaml.load(path) if path.exists() else [] + + other_gateway_entries: t.List[t.Dict[str, t.Any]] = [] + same_gateway_entries: t.List[t.Dict[str, t.Any]] = [] + + for entry in existing_entries: + if entry.get("gateway", None) != gateway: + other_gateway_entries.append(entry) + else: + same_gateway_entries.append(entry) + + same_gateway_by_name: t.Dict[str, t.Dict[str, t.Any]] = {} + for entry in same_gateway_entries: + same_gateway_by_name[entry["name"]] = entry + + result = list(other_gateway_entries) + matched_in_schemas: t.Set[str] = {s["name"] for s in schemas} + + for new_schema in schemas: + name = new_schema["name"] + if name in same_gateway_by_name: + existing = same_gateway_by_name[name] + existing["columns"] = new_schema["columns"] + result.append(existing) + else: + result.append(new_schema) + + stale_count = 0 + for entry in same_gateway_entries: + name = entry["name"] + if name not in matched_in_schemas: + stale_count += 1 + if mode.is_sync: + get_console().log_warning( + f"External model '{name}' is no longer referenced but was preserved. " + "Use --mode sync_prune to automatically remove stale entries." + ) + result.append(entry) + + if stale_count and mode.is_sync_prune: + get_console().log_warning( + f"Removed {stale_count} stale external model(s) that are no longer referenced." ) - with open(path, "w", encoding="utf-8") as file: - yaml.dump(entries_to_keep + schemas, file) + with open(path, "w", encoding="utf-8") as file: + yaml.dump(result, file) def get_columns( diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py index 3a59fc4f7b..7a8b01fc71 100644 --- a/sqlmesh/magics.py +++ b/sqlmesh/magics.py @@ -738,12 +738,20 @@ def migrate(self, context: Context, line: str) -> None: action="store_true", help="Raise an error if the external model is missing in the database", ) + @argument( + "--mode", + choices=["overwrite", "sync", "sync_prune"], + default="overwrite", + help="The mode for updating external models. overwrite replaces all entries (default), " + "sync syncs columns while preserving metadata and warns on stale entries, " + "sync_prune also removes stale entries.", + ) @line_magic @pass_sqlmesh_context def create_external_models(self, context: Context, line: str) -> None: """Create a schema file containing external model schemas.""" args = parse_argstring(self.create_external_models, line) - context.create_external_models(strict=args.strict) + context.create_external_models(strict=args.strict, mode=args.mode) @magic_arguments() @argument( diff --git a/tests/core/test_schema_loader.py b/tests/core/test_schema_loader.py index 8b944793be..2d5992256e 100644 --- a/tests/core/test_schema_loader.py +++ b/tests/core/test_schema_loader.py @@ -9,11 +9,15 @@ from sqlmesh.core import constants as c from sqlmesh.core.config import Config, DuckDBConnectionConfig, GatewayConfig +from sqlmesh.core.console import get_console from sqlmesh.core.context import Context from sqlmesh.core.dialect import parse from sqlmesh.core.model import SqlModel, create_external_model, load_sql_based_model from sqlmesh.core.model.definition import ExternalModel -from sqlmesh.core.schema_loader import create_external_models_file +from sqlmesh.core.schema_loader import ( + CreateExternalModelsMode, + create_external_models_file, +) from sqlmesh.core.snapshot import SnapshotChangeCategory from sqlmesh.utils import yaml from sqlmesh.utils.errors import SQLMeshError @@ -342,6 +346,230 @@ def test_no_internal_model_conversion(tmp_path: Path, mocker: MockerFixture): create_external_model(**row, dialect="bigquery") +def test_create_external_models_sync_preserves_metadata(tmp_path: Path, mocker: MockerFixture): + initial_entries = [ + { + "name": '"schema"."existing_table"', + "description": "A table with customer data", + "dialect": "duckdb", + "start": "1 week ago", + "audits": [ + {"name": "not_null", "columns": "[customer_id]"}, + ], + "columns": {"a": "int"}, + }, + { + "name": '"schema"."other_table"', + "description": "Another table", + "columns": {"x": "text"}, + }, + ] + with open(tmp_path / c.EXTERNAL_MODELS_YAML, "w", encoding="utf-8") as fd: + yaml.dump(initial_entries, fd) + + engine_adapter_mock = mocker.Mock() + engine_adapter_mock.columns.return_value = { + "id": exp.DataType.build("bigint"), + "name": exp.DataType.build("text"), + } + + state_reader_mock = mocker.Mock() + state_reader_mock.nodes_exist.return_value = set() + + model = SqlModel( + name="a", + query=parse_one("select * FROM schema.existing_table, schema.new_table"), + ) + + filename = tmp_path / c.EXTERNAL_MODELS_YAML + + with patch.object(get_console(), "log_warning") as mock_warning: + create_external_models_file( + filename, + {"a": model}, # type: ignore + engine_adapter_mock, + state_reader_mock, + "", + mode=CreateExternalModelsMode.SYNC, + ) + + schema = yaml.load(filename) + + assert len(schema) == 3 + + existing_entry = [e for e in schema if e["name"] == '"schema"."existing_table"'][0] + assert existing_entry["description"] == "A table with customer data" + assert existing_entry["dialect"] == "duckdb" + assert existing_entry["start"] == "1 week ago" + assert len(existing_entry["audits"]) == 1 + assert existing_entry["audits"][0]["name"] == "not_null" + assert existing_entry["columns"] == {"id": "BIGINT", "name": "TEXT"} + + other_entry = [e for e in schema if e["name"] == '"schema"."other_table"'][0] + assert other_entry["description"] == "Another table" + assert other_entry["columns"] == {"x": "text"} + + new_entry = [e for e in schema if e["name"] == '"schema"."new_table"'][0] + assert new_entry["columns"] == {"id": "BIGINT", "name": "TEXT"} + + assert any( + "other_table" in str(call) and "no longer referenced" in str(call) + for call in mock_warning.call_args_list + ) + + +def test_create_external_models_sync_appends_new_entries(tmp_path: Path, mocker: MockerFixture): + initial_entries = [ + { + "name": '"schema"."existing_table"', + "columns": {"a": "int"}, + }, + ] + with open(tmp_path / c.EXTERNAL_MODELS_YAML, "w", encoding="utf-8") as fd: + yaml.dump(initial_entries, fd) + + engine_adapter_mock = mocker.Mock() + engine_adapter_mock.columns.return_value = { + "id": exp.DataType.build("bigint"), + } + + state_reader_mock = mocker.Mock() + state_reader_mock.nodes_exist.return_value = set() + + model = SqlModel( + name="b", + query=parse_one("select * FROM schema.new_table"), + ) + + filename = tmp_path / c.EXTERNAL_MODELS_YAML + + with patch.object(get_console(), "log_warning"): + create_external_models_file( + filename, + {"b": model}, # type: ignore + engine_adapter_mock, + state_reader_mock, + "", + mode=CreateExternalModelsMode.SYNC, + ) + + schema = yaml.load(filename) + + assert len(schema) == 2 + names = {e["name"] for e in schema} + assert names == {'"schema"."existing_table"', '"schema"."new_table"'} + + +def test_create_external_models_sync_prune_removes_stale(tmp_path: Path, mocker: MockerFixture): + initial_entries = [ + { + "name": '"schema"."stale_table"', + "description": "This table no longer exists in the project", + "columns": {"a": "int"}, + }, + { + "name": '"schema"."active_table"', + "description": "Still in use", + "columns": {"b": "text"}, + }, + ] + with open(tmp_path / c.EXTERNAL_MODELS_YAML, "w", encoding="utf-8") as fd: + yaml.dump(initial_entries, fd) + + engine_adapter_mock = mocker.Mock() + engine_adapter_mock.columns.return_value = { + "b": exp.DataType.build("varchar"), + "c": exp.DataType.build("bigint"), + } + + state_reader_mock = mocker.Mock() + state_reader_mock.nodes_exist.return_value = set() + + model = SqlModel( + name="x", + query=parse_one("select * FROM schema.active_table"), + ) + + filename = tmp_path / c.EXTERNAL_MODELS_YAML + + with patch.object(get_console(), "log_warning") as mock_warning: + create_external_models_file( + filename, + {"x": model}, # type: ignore + engine_adapter_mock, + state_reader_mock, + "", + mode=CreateExternalModelsMode.SYNC_PRUNE, + ) + + schema = yaml.load(filename) + + assert len(schema) == 1 + assert schema[0]["name"] == '"schema"."active_table"' + assert schema[0]["description"] == "Still in use" + assert schema[0]["columns"] == {"b": "VARCHAR", "c": "BIGINT"} + assert any("Removed 1 stale" in str(call) for call in mock_warning.call_args_list) + + +def test_create_external_models_sync_preserves_other_gateways( + tmp_path: Path, mocker: MockerFixture +): + initial_entries = [ + { + "name": '"schema"."my_table"', + "gateway": "dev", + "description": "Dev-specific config", + "columns": {"a": "int"}, + }, + { + "name": '"schema"."my_table"', + "gateway": "prod", + "description": "Prod-specific config", + "columns": {"a": "int"}, + }, + ] + with open(tmp_path / c.EXTERNAL_MODELS_YAML, "w", encoding="utf-8") as fd: + yaml.dump(initial_entries, fd) + + engine_adapter_mock = mocker.Mock() + engine_adapter_mock.columns.return_value = { + "x": exp.DataType.build("int"), + } + + state_reader_mock = mocker.Mock() + state_reader_mock.nodes_exist.return_value = set() + + model = SqlModel( + name="m", + query=parse_one("select * FROM schema.my_table"), + ) + + filename = tmp_path / c.EXTERNAL_MODELS_YAML + + with patch.object(get_console(), "log_warning"): + create_external_models_file( + filename, + {"m": model}, # type: ignore + engine_adapter_mock, + state_reader_mock, + "", + gateway="dev", + mode=CreateExternalModelsMode.SYNC_PRUNE, + ) + + schema = yaml.load(filename) + + assert len(schema) == 2 + + dev_entry = [e for e in schema if e.get("gateway") == "dev"][0] + prod_entry = [e for e in schema if e.get("gateway") == "prod"][0] + + assert dev_entry["description"] == "Dev-specific config" + assert dev_entry["columns"] == {"x": "INT"} + assert prod_entry["description"] == "Prod-specific config" + assert prod_entry["columns"] == {"a": "int"} + + def test_missing_table(tmp_path: Path): config = Config(gateways=GatewayConfig(connection=DuckDBConnectionConfig())) context = Context(paths=[str(tmp_path.absolute())], config=config)