Repository: TobikoData/sqlmesh
Pull request #3304 (2024-10-28 18:31:18)

Fix: Skip creation call for tables that already exist
This update aims to eliminate unnecessary calls to create schemas and tables when the required deployable or non-deployable versions already exist, by identifying if the tables that need to be created are missing, instead of checking for both versions.

Themiscodes: Thank you for the review @izeigerman addressed comments

Patch:
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
index 3b82e8f1..4a4d7edc 100644
--- a/sqlmesh/core/snapshot/evaluator.py
+++ b/sqlmesh/core/snapshot/evaluator.py
@@ -259,14 +259,23 @@ class SnapshotEvaluator:
         """
         snapshots_with_table_names = defaultdict(set)
         tables_by_schema = defaultdict(set)
+        table_deployability: t.Dict[str, bool] = {}
         for snapshot in target_snapshots:
             if not snapshot.is_model or snapshot.is_symbolic:
                 continue
-            for is_deployable in (True, False):
+            deployability_flags = [True]
+            if (
+                snapshot.reuses_previous_version
+                or snapshot.is_managed
+                or (deployability_index and not deployability_index.is_deployable(snapshot))
+            ):
+                deployability_flags.append(False)
+            for is_deployable in deployability_flags:
                 table = exp.to_table(
                     snapshot.table_name(is_deployable), dialect=snapshot.model.dialect
                 )
                 snapshots_with_table_names[snapshot].add(table.name)
+                table_deployability[table.name] = is_deployable
                 tables_by_schema[d.schema_(table.db, catalog=table.catalog)].add(table.name)
 
         def _get_data_objects(schema: exp.Table) -> t.Set[str]:
@@ -284,9 +293,16 @@ class SnapshotEvaluator:
             }
 
         snapshots_to_create = []
+        target_deployability_flags: t.Dict[str, t.List[bool]] = defaultdict(list)
         for snapshot, table_names in snapshots_with_table_names.items():
-            if table_names - existing_objects or (snapshot.is_seed and not snapshot.intervals):
+            missing_tables = table_names - existing_objects
+            if missing_tables or (snapshot.is_seed and not snapshot.intervals):
                 snapshots_to_create.append(snapshot)
+                for table_name in missing_tables or table_names:
+                    target_deployability_flags[snapshot.name].append(
+                        table_deployability[table_name]
+                    )
+                target_deployability_flags[snapshot.name].sort()
             elif on_complete:
                 on_complete(snapshot)
 
@@ -298,7 +314,12 @@ class SnapshotEvaluator:
             concurrent_apply_to_snapshots(
                 snapshots_to_create,
                 lambda s: self._create_snapshot(
-                    s, snapshots, deployability_index, on_complete, allow_destructive_snapshots
+                    s,
+                    snapshots,
+                    target_deployability_flags[s.name],
+                    deployability_index,
+                    on_complete,
+                    allow_destructive_snapshots,
                 ),
                 self.ddl_concurrent_tasks,
             )
@@ -625,6 +646,7 @@ class SnapshotEvaluator:
         self,
         snapshot: Snapshot,
         snapshots: t.Dict[SnapshotId, Snapshot],
+        deployability_flags: t.List[bool],
         deployability_index: t.Optional[DeployabilityIndex],
         on_complete: t.Optional[t.Callable[[SnapshotInfoLike], None]],
         allow_destructive_snapshots: t.Set[str],
@@ -697,22 +719,8 @@ class SnapshotEvaluator:
                 finally:
                     self.adapter.drop_table(tmp_table_name)
             else:
-                dry_run = True
-                table_deployability_flags = (
-                    []
-                    if (
-                        is_snapshot_deployable
-                        and not snapshot.reuses_previous_version
-                        and not snapshot.is_managed
-                    )
-                    else [False]
-                )
-                if not snapshot.reuses_previous_version:
-                    table_deployability_flags.append(True)
-                    if len(table_deployability_flags) > 1:
-                        dry_run = False
-
-                for is_table_deployable in table_deployability_flags:
+                dry_run = len(deployability_flags) == 1
+                for is_table_deployable in deployability_flags:
                     evaluation_strategy.create(
                         table_name=snapshot.table_name(is_deployable=is_table_deployable),
                         model=snapshot.model,


Test patch:
diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
index 0355cad8..a41198ba 100644
--- a/tests/core/test_snapshot_evaluator.py
+++ b/tests/core/test_snapshot_evaluator.py
@@ -661,7 +661,7 @@ def test_evaluate_incremental_unmanaged_no_intervals(
     )
 
 
-def test_create_tables_exists(mocker: MockerFixture, adapter_mock, make_snapshot):
+def test_create_prod_table_exists(mocker: MockerFixture, adapter_mock, make_snapshot):
     model = load_sql_based_model(
         parse(  # type: ignore
             """
