Fix: Make sure that the deployable table is used for pre-/post- statements during the snapshot table creation (#2282)

Author: izeigerman

sqlmesh/core/snapshot/definition.py

@@ -1140,12 +1140,19 @@ def is_representative(self, snapshot: SnapshotIdLike) -> bool:
 
     def with_non_deployable(self, snapshot: SnapshotIdLike) -> DeployabilityIndex:
         """Creates a new index with the given snapshot marked as non-deployable."""
-        snapshot_id = self._snapshot_id_key(snapshot.snapshot_id)
+        return self._add_snapshot(snapshot, False)
+
+    def with_deployable(self, snapshot: SnapshotIdLike) -> DeployabilityIndex:
+        """Creates a new index with the given snapshot marked as deployable."""
+        return self._add_snapshot(snapshot, True)
+
+    def _add_snapshot(self, snapshot: SnapshotIdLike, deployable: bool) -> DeployabilityIndex:
+        snapshot_id = {self._snapshot_id_key(snapshot.snapshot_id)}
         indexed_ids = self.indexed_ids
         if self.is_opposite_index:
-            indexed_ids = indexed_ids | {snapshot_id}
+            indexed_ids = indexed_ids - snapshot_id if deployable else indexed_ids | snapshot_id
         else:
-            indexed_ids = indexed_ids - {snapshot_id}
+            indexed_ids = indexed_ids | snapshot_id if deployable else indexed_ids - snapshot_id
 
         return DeployabilityIndex(
             indexed_ids=indexed_ids,

sqlmesh/core/snapshot/evaluator.py

@@ -598,20 +598,25 @@ def _create_snapshot(
         deployability_index = deployability_index or DeployabilityIndex.all_deployable()
         is_snapshot_deployable = deployability_index.is_deployable(snapshot)
 
-        # Refers to self as non-deployable to successfully create self-referential tables / views.
-        deployability_index = deployability_index.with_non_deployable(snapshot)
-
-        render_kwargs: t.Dict[str, t.Any] = dict(
+        common_render_kwargs: t.Dict[str, t.Any] = dict(
             engine_adapter=self.adapter,
             snapshots=parent_snapshots_by_name,
-            deployability_index=deployability_index,
             runtime_stage=RuntimeStage.CREATING,
         )
+        pre_post_render_kwargs = dict(
+            **common_render_kwargs,
+            deployability_index=deployability_index.with_deployable(snapshot),
+        )
+        create_render_kwargs = dict(
+            **common_render_kwargs,
+            # Refers to self as non-deployable to successfully create self-referential tables / views.
+            deployability_index=deployability_index.with_non_deployable(snapshot),
+        )
 
         evaluation_strategy = _evaluation_strategy(snapshot, self.adapter)
 
         with self.adapter.transaction(), self.adapter.session(snapshot.model.session_properties):
-            self.adapter.execute(snapshot.model.render_pre_statements(**render_kwargs))
+            self.adapter.execute(snapshot.model.render_pre_statements(**pre_post_render_kwargs))
 
             if (
                 snapshot.is_forward_only
@@ -626,7 +631,7 @@ def _create_snapshot(
                 logger.info(f"Cloning table '{source_table_name}' into '{target_table_name}'")
 
                 evaluation_strategy.create(
-                    snapshot, tmp_table_name, False, is_snapshot_deployable, **render_kwargs
+                    snapshot, tmp_table_name, False, is_snapshot_deployable, **create_render_kwargs
                 )
                 try:
                     self.adapter.clone_table(target_table_name, snapshot.table_name(), replace=True)
@@ -643,10 +648,10 @@ def _create_snapshot(
                         snapshot.table_name(is_deployable=is_table_deployable),
                         is_table_deployable,
                         is_snapshot_deployable,
-                        **render_kwargs,
+                        **create_render_kwargs,
                     )
 
-            self.adapter.execute(snapshot.model.render_post_statements(**render_kwargs))
+            self.adapter.execute(snapshot.model.render_post_statements(**pre_post_render_kwargs))
 
         if on_complete is not None:
             on_complete(snapshot)

tests/core/test_snapshot_evaluator.py

@@ -1583,3 +1583,43 @@ def test_audit_wap(adapter_mock, make_snapshot):
 
     adapter_mock.wap_table_name.assert_called_once_with(snapshot.table_name(), wap_id)
     adapter_mock.wap_publish.assert_called_once_with(snapshot.table_name(), wap_id)
+
+
+def test_create_post_statements_use_deployable_table(
+    mocker: MockerFixture, adapter_mock, make_snapshot
+):
+    evaluator = SnapshotEvaluator(adapter_mock)
+
+    model = load_sql_based_model(
+        parse(  # type: ignore
+            """
+            MODEL (
+                name test_schema.test_model,
+                kind FULL,
+                dialect postgres,
+            );
+
+            CREATE INDEX IF NOT EXISTS test_idx ON test_schema.test_model(a);
+
+            SELECT a::int FROM tbl;
+
+            CREATE INDEX IF NOT EXISTS test_idx ON test_schema.test_model(a);
+            """
+        ),
+    )
+
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    expected_call = f'CREATE INDEX IF NOT EXISTS "test_idx" ON "sqlmesh__test_schema"."test_schema__test_model__{snapshot.version}" /* test_schema.test_model */("a" NULLS FIRST)'
+
+    evaluator.create([snapshot], {}, DeployabilityIndex.none_deployable())
+
+    call_args = adapter_mock.execute.call_args_list
+    pre_calls = call_args[0][0][0]
+    assert len(pre_calls) == 1
+    assert pre_calls[0].sql(dialect="postgres") == expected_call
+
+    post_calls = call_args[1][0][0]
+    assert len(post_calls) == 1
+    assert post_calls[0].sql(dialect="postgres") == expected_call