Repository: TobikoData/sqlmesh
Pull request #4187 (2025-04-19 03:10:20)

feat: add the ability to check intervals


Patch:
diff --git a/docs/guides/signals.md b/docs/guides/signals.md
index bc76d9d6..693594a5 100644
--- a/docs/guides/signals.md
+++ b/docs/guides/signals.md
@@ -131,3 +131,19 @@ from sqlmesh import signal, DatetimeRanges, ExecutionContext
 def one_week_ago(batch: DatetimeRanges, context: ExecutionContext) -> t.Union[bool, DatetimeRanges]:
     return len(context.engine_adapter.fetchdf("SELECT 1")) > 1
 ```
+
+### Testing Signals
+Signals only evaluate on `run` or with `check_intervals`.
+
+To test signals with the [check_intervals](../reference/cli.md#check_intervals) command:
+
+1. Deploy your changes to an environment with `sqlmesh plan my_dev`.
+2. Run `sqlmesh check_intervals my_dev`.
+
+   * To check a subset of models use the --select-model flag.
+   * To turn off signals and just check missing intervals, use the --no-signals flag.
+
+3. To iterate, make changes to the signal, and redeploy with step 1.
+
+!!! note
+    `check_intervals` only works on remote models in an environment. Local signal changes are never run.
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index dd958d17..35586b95 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -65,6 +65,24 @@ Options:
   --help                 Show this message and exit.
 ```
 
+## check_intervals
+
+```
+Usage: sqlmesh check_intervals [OPTIONS] [ENVIRONMENT]
+
+  Show missing intervals in an environment, respecting signals.
+
+Options:
+  --no-signals         Disable signal checks and only show missing intervals.
+  --select-model TEXT  Select specific models to show missing intervals for.
+  -s, --start TEXT     The start datetime of the interval for which this
+                       command will be applied.
+  -e, --end TEXT       The end datetime of the interval for which this command
+                       will be applied.
+  --help               Show this message and exit.
+```
+
+
 ## clean
 
 ```
@@ -584,4 +602,4 @@ Options:
   --model TEXT           A model to lint. Multiple models can be linted.  If no models are specified, every model will be linted.
   --help                 Show this message and exit.
 
-```
\ No newline at end of file
+```
diff --git a/docs/reference/notebook.md b/docs/reference/notebook.md
index 50564be3..60a7e16e 100644
--- a/docs/reference/notebook.md
+++ b/docs/reference/notebook.md
@@ -438,6 +438,27 @@ options:
                         Execution time.
 ```
 
+#### check_intervals
+```
+%check_intervals [--no-signals] [--select-model [SELECT_MODEL ...]]
+                   [--start START] [--end END]
+                   [environment]
+
+Show missing intervals in an environment, respecting signals.
+
+positional arguments:
+  environment           The environment to check intervals for.
+
+options:
+  --no-signals          Disable signal checks and only show missing intervals.
+  --select-model <[SELECT_MODEL ...]>
+                        Select specific model changes that should be included
+                        in the plan.
+  --start START, -s START
+                        Start date of intervals to check for.
+  --end END, -e END     End date of intervals to check for.
+```
+
 #### rollback
 ```
 %rollback
@@ -499,10 +520,10 @@ options:
 
 #### lint
 ```
-%lint [--models ...] 
+%lint [--models ...]
 
 Run the linter on the target models(s)
 
 positional arguments:
   --models                A model to lint. Multiple models can be linted. If no models are specified, every model will be linted.
-```
\ No newline at end of file
+```
diff --git a/mkdocs.yml b/mkdocs.yml
index 212bb495..f6858b66 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -36,6 +36,7 @@ nav:
       - guides/observer.md
     - Advanced usage:
       - guides/customizing_sqlmesh.md