@@ -679,11 +679,6 @@ def test_create_tables_exists(mocker: MockerFixture, adapter_mock, make_snapshot
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
     adapter_mock.get_data_objects.return_value = [
-        DataObject(
-            name=f"test_schema__test_model__{snapshot.version}__temp",
-            schema="sqlmesh__test_schema",
-            type=DataObjectType.VIEW,
-        ),
         DataObject(
             name=f"test_schema__test_model__{snapshot.version}",
             schema="sqlmesh__test_schema",
@@ -694,10 +689,10 @@ def test_create_tables_exists(mocker: MockerFixture, adapter_mock, make_snapshot
 
     evaluator.create([snapshot], {})
     adapter_mock.create_view.assert_not_called()
+    adapter_mock.create_schema.assert_not_called()
     adapter_mock.get_data_objects.assert_called_once_with(
         schema_("sqlmesh__test_schema"),
         {
-            f"test_schema__test_model__{snapshot.version}__temp",
             f"test_schema__test_model__{snapshot.version}",
         },
     )
@@ -730,9 +725,129 @@ def test_create_only_dev_table_exists(mocker: MockerFixture, adapter_mock, make_
     adapter_mock.table_exists.return_value = True
     evaluator = SnapshotEvaluator(adapter_mock)
 
+    evaluator.create([snapshot], {})
+    adapter_mock.create_schema.assert_called_once_with(to_schema("sqlmesh__test_schema"))
+    adapter_mock.create_view.assert_not_called()
+    adapter_mock.get_data_objects.assert_called_once_with(
+        schema_("sqlmesh__test_schema"),
+        {
+            f"test_schema__test_model__{snapshot.version}",
+        },
+    )
+
+
+@pytest.mark.parametrize(
+    "deployability_index,  snapshot_category, deployability_flags",
+    [
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.BREAKING, [False]),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.NON_BREAKING, [False]),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.FORWARD_ONLY, [True]),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.INDIRECT_BREAKING, [False]),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.INDIRECT_NON_BREAKING, [True]),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.METADATA, [True]),
+        (
+            DeployabilityIndex.none_deployable(),
+            SnapshotChangeCategory.BREAKING,
+            [True, False],
+        ),
+        (
+            DeployabilityIndex.none_deployable(),
+            SnapshotChangeCategory.NON_BREAKING,
+            [True, False],
+        ),
+        (
+            DeployabilityIndex.none_deployable(),
+            SnapshotChangeCategory.FORWARD_ONLY,
+            [True],
+        ),
+        (
+            DeployabilityIndex.none_deployable(),
+            SnapshotChangeCategory.INDIRECT_BREAKING,
+            [True, False],
+        ),
+        (
+            DeployabilityIndex.none_deployable(),
+            SnapshotChangeCategory.INDIRECT_NON_BREAKING,
+            [True],
+        ),
+        (
+            DeployabilityIndex.none_deployable(),
+            SnapshotChangeCategory.METADATA,
+            [True],
+        ),
+    ],
+)
+def test_create_tables_exist(
+    snapshot: Snapshot,
+    mocker: MockerFixture,
+    adapter_mock,
+    deployability_index: DeployabilityIndex,
+    deployability_flags: t.List[bool],
+    snapshot_category: SnapshotChangeCategory,
+):
+    adapter_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter")
+    adapter_mock.dialect = "duckdb"
+
+    evaluator = SnapshotEvaluator(adapter_mock)
+    snapshot.categorize_as(category=snapshot_category)
+
+    adapter_mock.get_data_objects.return_value = [
+        DataObject(
+            name=f"db__model__{snapshot.version}__temp",
+            schema="sqlmesh__db",
+            type=DataObjectType.TABLE,
+        ),
+        DataObject(
+            name=f"db__model__{snapshot.version}",
+            schema="sqlmesh__db",
+            type=DataObjectType.TABLE,
+        ),
+    ]
+
+    evaluator.create(
+        target_snapshots=[snapshot],
+        snapshots={},
+        deployability_index=deployability_index,
+    )
+
+    adapter_mock.get_data_objects.assert_called_once_with(
+        schema_("sqlmesh__db"),
+        {
+            f"db__model__{snapshot.version}" if not flag else f"db__model__{snapshot.version}__temp"
+            for flag in set(deployability_flags + [False])
+        },
+    )
+    adapter_mock.create_schema.assert_not_called()
+    adapter_mock.create_table.assert_not_called()
+
+
+def test_create_prod_table_exists_forward_only(mocker: MockerFixture, adapter_mock, make_snapshot):
+    model = load_sql_based_model(
+        parse(  # type: ignore
+            """
+            MODEL (
+                name test_schema.test_model,
+                kind FULL
+            );
+
+            SELECT a::int FROM tbl;
+            """
+        ),
+    )
+
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.FORWARD_ONLY)
+
+    adapter_mock.get_data_objects.return_value = [
+        DataObject(
+            name=f"test_schema__test_model__{snapshot.version}",
+            schema="sqlmesh__test_schema",
+            type=DataObjectType.TABLE,
+        ),
+    ]
+    evaluator = SnapshotEvaluator(adapter_mock)
     evaluator.create([snapshot], {})
 
