Repository: TobikoData/sqlmesh
Pull request #3311 (2024-10-30 18:21:28)

feat: add temp_schema option for table_diff
So far, table_diff will create a temp table in `sqlmesh_temp` schema.
However in some environments, the user might want to use an existing schema for that because he either don't want to create and manage another schema or he doesn't have the rights to do so.

The PR introduces a CLI parameter `--temp-schema <catalog>.<schema>` or `--temp-schema <schema>` to write the temp tables inside that catalog and/or schema.

To do list:
- [x] Check that works as expected
- [ ] Ensure it works with API Endpoint
- [x] Add related documentation


Patch:
diff --git a/docs/concepts/models/overview.md b/docs/concepts/models/overview.md
index 03b1efde..4902b809 100644
--- a/docs/concepts/models/overview.md
+++ b/docs/concepts/models/overview.md
@@ -368,7 +368,7 @@ Some properties are only available in specific model kinds - see the [model conf
 :   The maximum number of [batches](#batch_size) that can run concurrently for this model. If not specified, the concurrency is only constrained by the number of concurrent tasks set in the connection settings.
 
 ### lookback
-:   Lookback is used with [incremental by time range](model_kinds.md#incremental_by_time_range) models to capture late-arriving data. It must be a positive integer and specifies the number of interval time units prior to the current interval the model should include.
+:   Lookback is used with [incremental by time range](model_kinds.md#incremental_by_time_range) and [incremental by unique key](model_kinds.md#incremental_by_unique_key) models to capture late-arriving data. It must be a positive integer and specifies the number of interval time units prior to the current interval the model should include.
 
     For example, a model with cron `@daily` and `lookback` of 7 would include the previous 7 days each time it ran, while a model with cron `@weekly` and `lookback` of 7 would include the previous 7 weeks each time it ran.
 
diff --git a/docs/guides/tablediff.md b/docs/guides/tablediff.md
index 5dd6debd..193c253a 100644
--- a/docs/guides/tablediff.md
+++ b/docs/guides/tablediff.md
@@ -118,6 +118,10 @@ Grain should have unique and not-null audits for accurate results.
 
 ```
 
+Under the hood, SQLMesh stores temporary data in the database to perform the comparison.
+The default schema for these temporary tables is `sqlmesh_temp` but can be changed with the `--temp-schema` option.
+The schema can be specified as a `CATALOG.SCHEMA` or `SCHEMA`.
+
 ## Diffing tables or views
 
 Compare specific tables or views with the SQLMesh CLI interface by using the command `sqlmesh table_diff [source table]:[target table]`.
diff --git a/docs/integrations/dbt.md b/docs/integrations/dbt.md
index 2f952e4f..6e926bcf 100644
--- a/docs/integrations/dbt.md
+++ b/docs/integrations/dbt.md
@@ -292,7 +292,6 @@ The dbt jinja methods that are not currently supported are:
 * schemas
 * graph.nodes.values
 * graph.metrics.values
-* version - learn more about why SQLMesh doesn't support model versions at the [Tobiko Data blog](https://tobikodata.com/the-false-promise-of-dbt-contracts.html)
 
 ## Missing something you need?
 
diff --git a/docs/reference/notebook.md b/docs/reference/notebook.md
index 8053de2a..2168fcdd 100644
--- a/docs/reference/notebook.md
+++ b/docs/reference/notebook.md
@@ -273,6 +273,7 @@ Create a schema file containing external model schemas.
 %table_diff [--on [ON ...]] [--skip-columns [SKIP_COLUMNS ...]]
                 [--model MODEL] [--where WHERE] [--limit LIMIT]
                 [--show-sample] [--decimals DECIMALS] [--skip-grain-check]
+                [--temp-schema SCHEMA]
                 SOURCE:TARGET
 
 Show the diff between two tables.
@@ -298,6 +299,7 @@ options:
                         floating point columns. Default: 3
   --skip-grain-check    Disable the check for a primary key (grain) that is
                         missing or is not unique.
+  --temp-schema SCHEMA  The schema to use for temporary tables.
 ```
 
 #### model
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 60dfd31e..116b9415 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -804,6 +804,11 @@ def create_external_models(obj: Context, **kwargs: t.Any) -> None:
     is_flag=True,
     help="Disable the check for a primary key (grain) that is missing or is not unique.",
 )
+@click.option(
+    "--temp-schema",
+    type=str,
+    help="Schema used for temporary tables. It can be `CATALOG.SCHEMA` or `SCHEMA`. Default: `sqlmesh_temp`",
+)
 @click.pass_obj
 @error_handler
 @cli_analytics
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 5f73c9df..f9fb2b4b 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -617,6 +617,8 @@ class GenericContext(BaseContext, t.Generic[C]):
             True if the run was successful, False otherwise.
         """
         environment = environment or self.config.default_target_environment
+        if not skip_janitor and environment.lower() == c.PROD:
+            self._run_janitor()
 
         self.notification_target_manager.notify(
             NotificationEvent.RUN_START, environment=environment
@@ -627,30 +629,76 @@ class GenericContext(BaseContext, t.Generic[C]):
         )
         self._load_materializations_and_signals()
 
+        env_check_attempts_num = max(
+            1,
+            self.config.run.environment_check_max_wait
+            // self.config.run.environment_check_interval,
+        )
+
+        def _block_until_finalized() -> str:
+            for _ in range(env_check_attempts_num):
+                assert environment is not None  # mypy
+                environment_state = self.state_sync.get_environment(environment)
+                if not environment_state:
+                    raise SQLMeshError(f"Environment '{environment}' was not found.")
+                if environment_state.finalized_ts:
+                    return environment_state.plan_id
+                logger.warning(
+                    "Environment '%s' is being updated by plan '%s'. Retrying in %s seconds...",
+                    environment,
+                    environment_state.plan_id,
+                    self.config.run.environment_check_interval,
+                )
+                time.sleep(self.config.run.environment_check_interval)
+            raise SQLMeshError(
+                f"Exceeded the maximum wait time for environment '{environment}' to be ready. "
+                "This means that the environment either failed to update or the update is taking longer than expected. "
+                "See https://sqlmesh.readthedocs.io/en/stable/reference/configuration/#run to adjust the timeout settings."
+            )
+
         success = False
         interrupted = False
-        try:
-            success = self._run(
-                environment=environment,
-                start=start,
-                end=end,
-                execution_time=execution_time,
-                skip_janitor=skip_janitor,
-                ignore_cron=ignore_cron,
-                select_models=select_models,
-                exit_on_env_update=exit_on_env_update is not None,
-            )
-        except CircuitBreakerError:
-            interrupted = True
-        except Exception as e:
-            self.notification_target_manager.notify(
-                NotificationEvent.RUN_FAILURE, traceback.format_exc()
-            )
-            logger.error(f"Run Failure: {traceback.format_exc()}")
-            analytics.collector.on_run_end(
-                run_id=analytics_run_id, succeeded=False, interrupted=False, error=e
-            )
-            raise e
+        done = False
+        while not done:
+            plan_id_at_start = _block_until_finalized()
+
+            def _has_environment_changed() -> bool:
+                assert environment is not None  # mypy
+                current_environment_state = self.state_sync.get_environment(environment)
+                return (
+                    not current_environment_state
+                    or current_environment_state.plan_id != plan_id_at_start
+                    or not current_environment_state.finalized_ts
+                )
+
+            try:
+                success = self._run(
+                    environment,
+                    start=start,
+                    end=end,
+                    execution_time=execution_time,
+                    ignore_cron=ignore_cron,
+                    select_models=select_models,
+                    circuit_breaker=_has_environment_changed,
+                )
+                done = True
+            except CircuitBreakerError:
+                logger.warning(
+                    "Environment '%s' has been modified while running. Restarting the run...",
+                    environment,
+                )
+                if exit_on_env_update:
+                    interrupted = True
+                    done = True
+            except Exception as e:
+                self.notification_target_manager.notify(
+                    NotificationEvent.RUN_FAILURE, traceback.format_exc()
+                )
+                logger.error(f"Run Failure: {traceback.format_exc()}")
+                analytics.collector.on_run_end(
+                    run_id=analytics_run_id, succeeded=False, interrupted=False, error=e
+                )
+                raise e
 
         if success or interrupted:
             self.notification_target_manager.notify(
@@ -1382,6 +1430,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         show_sample: bool = True,
         decimals: int = 3,
         skip_grain_check: bool = False,
+        temp_schema: t.Optional[str] = None,
     ) -> TableDiff:
         """Show a diff between two tables.
 
@@ -1398,6 +1447,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             show_sample: Show the sample dataframe in the console. Requires show=True.
             decimals: The number of decimal places to keep when comparing floating point columns.
             skip_grain_check: Skip check for rows that contain null or duplicate grains.
+            temp_schema: The schema to use for temporary tables.
 
         Returns:
             The TableDiff object containing schema and summary differences.
@@ -1456,7 +1506,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         if show:
             self.console.show_schema_diff(table_diff.schema_diff())
             self.console.show_row_diff(
-                table_diff.row_diff(skip_grain_check=skip_grain_check),
+                table_diff.row_diff(temp_schema=temp_schema, skip_grain_check=skip_grain_check),
                 show_sample=show_sample,
                 skip_grain_check=skip_grain_check,
             )
@@ -1822,41 +1872,10 @@ class GenericContext(BaseContext, t.Generic[C]):
         start: t.Optional[TimeLike],
         end: t.Optional[TimeLike],
         execution_time: t.Optional[TimeLike],
-        skip_janitor: bool,
         ignore_cron: bool,
         select_models: t.Optional[t.Collection[str]],
-        exit_on_env_update: bool,
+        circuit_breaker: t.Optional[t.Callable[[], bool]],
     ) -> bool:
-        if not skip_janitor and environment.lower() == c.PROD:
-            self._run_janitor()
-
-        env_check_attempts_num = max(
-            1,
-            self.config.run.environment_check_max_wait
-            // self.config.run.environment_check_interval,
-        )
-
-        def _block_until_finalized() -> str:
-            for _ in range(env_check_attempts_num):
-                assert environment is not None  # mypy
-                environment_state = self.state_sync.get_environment(environment)
-                if not environment_state:
-                    raise SQLMeshError(f"Environment '{environment}' was not found.")
-                if environment_state.finalized_ts:
-                    return environment_state.plan_id
-                logger.warning(
-                    "Environment '%s' is being updated by plan '%s'. Retrying in %s seconds...",
-                    environment,
-                    environment_state.plan_id,
-                    self.config.run.environment_check_interval,
-                )
-                time.sleep(self.config.run.environment_check_interval)
-            raise SQLMeshError(
-                f"Exceeded the maximum wait time for environment '{environment}' to be ready. "
-                "This means that the environment either failed to update or the update is taking longer than expected. "
-                "See https://sqlmesh.readthedocs.io/en/stable/reference/configuration/#run to adjust the timeout settings."
-            )
-
         scheduler = self.scheduler(environment=environment)
         snapshots = scheduler.snapshots
 
@@ -1870,39 +1889,15 @@ class GenericContext(BaseContext, t.Generic[C]):
             model_selector = self._new_selector(models=models, dag=dag)
             select_models = set(dag.subdag(*model_selector.expand_model_selections(select_models)))
 
-        done = False
-        while not done:
-            plan_id_at_start = _block_until_finalized()
-
-            def _has_environment_changed() -> bool:
-                assert environment is not None  # mypy
-                current_environment_state = self.state_sync.get_environment(environment)
-                return (
-                    not current_environment_state
-                    or current_environment_state.plan_id != plan_id_at_start
-                    or not current_environment_state.finalized_ts
-                )
-
-            try:
-                success = scheduler.run(
-                    environment,
-                    start=start,
-                    end=end,
-                    execution_time=execution_time,
-                    ignore_cron=ignore_cron,
-                    circuit_breaker=_has_environment_changed,
-                    selected_snapshots=select_models,
-                )
-                done = True
-            except CircuitBreakerError:
-                logger.warning(
-                    "Environment '%s' has been modified while running. Restarting the run...",
-                    environment,
-                )
-                if exit_on_env_update:
-                    raise
-
-        return success
+        return scheduler.run(
+            environment,
+            start=start,
+            end=end,
+            execution_time=execution_time,
+            ignore_cron=ignore_cron,
+            circuit_breaker=circuit_breaker,
+            selected_snapshots=select_models,
+        )
 
     def _apply(self, plan: Plan, circuit_breaker: t.Optional[t.Callable[[], bool]]) -> None:
         self._scheduler.create_plan_evaluator(self).evaluate(
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index db90e1ba..9caa270d 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -1133,10 +1133,12 @@ def transform_values(
         yield _transform_value(col_value, col_type)
 
 
-def to_schema(sql_path: str | exp.Table) -> exp.Table:
+def to_schema(sql_path: str | exp.Table, dialect: DialectType = None) -> exp.Table:
     if isinstance(sql_path, exp.Table) and sql_path.this is None:
         return sql_path
-    table = exp.to_table(sql_path.copy() if isinstance(sql_path, exp.Table) else sql_path)
+    table = exp.to_table(
+        sql_path.copy() if isinstance(sql_path, exp.Table) else sql_path, dialect=dialect
+    )
     table.set("catalog", table.args.get("db"))
     table.set("db", table.args.get("this"))
     table.set("this", None)
diff --git a/sqlmesh/core/engine_adapter/bigquery.py b/sqlmesh/core/engine_adapter/bigquery.py
index ba1512c3..4c4e438b 100644
--- a/sqlmesh/core/engine_adapter/bigquery.py
+++ b/sqlmesh/core/engine_adapter/bigquery.py
@@ -223,17 +223,35 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
                 return "JSON"
             return kind.name
 
-        table = self._get_table(table_name)
-        columns = {
-            field.name: exp.DataType.build(
-                dtype_to_sql(field.to_standard_sql().type), dialect=self.dialect
-            )
-            for field in table.schema
-        }
-        if include_pseudo_columns and table.time_partitioning and not table.time_partitioning.field:
-            columns["_PARTITIONTIME"] = exp.DataType.build("TIMESTAMP")
-            if table.time_partitioning.type_ == "DAY":
-                columns["_PARTITIONDATE"] = exp.DataType.build("DATE")
+        def create_mapping_schema(
+            schema: t.Sequence[bigquery.SchemaField],
+        ) -> t.Dict[str, exp.DataType]:
+            return {
+                field.name: exp.DataType.build(
+                    dtype_to_sql(field.to_standard_sql().type), dialect=self.dialect
+                )
+                for field in schema
+            }
+
+        table = exp.to_table(table_name)
+        if len(table.parts) > 3:
+            # The client's `get_table` method can't handle paths with >3 identifiers
+            self.execute(exp.select("*").from_(table).limit(1))
+            query_results = self._query_job._query_results
+            columns = create_mapping_schema(query_results.schema)
+        else:
+            bq_table = self._get_table(table)
+            columns = create_mapping_schema(bq_table.schema)
+
+            if (
+                include_pseudo_columns
+                and bq_table.time_partitioning
+                and not bq_table.time_partitioning.field
+            ):
+                columns["_PARTITIONTIME"] = exp.DataType.build("TIMESTAMP")
+                if bq_table.time_partitioning.type_ == "DAY":
+                    columns["_PARTITIONDATE"] = exp.DataType.build("DATE")
+
         return columns
 
     def alter_table(
diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py
index 36f2b675..e8020c72 100644
--- a/sqlmesh/core/table_diff.py
+++ b/sqlmesh/core/table_diff.py
@@ -5,6 +5,8 @@ import typing as t
 from functools import cached_property
 
 import pandas as pd
+
+from sqlmesh.core.dialect import to_schema
 from sqlmesh.core.engine_adapter.mixins import RowDiffMixin
 from sqlglot import exp, parse_one
 from sqlglot.helper import ensure_list
@@ -253,7 +255,9 @@ class TableDiff:
             model_name=self.model_name,
         )
 
-    def row_diff(self, skip_grain_check: bool = False) -> RowDiff:
+    def row_diff(
+        self, temp_schema: t.Optional[str] = None, skip_grain_check: bool = False
+    ) -> RowDiff:
         if self._row_diff is None:
             source_schema = {
                 c: t for c, t in self.source_schema.items() if c not in self.skip_columns
@@ -399,7 +403,12 @@ class TableDiff:
             query = self.adapter.ensure_nulls_for_unmatched_after_join(
                 quote_identifiers(base_query.copy(), dialect=self.model_dialect or self.dialect)
             )
-            temp_table = exp.table_("diff", db="sqlmesh_temp", quoted=True)
+
+            if not temp_schema:
+                temp_schema = "sqlmesh_temp"
+
+            schema = to_schema(temp_schema, dialect=self.dialect)
+            temp_table = exp.table_("diff", db=schema.db, catalog=schema.catalog, quoted=True)
 
             with self.adapter.temp_table(query, name=temp_table) as table:
                 summary_sums = [
diff --git a/sqlmesh/dbt/target.py b/sqlmesh/dbt/target.py
index 79c27c12..6c523703 100644
--- a/sqlmesh/dbt/target.py
+++ b/sqlmesh/dbt/target.py
@@ -526,6 +526,7 @@ class BigQueryConfig(TargetConfig):
     dataset: t.Optional[str] = None
     project: t.Optional[str] = None
     execution_project: t.Optional[str] = None
+    quota_project: t.Optional[str] = None
     location: t.Optional[str] = None
     keyfile: t.Optional[str] = None
     keyfile_json: t.Optional[t.Dict[str, t.Any]] = None
@@ -586,6 +587,7 @@ class BigQueryConfig(TargetConfig):
             method=self.method,
             project=self.database,
             execution_project=self.execution_project,
+            quota_project=self.quota_project,
             location=self.location,
             concurrent_tasks=self.threads,
             keyfile=self.keyfile,
diff --git a/web/server/api/endpoints/table_diff.py b/web/server/api/endpoints/table_diff.py
index d60bc676..7d0da2cc 100644
--- a/web/server/api/endpoints/table_diff.py
+++ b/web/server/api/endpoints/table_diff.py
@@ -20,6 +20,7 @@ def get_table_diff(
     on: t.Optional[str] = None,
     model_or_snapshot: t.Optional[str] = None,
     where: t.Optional[str] = None,
+    temp_schema: t.Optional[str] = None,
     limit: int = 20,
     context: Context = Depends(get_loaded_context),
 ) -> TableDiff:
@@ -34,7 +35,7 @@ def get_table_diff(
         show=False,
     )
     _schema_diff = diff.schema_diff()
-    _row_diff = diff.row_diff()
+    _row_diff = diff.row_diff(temp_schema=temp_schema)
     schema_diff = SchemaDiff(
         source=_schema_diff.source,
         target=_schema_diff.target,


Test patch:
diff --git a/tests/core/engine_adapter/integration/__init__.py b/tests/core/engine_adapter/integration/__init__.py
index 7773ccf9..1c47e628 100644
--- a/tests/core/engine_adapter/integration/__init__.py
+++ b/tests/core/engine_adapter/integration/__init__.py
@@ -23,6 +23,7 @@ from sqlmesh.utils.pydantic import PydanticModel
 from tests.utils.pandas import compare_dataframes
 
 if t.TYPE_CHECKING:
+    from sqlmesh.core._typing import TableName
     from sqlmesh.core.engine_adapter._typing import Query
 
 TEST_SCHEMA = "test_schema"
@@ -212,12 +213,16 @@ class TestContext:
     def output_data(self, data: pd.DataFrame) -> pd.DataFrame:
         return self._format_df(data)
 
-    def table(self, table_name: str, schema: str = TEST_SCHEMA) -> exp.Table:
+    def table(self, table_name: TableName, schema: str = TEST_SCHEMA) -> exp.Table:
         schema = self.add_test_suffix(schema)
         self._schemas.append(schema)
+
+        table = exp.to_table(table_name)
+        table.set("db", exp.parse_identifier(schema, dialect=self.dialect))
+
         return exp.to_table(
             normalize_model_name(
-                ".".join([schema, table_name]),
+                table,
                 default_catalog=self.engine_adapter.default_catalog,
                 dialect=self.dialect,
             )
diff --git a/tests/core/engine_adapter/integration/test_integration_bigquery.py b/tests/core/engine_adapter/integration/test_integration_bigquery.py
index 55ffd3a7..adebb6af 100644
--- a/tests/core/engine_adapter/integration/test_integration_bigquery.py
+++ b/tests/core/engine_adapter/integration/test_integration_bigquery.py
@@ -178,3 +178,43 @@ def test_mutating_clustered_by_forward_only(
 
     metadata = _get_data_object(target_table_1)
     assert not metadata.is_clustered
+
+
+def test_fetch_schema_of_information_schema_tables(
+    ctx: TestContext, engine_adapter: BigQueryEngineAdapter
+):
+    # We produce Table(this=Dot(this=INFORMATION_SCHEMA, expression=TABLES)) here,
+    # otherwise `db` or `catalog` would be set, which is not the right representation
+    information_schema_tables = exp.to_table("_._.INFORMATION_SCHEMA.TABLES")
+    information_schema_tables.set("db", None)
+    information_schema_tables.set("catalog", None)
+
+    source = ctx.table(information_schema_tables)
+
+    expected_columns_to_types = {
+        "table_catalog": exp.DataType.build("TEXT"),
+        "table_schema": exp.DataType.build("TEXT"),
+        "table_name": exp.DataType.build("TEXT"),
+        "table_type": exp.DataType.build("TEXT"),
+        "is_insertable_into": exp.DataType.build("TEXT"),
+        "is_typed": exp.DataType.build("TEXT"),
+        "creation_time": exp.DataType.build("TIMESTAMPTZ"),
+        "base_table_catalog": exp.DataType.build("TEXT"),
+        "base_table_schema": exp.DataType.build("TEXT"),
+        "base_table_name": exp.DataType.build("TEXT"),
+        "snapshot_time_ms": exp.DataType.build("TIMESTAMPTZ"),
+        "ddl": exp.DataType.build("TEXT"),
+        "default_collation_name": exp.DataType.build("TEXT"),
+        "upsert_stream_apply_watermark": exp.DataType.build("TIMESTAMPTZ"),
+        "replica_source_catalog": exp.DataType.build("TEXT"),
+        "replica_source_schema": exp.DataType.build("TEXT"),
+        "replica_source_name": exp.DataType.build("TEXT"),
+        "replication_status": exp.DataType.build("TEXT"),
+        "replication_error": exp.DataType.build("TEXT"),
+        "is_change_history_enabled": exp.DataType.build("TEXT"),
+        "sync_status": exp.DataType.build(
+            "STRUCT<last_completion_time TIMESTAMPTZ, error_time TIMESTAMPTZ, error STRUCT<reason TEXT, location TEXT, message TEXT>>"
+        ),
+    }
+
+    assert expected_columns_to_types == engine_adapter.columns(source.sql())
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index d7f34d43..b8caf69b 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -5945,12 +5945,12 @@ MODEL (
   dialect snowflake,
 );
 
-SELECT * FROM (@custom_macro(@foo, @bar)) AS q
+SELECT * FROM (@custom_macro(@a, @b)) AS q
     """)
 
     config = Config(
         model_defaults=ModelDefaultsConfig(dialect="duckdb"),
-        variables={"foo": "foo", "bar": "boo"},
+        variables={"a": "a", "b": "b"},
     )
     context = Context(paths=tmp_path, config=config)
 
diff --git a/tests/core/test_table_diff.py b/tests/core/test_table_diff.py
index 811ff762..c59df5b9 100644
--- a/tests/core/test_table_diff.py
+++ b/tests/core/test_table_diff.py
@@ -266,11 +266,11 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
         ),
     )
 
-    query_sql = 'CREATE TABLE IF NOT EXISTS "memory"."sqlmesh_temp"."__temp_diff_abcdefgh" AS WITH "__source" AS (SELECT "key", "value", "key" AS "__sqlmesh_join_key" FROM "table_diff_source"), "__target" AS (SELECT "key", "value", "key" AS "__sqlmesh_join_key" FROM "table_diff_target"), "__stats" AS (SELECT "s"."key" AS "s__key", "s"."value" AS "s__value", "s"."__sqlmesh_join_key" AS "s____sqlmesh_join_key", "t"."key" AS "t__key", "t"."value" AS "t__value", "t"."__sqlmesh_join_key" AS "t____sqlmesh_join_key", CASE WHEN NOT "s"."key" IS NULL THEN 1 ELSE 0 END AS "s_exists", CASE WHEN NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "t_exists", CASE WHEN "s"."__sqlmesh_join_key" = "t"."__sqlmesh_join_key" AND (NOT "s"."key" IS NULL AND NOT "t"."key" IS NULL) THEN 1 ELSE 0 END AS "row_joined", CASE WHEN "s"."key" IS NULL AND "t"."key" IS NULL THEN 1 ELSE 0 END AS "null_grain", CASE WHEN "s"."key" = "t"."key" THEN 1 WHEN ("s"."key" IS NULL) AND ("t"."key" IS NULL) THEN 1 WHEN ("s"."key" IS NULL) OR ("t"."key" IS NULL) THEN 0 ELSE 0 END AS "key_matches", CASE WHEN ROUND("s"."value", 3) = ROUND("t"."value", 3) THEN 1 WHEN ("s"."value" IS NULL) AND ("t"."value" IS NULL) THEN 1 WHEN ("s"."value" IS NULL) OR ("t"."value" IS NULL) THEN 0 ELSE 0 END AS "value_matches" FROM "__source" AS "s" FULL JOIN "__target" AS "t" ON "s"."__sqlmesh_join_key" = "t"."__sqlmesh_join_key") SELECT *, CASE WHEN "key_matches" = 1 AND "value_matches" = 1 THEN 1 ELSE 0 END AS "row_full_match" FROM "__stats"'
-    summary_query_sql = 'SELECT SUM("s_exists") AS "s_count", SUM("t_exists") AS "t_count", SUM("row_joined") AS "join_count", SUM("null_grain") AS "null_grain_count", SUM("row_full_match") AS "full_match_count", SUM("key_matches") AS "key_matches", SUM("value_matches") AS "value_matches", COUNT(DISTINCT ("s____sqlmesh_join_key")) AS "distinct_count_s", COUNT(DISTINCT ("t____sqlmesh_join_key")) AS "distinct_count_t" FROM "memory"."sqlmesh_temp"."__temp_diff_abcdefgh"'
-    compare_sql = 'SELECT ROUND(100 * (CAST(SUM("key_matches") AS DECIMAL) / COUNT("key_matches")), 1) AS "key_matches", ROUND(100 * (CAST(SUM("value_matches") AS DECIMAL) / COUNT("value_matches")), 1) AS "value_matches" FROM "memory"."sqlmesh_temp"."__temp_diff_abcdefgh" WHERE "row_joined" = 1'
-    sample_query_sql = 'SELECT "s_exists", "t_exists", "row_joined", "row_full_match", "s__key", "s__value", "s____sqlmesh_join_key", "t__key", "t__value", "t____sqlmesh_join_key" FROM "memory"."sqlmesh_temp"."__temp_diff_abcdefgh" WHERE "key_matches" = 0 OR "value_matches" = 0 ORDER BY "s__key" NULLS FIRST, "t__key" NULLS FIRST LIMIT 20'
-    drop_sql = 'DROP TABLE IF EXISTS "memory"."sqlmesh_temp"."__temp_diff_abcdefgh"'
+    query_sql = 'CREATE TABLE IF NOT EXISTS "memory"."sqlmesh_temp_test"."__temp_diff_abcdefgh" AS WITH "__source" AS (SELECT "key", "value", "key" AS "__sqlmesh_join_key" FROM "table_diff_source"), "__target" AS (SELECT "key", "value", "key" AS "__sqlmesh_join_key" FROM "table_diff_target"), "__stats" AS (SELECT "s"."key" AS "s__key", "s"."value" AS "s__value", "s"."__sqlmesh_join_key" AS "s____sqlmesh_join_key", "t"."key" AS "t__key", "t"."value" AS "t__value", "t"."__sqlmesh_join_key" AS "t____sqlmesh_join_key", CASE WHEN NOT "s"."key" IS NULL THEN 1 ELSE 0 END AS "s_exists", CASE WHEN NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "t_exists", CASE WHEN "s"."__sqlmesh_join_key" = "t"."__sqlmesh_join_key" AND (NOT "s"."key" IS NULL AND NOT "t"."key" IS NULL) THEN 1 ELSE 0 END AS "row_joined", CASE WHEN "s"."key" IS NULL AND "t"."key" IS NULL THEN 1 ELSE 0 END AS "null_grain", CASE WHEN "s"."key" = "t"."key" THEN 1 WHEN ("s"."key" IS NULL) AND ("t"."key" IS NULL) THEN 1 WHEN ("s"."key" IS NULL) OR ("t"."key" IS NULL) THEN 0 ELSE 0 END AS "key_matches", CASE WHEN ROUND("s"."value", 3) = ROUND("t"."value", 3) THEN 1 WHEN ("s"."value" IS NULL) AND ("t"."value" IS NULL) THEN 1 WHEN ("s"."value" IS NULL) OR ("t"."value" IS NULL) THEN 0 ELSE 0 END AS "value_matches" FROM "__source" AS "s" FULL JOIN "__target" AS "t" ON "s"."__sqlmesh_join_key" = "t"."__sqlmesh_join_key") SELECT *, CASE WHEN "key_matches" = 1 AND "value_matches" = 1 THEN 1 ELSE 0 END AS "row_full_match" FROM "__stats"'
+    summary_query_sql = 'SELECT SUM("s_exists") AS "s_count", SUM("t_exists") AS "t_count", SUM("row_joined") AS "join_count", SUM("null_grain") AS "null_grain_count", SUM("row_full_match") AS "full_match_count", SUM("key_matches") AS "key_matches", SUM("value_matches") AS "value_matches", COUNT(DISTINCT ("s____sqlmesh_join_key")) AS "distinct_count_s", COUNT(DISTINCT ("t____sqlmesh_join_key")) AS "distinct_count_t" FROM "memory"."sqlmesh_temp_test"."__temp_diff_abcdefgh"'
+    compare_sql = 'SELECT ROUND(100 * (CAST(SUM("key_matches") AS DECIMAL) / COUNT("key_matches")), 1) AS "key_matches", ROUND(100 * (CAST(SUM("value_matches") AS DECIMAL) / COUNT("value_matches")), 1) AS "value_matches" FROM "memory"."sqlmesh_temp_test"."__temp_diff_abcdefgh" WHERE "row_joined" = 1'
+    sample_query_sql = 'SELECT "s_exists", "t_exists", "row_joined", "row_full_match", "s__key", "s__value", "s____sqlmesh_join_key", "t__key", "t__value", "t____sqlmesh_join_key" FROM "memory"."sqlmesh_temp_test"."__temp_diff_abcdefgh" WHERE "key_matches" = 0 OR "value_matches" = 0 ORDER BY "s__key" NULLS FIRST, "t__key" NULLS FIRST LIMIT 20'
+    drop_sql = 'DROP TABLE IF EXISTS "memory"."sqlmesh_temp_test"."__temp_diff_abcdefgh"'
 
     spy_execute = mocker.spy(engine_adapter, "_execute")
     mocker.patch("sqlmesh.core.engine_adapter.base.random_id", return_value="abcdefgh")
@@ -280,6 +280,7 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
         target="table_diff_target",
         on=["key"],
         skip_columns=["ignored"],
+        temp_schema="sqlmesh_temp_test",
     )
 
     spy_execute.assert_any_call(query_sql)