+      - guides/signals.md
   - Concepts:
     - concepts/overview.md
     - Development:
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 50ee33e8..0fbf2f88 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -693,6 +693,42 @@ def audit(
     obj.audit(models=models, start=start, end=end, execution_time=execution_time)
 
 
+@cli.command("check_intervals")
+@click.option(
+    "--no-signals",
+    is_flag=True,
+    help="Disable signal checks and only show missing intervals.",
+    default=False,
+)
+@click.argument("environment", required=False)
+@click.option(
+    "--select-model",
+    type=str,
+    multiple=True,
+    help="Select specific models to show missing intervals for.",
+)
+@opt.start_time
+@opt.end_time
+@click.pass_context
+@error_handler
+@cli_analytics
+def check_intervals(
+    ctx: click.Context,
+    environment: t.Optional[str],
+    no_signals: bool,
+    select_model: t.List[str],
+    start: TimeLike,
+    end: TimeLike,
+) -> None:
+    """Show missing intervals in an environment, respecting signals."""
+    context = ctx.obj
+    context.console.show_intervals(
+        context.check_intervals(
+            environment, no_signals=no_signals, select_models=select_model, start=start, end=end
+        )
+    )
+
+
 @cli.command("fetchdf")
 @click.argument("sql")
 @click.pass_context
diff --git a/sqlmesh/core/console.py b/sqlmesh/core/console.py
index abf7916a..6991e902 100644
--- a/sqlmesh/core/console.py
+++ b/sqlmesh/core/console.py
@@ -191,6 +191,10 @@ class EnvironmentsConsole(abc.ABC):
     def print_environments(self, environments_summary: t.Dict[str, int]) -> None:
         """Prints all environment names along with expiry datetime."""
 
+    @abc.abstractmethod
+    def show_intervals(self, snapshot_intervals: t.Dict[Snapshot, SnapshotIntervals]) -> None:
+        """Show ready intervals"""
+
 
 class DifferenceConsole(abc.ABC):
     """Console for displaying environment differences"""
@@ -658,6 +662,9 @@ class NoopConsole(Console):
     def print_environments(self, environments_summary: t.Dict[str, int]) -> None:
         pass
 
+    def show_intervals(self, snapshot_intervals: t.Dict[Snapshot, SnapshotIntervals]) -> None:
+        pass
+
     def show_linter_violations(
         self, violations: t.List[RuleViolation], model: Model, is_error: bool = False
     ) -> None:
@@ -2091,6 +2098,24 @@ class TerminalConsole(Console):
         output_str = "\n".join([str(len(output)), *output])
         self.log_status_update(f"Number of SQLMesh environments are: {output_str}")
 
+    def show_intervals(self, snapshot_intervals: t.Dict[Snapshot, SnapshotIntervals]) -> None:
+        complete = Tree(f"[b]Complete Intervals[/b]")
+        incomplete = Tree(f"[b]Missing Intervals[/b]")
+
+        for snapshot, intervals in sorted(snapshot_intervals.items(), key=lambda s: s[0].node.name):
+            if intervals.intervals:
+                incomplete.add(
+                    f"{snapshot.node.name}: [{intervals.format_intervals(snapshot.node.interval_unit)}]"
+                )
+            else:
+                complete.add(snapshot.node.name)
+
+        if complete.children:
+            self._print(complete)
+
+        if incomplete.children:
+            self._print(incomplete)
+
     def print_connection_config(self, config: ConnectionConfig, title: str = "Connection") -> None:
         tree = Tree(f"[b]{title}:[/b]")
         tree.add(f"Type: [bold cyan]{config.type_}[/bold cyan]")
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 0aee1770..81b30425 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -87,7 +87,7 @@ from sqlmesh.core.notification_target import (
     NotificationTarget,
     NotificationTargetManager,
 )
-from sqlmesh.core.plan import Plan, PlanBuilder
+from sqlmesh.core.plan import Plan, PlanBuilder, SnapshotIntervals
 from sqlmesh.core.reference import ReferenceGraph
 from sqlmesh.core.scheduler import Scheduler, CompletionStatus
 from sqlmesh.core.schema_loader import create_external_models_file
@@ -97,6 +97,7 @@ from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotEvaluator,
     SnapshotFingerprint,
+    missing_intervals,
     to_table_mapping,
 )
 from sqlmesh.core.snapshot.definition import get_next_model_interval_start