-    adapter_mock.create_view.assert_not_called
     adapter_mock.get_data_objects.assert_called_once_with(
         schema_("sqlmesh__test_schema"),
         {
@@ -741,6 +856,20 @@ def test_create_only_dev_table_exists(mocker: MockerFixture, adapter_mock, make_
         },
     )
 
+    adapter_mock.create_schema.assert_called_once_with(to_schema("sqlmesh__test_schema"))
+    adapter_mock.create_table.assert_called_once_with(
+        f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__temp",
+        columns_to_types={"a": exp.DataType.build("int")},
+        table_format=None,
+        storage_format=None,
+        partitioned_by=[],
+        partition_interval_unit=IntervalUnit.DAY,
+        clustered_by=[],
+        table_properties={},
+        table_description=None,
+        column_descriptions=None,
+    )
+
 
 def test_create_view_non_deployable_snapshot(mocker: MockerFixture, adapter_mock, make_snapshot):
     model = load_sql_based_model(
@@ -897,6 +1026,52 @@ def test_promote_model_info(mocker: MockerFixture, make_snapshot):
     )
 
 
+def test_promote_deployable(mocker: MockerFixture, make_snapshot):
+    adapter_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter")
+    adapter_mock.dialect = "duckdb"
+
+    evaluator = SnapshotEvaluator(adapter_mock)
+
+    model = SqlModel(
+        name="test_schema.test_model",
+        kind=FullKind(),
+        query=parse_one("SELECT a FROM tbl"),
+    )
+
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    adapter_mock.get_data_objects.return_value = [
+        DataObject(
+            name=f"test_schema__test_model__{snapshot.version}",
+            schema="sqlmesh__test_schema",
+            type=DataObjectType.TABLE,
+        ),
+    ]
+
+    evaluator.create([snapshot], {})
+    adapter_mock.get_data_objects.assert_called_once_with(
+        schema_("sqlmesh__test_schema"),
+        {
+            f"test_schema__test_model__{snapshot.version}",
+        },
+    )
+    adapter_mock.create_table.assert_not_called()
+
+    evaluator.promote([snapshot], EnvironmentNamingInfo(name="test_env"))
+
+    adapter_mock.create_schema.assert_called_once_with(to_schema("test_schema__test_env"))
+    adapter_mock.create_view.assert_called_once_with(
+        "test_schema__test_env.test_model",
+        parse_one(
+            f"SELECT * FROM sqlmesh__test_schema.test_schema__test_model__{snapshot.version}"
+        ),
+        table_description=None,
+        column_descriptions=None,
+        view_properties={},
+    )
+
+
 def test_migrate(mocker: MockerFixture, make_snapshot):
     connection_mock = mocker.NonCallableMock()
     cursor_mock = mocker.Mock()
@@ -2635,6 +2810,14 @@ def test_create_managed_forward_only_with_previous_version_doesnt_clone_for_dev_
         ),
     )
 
+    adapter_mock.get_data_objects.return_value = [
+        DataObject(
+            name=f"test_schema__test_model__{snapshot.version}",
+            schema="sqlmesh__test_schema",
+            type=DataObjectType.MANAGED_TABLE,
+        ),
+    ]
+
     evaluator.create(target_snapshots=[snapshot], snapshots={})
 
     # We dont clone managed tables to create dev previews, we use normal tables
@@ -2648,43 +2831,47 @@ def test_create_managed_forward_only_with_previous_version_doesnt_clone_for_dev_
 
 
 @pytest.mark.parametrize(
-    "deployability_index,  snapshot_category, temp_table",
+    "deployability_index,  snapshot_category, deployability_flags",
     [
-        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.BREAKING, [False]),
-        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.NON_BREAKING, [False]),
-        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.FORWARD_ONLY, [True]),
-        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.INDIRECT_BREAKING, [False]),
-        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.INDIRECT_NON_BREAKING, [True]),
-        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.METADATA, [True]),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.BREAKING, [True]),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.NON_BREAKING, [True]),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.FORWARD_ONLY, [False]),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.INDIRECT_BREAKING, [True]),
+        (
+            DeployabilityIndex.all_deployable(),
+            SnapshotChangeCategory.INDIRECT_NON_BREAKING,
+            [False],
+        ),
+        (DeployabilityIndex.all_deployable(), SnapshotChangeCategory.METADATA, [False]),
         (
             DeployabilityIndex.none_deployable(),
             SnapshotChangeCategory.BREAKING,
-            [True, False],
+            [False, True],
         ),
         (
             DeployabilityIndex.none_deployable(),
             SnapshotChangeCategory.NON_BREAKING,
-            [True, False],
+            [False, True],
         ),
         (
             DeployabilityIndex.none_deployable(),
             SnapshotChangeCategory.FORWARD_ONLY,
-            [True],
+            [False],
         ),
         (
             DeployabilityIndex.none_deployable(),
             SnapshotChangeCategory.INDIRECT_BREAKING,
-            [True, False],
+            [False, True],
         ),
         (
             DeployabilityIndex.none_deployable(),
             SnapshotChangeCategory.INDIRECT_NON_BREAKING,
-            [True],
+            [False],
         ),
         (
             DeployabilityIndex.none_deployable(),
             SnapshotChangeCategory.METADATA,
-            [True],
+            [False],
         ),
     ],
 )
@@ -2693,7 +2880,7 @@ def test_create_snapshot(
     mocker: MockerFixture,
     adapter_mock,
     deployability_index: DeployabilityIndex,
-    temp_table: t.List[bool],
+    deployability_flags: t.List[bool],
     snapshot_category: SnapshotChangeCategory,
 ):
     adapter_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter")
@@ -2704,6 +2891,7 @@ def test_create_snapshot(
     evaluator._create_snapshot(
         snapshot=snapshot,
         snapshots={},
+        deployability_flags=deployability_flags,
         deployability_index=deployability_index,
         on_complete=None,
         allow_destructive_snapshots=set(),
@@ -2722,11 +2910,11 @@ def test_create_snapshot(
 
     tables_created = [
         call(
-            snapshot.table_name(is_deployable=not temp),
-            column_descriptions=(None if temp else {}),
+            snapshot.table_name(is_deployable=is_deployable),
+            column_descriptions=(None if not is_deployable else {}),
             **common_kwargs,
         )
-        for temp in temp_table
+        for is_deployable in deployability_flags
     ]
 
     adapter_mock.create_table.assert_has_calls(tables_created)