@@ -469,11 +470,12 @@ class GenericContext(BaseContext, t.Generic[C]):
         self,
         deployability_index: t.Optional[DeployabilityIndex] = None,
         engine_adapter: t.Optional[EngineAdapter] = None,
+        snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
     ) -> ExecutionContext:
         """Returns an execution context."""
         return ExecutionContext(
             engine_adapter=engine_adapter or self.engine_adapter,
-            snapshots=self.snapshots,
+            snapshots=snapshots or self.snapshots,
             deployability_index=deployability_index,
             default_dialect=self.default_dialect,
             default_catalog=self.default_catalog,
@@ -1893,6 +1895,61 @@ class GenericContext(BaseContext, t.Generic[C]):
             dialect=dialect or self.default_dialect,
         )
 
+    @python_api_analytics
+    def check_intervals(
+        self,
+        environment: t.Optional[str],
+        no_signals: bool,
+        select_models: t.Collection[str],
+        start: t.Optional[TimeLike] = None,
+        end: t.Optional[TimeLike] = None,
+    ) -> t.Dict[Snapshot, SnapshotIntervals]:
+        """Check intervals for a given environment.
+
+        Args:
+            environment: The environment or prod if None.
+            select_models: A list of model selection strings to show intervals for.
+            start: The start of the intervals to check.
+            end: The end of the intervals to check.
+        """
+
+        environment = environment or c.PROD
+        env = self.state_reader.get_environment(environment)
+        if not env:
+            raise SQLMeshError(f"Environment '{environment}' was not found.")
+
+        snapshots = {k.name: v for k, v in self.state_sync.get_snapshots(env.snapshots).items()}
+
+        missing = {
+            k.name: v
+            for k, v in missing_intervals(
+                snapshots.values(), start=start, end=end, execution_time=end
+            ).items()
+        }
+
+        if select_models:
+            selected: t.Collection[str] = self._select_models_for_run(
+                select_models, True, snapshots.values()
+            )
+        else:
+            selected = snapshots.keys()
+
+        results = {}
+        execution_context = self.execution_context(snapshots=snapshots)
+
+        for fqn in selected:
+            snapshot = snapshots[fqn]
+            intervals = missing.get(fqn) or []
+
+            results[snapshot] = SnapshotIntervals(
+                snapshot.snapshot_id,
+                intervals
+                if no_signals
+                else snapshot.check_ready_intervals(intervals, execution_context),
+            )
+
+        return results
+
     @python_api_analytics
     def migrate(self) -> None:
         """Migrates SQLMesh to the current running version.
diff --git a/sqlmesh/core/engine_adapter/bigquery.py b/sqlmesh/core/engine_adapter/bigquery.py
index 12338d4a..e60e606a 100644
--- a/sqlmesh/core/engine_adapter/bigquery.py
+++ b/sqlmesh/core/engine_adapter/bigquery.py
@@ -611,7 +611,7 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
             if columns_to_types is None or columns_to_types[
                 partition_column.name
             ] == exp.DataType.build("unknown"):
-                columns_to_types = self.columns(temp_table_name)
+                columns_to_types = self.columns(table_name)
 
             partition_type_sql = columns_to_types[partition_column.name].sql(dialect=self.dialect)
 
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index 51f5bb1c..e9ffd9aa 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -1007,6 +1007,38 @@ class SQLMeshMagics(Magics):
             models=args.models, start=args.start, end=args.end, execution_time=args.execution_time
         )
 
+    @magic_arguments()
+    @argument("environment", nargs="?", type=str, help="The environment to check intervals for.")
+    @argument(
+        "--no-signals",
+        action="store_true",
+        help="Disable signal checks and only show missing intervals.",
+        default=False,
+    )
+    @argument(
+        "--select-model",
+        type=str,
+        nargs="*",
+        help="Select specific model changes that should be included in the plan.",
+    )
+    @argument("--start", "-s", type=str, help="Start date of intervals to check for.")
+    @argument("--end", "-e", type=str, help="End date of intervals to check for.")
+    @line_magic
+    @pass_sqlmesh_context
+    def check_intervals(self, context: Context, line: str) -> None:
+        """Show missing intervals in an environment, respecting signals."""
+        args = parse_argstring(self.check_intervals, line)
+
+        context.console.show_intervals(
+            context.check_intervals(
+                environment=args.environment,
+                no_signals=args.no_signals,
+                select_models=args.select_model,
+                start=args.start,
+                end=args.end,
+            )
+        )
+
     @magic_arguments()
     @argument(
         "--skip-connection",
diff --git a/vscode/extension/src/commands/format.ts b/vscode/extension/src/commands/format.ts
index 6202ffc0..ebedd4c4 100644
--- a/vscode/extension/src/commands/format.ts
+++ b/vscode/extension/src/commands/format.ts
@@ -1,29 +1,36 @@
 import { traceError, traceLog } from "../utilities/common/log"
-import { execSync } from 'child_process'
+import { execSync } from "child_process"
 import { sqlmesh_exec } from "../utilities/sqlmesh/sqlmesh"
 import { isErr } from "../utilities/functional/result"
+import * as vscode from "vscode"
 
-interface CLICallout {
-    format: () => Promise<number>
+export const format = async () => {
+  traceLog("Calling format")
+  const out = await internalFormat()
+  if (out === 0) {
+    vscode.window.showInformationMessage("Project formatted successfully")
+  } else {
+    vscode.window.showErrorMessage("Project format failed")
+  }
 }
 
-export const actual_callout: CLICallout = {
-    format: async () => {
-        traceLog("Calling format")
-        try {
-            const exec = await sqlmesh_exec()
-            if (isErr(exec)) {
-                traceError(exec.error)
-                return 1
-            }
-            execSync(`${exec.value.bin} format`, { encoding: 'utf-8', cwd: exec.value.workspacePath, env: exec.value.env })
-            return 0
-        } catch (error: any) {
-            traceError('Error executing sqlmesh format:', error.message)
-            traceError(error.stdout)
-            traceError(error.stderr)
-            return error.status || 1
-        }
+const internalFormat = async (): Promise<number> => {
+  try {
+    const exec = await sqlmesh_exec()
+    if (isErr(exec)) {
+      traceError(exec.error)
+      return 1
     }
+    execSync(`${exec.value.bin} format`, {
+      encoding: "utf-8",
+      cwd: exec.value.workspacePath,
+      env: exec.value.env,
+    })
+    return 0
+  } catch (error: any) {
+    traceError("Error executing sqlmesh format:", error.message)
+    traceError(error.stdout)
+    traceError(error.stderr)
+    return error.status || 1
+  }
 }
-
diff --git a/vscode/extension/src/extension.ts b/vscode/extension/src/extension.ts
index 7a936720..d7fa8840 100644
--- a/vscode/extension/src/extension.ts
+++ b/vscode/extension/src/extension.ts
@@ -1,7 +1,5 @@
-// The module 'vscode' contains the VS Code extensibility API
-// Import the module and reference it with the alias vscode in your code below
 import * as vscode from "vscode"
-import { actual_callout } from "./commands/format"
+import { format } from "./commands/format"
 import {
   createOutputChannel,
   onDidChangeConfiguration,
@@ -48,23 +46,16 @@ export async function activate(context: vscode.ExtensionContext) {
   )
 
   context.subscriptions.push(
-    vscode.commands.registerCommand("sqlmesh.signinSpecifyFlow", signInSpecifyFlow(authProvider)
+    vscode.commands.registerCommand(
+      "sqlmesh.signinSpecifyFlow",
+      signInSpecifyFlow(authProvider)
     )
   )
-
   context.subscriptions.push(
     vscode.commands.registerCommand("sqlmesh.signout", signOut(authProvider))
   )
-
   context.subscriptions.push(
-    vscode.commands.registerCommand("sqlmesh.format", async () => {
-      const out = await actual_callout.format()
-      if (out === 0) {
-        vscode.window.showInformationMessage("Project formatted successfully")
-      } else {
-        vscode.window.showErrorMessage("Project format failed")
-      }
-    })
+    vscode.commands.registerCommand("sqlmesh.format", format)
   )
 
   lspClient = new LSPClient()
@@ -77,12 +68,11 @@ export async function activate(context: vscode.ExtensionContext) {
       await lspClient.restart()
     }
   }
-
   context.subscriptions.push(
     onDidChangePythonInterpreter(async () => {
       await restart()
     }),
-    onDidChangeConfiguration(async (_: vscode.ConfigurationChangeEvent) => {
+    onDidChangeConfiguration(async () => {
       await restart()
     }),
     registerCommand(`sqlmesh.restart`, async () => {


Test patch:
diff --git a/tests/core/engine_adapter/test_bigquery.py b/tests/core/engine_adapter/test_bigquery.py
index 048a23ca..fd5499bf 100644
--- a/tests/core/engine_adapter/test_bigquery.py
+++ b/tests/core/engine_adapter/test_bigquery.py
@@ -116,9 +116,7 @@ def test_insert_overwrite_by_partition_query_unknown_column_types(
         },
     )
 
-    columns_mock.assert_called_once_with(
-        exp.to_table(f"test_schema.__temp_test_table_{temp_table_id}")
-    )
+    columns_mock.assert_called_once_with(table_name)
 
     sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == [
diff --git a/tests/core/test_context.py b/tests/core/test_context.py
index f14f472b..e12cf53c 100644
--- a/tests/core/test_context.py
+++ b/tests/core/test_context.py
@@ -1993,3 +1993,36 @@ def test_plan_audit_intervals(tmp_path: pathlib.Path, capsys, caplog):
         """SELECT COUNT(*) FROM (SELECT ("date_id") AS "date_id" FROM (SELECT * FROM "sqlmesh__sqlmesh_audit"."sqlmesh_audit__date_example__4100277424" AS "sqlmesh_audit__date_example__4100277424" WHERE "date_id" BETWEEN CAST('2025-02-01' AS DATE) AND CAST('2025-02-01' AS DATE)) AS "_q_0" WHERE TRUE GROUP BY ("date_id") HAVING COUNT(*) > 1) AS "audit\""""
         in caplog.text
     )
+
+
+def test_check_intervals(sushi_context, mocker):
+    with pytest.raises(
+        SQLMeshError,
+        match="Environment 'dev' was not found",
+    ):
+        sushi_context.check_intervals(environment="dev", no_signals=False, select_models=[])
+
+    spy = mocker.spy(sqlmesh.core.snapshot.definition, "_check_ready_intervals")
+    intervals = sushi_context.check_intervals(environment=None, no_signals=False, select_models=[])
+
+    min_intervals = 19
+    assert spy.call_count == 1
+    assert len(intervals) >= min_intervals
+
+    for i in intervals.values():
+        assert not i.intervals
+
+    spy.reset_mock()
+    intervals = sushi_context.check_intervals(environment=None, no_signals=True, select_models=[])
+    assert spy.call_count == 0
+    assert len(intervals) >= min_intervals
+
+    intervals = sushi_context.check_intervals(
+        environment=None, no_signals=False, select_models=["*waiter_as_customer*"]
+    )
+    assert len(intervals) == 1
+
+    intervals = sushi_context.check_intervals(
+        environment=None, no_signals=False, select_models=["*waiter_as_customer*"], end="next week"
+    )
+    assert tuple(intervals.values())[0].intervals