Move all tracking into snapshot evaluator, remove seed tracker class

treysp · treysp · commit 335feb57e41d · 2025-08-22T16:05:36.000-05:00
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
@@ -40,7 +40,7 @@
 )
 from sqlmesh.core.model.kind import TimeColumn
 from sqlmesh.core.schema_diff import SchemaDiffer, TableAlterOperation
-from sqlmesh.core.execution_tracker import record_execution as track_execution_record
+from sqlmesh.core.execution_tracker import QueryExecutionTracker
 from sqlmesh.utils import (
     CorrelationId,
     columns_to_types_all_known,
@@ -2443,7 +2443,11 @@ def _log_sql(
     def _execute(self, sql: str, track_row_count: bool = False, **kwargs: t.Any) -> None:
         self.cursor.execute(sql, **kwargs)
 
-        if track_row_count and self.SUPPORTS_QUERY_EXECUTION_TRACKING:
+        if (
+            self.SUPPORTS_QUERY_EXECUTION_TRACKING
+            and track_row_count
+            and QueryExecutionTracker.is_tracking()
+        ):
             rowcount_raw = getattr(self.cursor, "rowcount", None)
             rowcount = None
             if rowcount_raw is not None:
@@ -2452,7 +2456,7 @@ def _execute(self, sql: str, track_row_count: bool = False, **kwargs: t.Any) ->
                 except (TypeError, ValueError):
                     pass
 
-            track_execution_record(sql, rowcount)
+            QueryExecutionTracker.record_execution(sql, rowcount)
 
     @contextlib.contextmanager
     def temp_table(
diff --git a/sqlmesh/core/engine_adapter/bigquery.py b/sqlmesh/core/engine_adapter/bigquery.py
@@ -21,7 +21,7 @@
     SourceQuery,
     set_catalog,
 )
-from sqlmesh.core.execution_tracker import record_execution as track_execution_record
+from sqlmesh.core.execution_tracker import QueryExecutionTracker
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.schema_diff import TableAlterOperation, NestedSupport
 from sqlmesh.utils import optional_import, get_source_columns_to_types
@@ -1104,7 +1104,7 @@ def _execute(
             elif query_job.statement_type in ["INSERT", "DELETE", "MERGE", "UPDATE"]:
                 num_rows = query_job.num_dml_affected_rows
 
-            track_execution_record(sql, num_rows)
+            QueryExecutionTracker.record_execution(sql, num_rows)
 
     def _get_data_objects(
         self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
diff --git a/sqlmesh/core/execution_tracker.py b/sqlmesh/core/execution_tracker.py
@@ -27,7 +27,6 @@ class QueryExecutionContext:
     queries_executed: t.List[t.Tuple[str, t.Optional[int], float]] = field(default_factory=list)
 
     def add_execution(self, sql: str, row_count: t.Optional[int]) -> None:
-        """Record a single query execution."""
         if row_count is not None and row_count >= 0:
             self.total_rows_processed += row_count
         self.query_count += 1
@@ -46,97 +45,49 @@ def get_execution_stats(self) -> t.Dict[str, t.Any]:
 
 class QueryExecutionTracker:
     """
-    Thread-local context manager for snapshot evaluation execution statistics, such as
+    Thread-local context manager for snapshot execution statistics, such as
     rows processed.
     """
 
     _thread_local = local()
+    _contexts: t.Dict[str, QueryExecutionContext] = {}
 
     @classmethod
-    def get_execution_context(cls) -> t.Optional[QueryExecutionContext]:
-        return getattr(cls._thread_local, "context", None)
+    def get_execution_context(cls, snapshot_id_batch: str) -> t.Optional[QueryExecutionContext]:
+        return cls._contexts.get(snapshot_id_batch)
 
     @classmethod
     def is_tracking(cls) -> bool:
-        return cls.get_execution_context() is not None
+        return getattr(cls._thread_local, "context", None) is not None
 
     @classmethod
     @contextmanager
-    def track_execution(cls, snapshot_name_batch: str) -> t.Iterator[QueryExecutionContext]:
+    def track_execution(
+        cls, snapshot_id_batch: str, condition: bool = True
+    ) -> t.Iterator[t.Optional[QueryExecutionContext]]:
         """
-        Context manager for tracking snapshot evaluation execution statistics.
+        Context manager for tracking snapshot execution statistics.
         """
-        context = QueryExecutionContext(id=snapshot_name_batch)
+        if not condition:
+            yield None
+            return
+
+        context = QueryExecutionContext(id=snapshot_id_batch)
         cls._thread_local.context = context
+        cls._contexts[snapshot_id_batch] = context
         try:
             yield context
         finally:
-            if hasattr(cls._thread_local, "context"):
-                delattr(cls._thread_local, "context")
+            cls._thread_local.context = None
 
     @classmethod
     def record_execution(cls, sql: str, row_count: t.Optional[int]) -> None:
-        context = cls.get_execution_context()
+        context = getattr(cls._thread_local, "context", None)
         if context is not None:
             context.add_execution(sql, row_count)
 
     @classmethod
-    def get_execution_stats(cls) -> t.Optional[t.Dict[str, t.Any]]:
-        context = cls.get_execution_context()
+    def get_execution_stats(cls, snapshot_id_batch: str) -> t.Optional[t.Dict[str, t.Any]]:
+        context = cls.get_execution_context(snapshot_id_batch)
+        cls._contexts.pop(snapshot_id_batch, None)
         return context.get_execution_stats() if context else None
-
-
-class SeedExecutionTracker:
-    _seed_contexts: t.Dict[str, QueryExecutionContext] = {}
-    _thread_local = local()
-
-    @classmethod
-    @contextmanager
-    def track_execution(cls, model_name: str) -> t.Iterator[QueryExecutionContext]:
-        """
-        Context manager for tracking seed creation execution statistics.
-        """
-        context = QueryExecutionContext(id=model_name)
-        cls._seed_contexts[model_name] = context
-        cls._thread_local.seed_id = model_name
-
-        try:
-            yield context
-        finally:
-            if hasattr(cls._thread_local, "seed_id"):
-                delattr(cls._thread_local, "seed_id")
-
-    @classmethod
-    def get_and_clear_seed_stats(cls, model_name: str) -> t.Optional[t.Dict[str, t.Any]]:
-        context = cls._seed_contexts.pop(model_name, None)
-        return context.get_execution_stats() if context else None
-
-    @classmethod
-    def clear_all_seed_stats(cls) -> None:
-        """Clear all remaining seed stats. Used for cleanup after evaluation completes."""
-        cls._seed_contexts.clear()
-
-    @classmethod
-    def is_tracking(cls) -> bool:
-        return hasattr(cls._thread_local, "seed_id")
-
-    @classmethod
-    def record_execution(cls, sql: str, row_count: t.Optional[int]) -> None:
-        seed_id = getattr(cls._thread_local, "seed_id", None)
-        if seed_id:
-            context = cls._seed_contexts.get(seed_id)
-            if context is not None:
-                context.add_execution(sql, row_count)
-
-
-def record_execution(sql: str, row_count: t.Optional[int]) -> None:
-    """
-    Record execution statistics for a single SQL statement.
-
-    Automatically infers which tracker is active based on the current thread.
-    """
-    if SeedExecutionTracker.is_tracking():
-        SeedExecutionTracker.record_execution(sql, row_count)
-        return
-    if QueryExecutionTracker.is_tracking():
-        QueryExecutionTracker.record_execution(sql, row_count)
diff --git a/sqlmesh/core/scheduler.py b/sqlmesh/core/scheduler.py
@@ -9,7 +9,7 @@
 from sqlmesh.core import constants as c
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.environment import EnvironmentNamingInfo, execute_environment_statements
-from sqlmesh.core.execution_tracker import QueryExecutionTracker, SeedExecutionTracker
+from sqlmesh.core.execution_tracker import QueryExecutionTracker
 from sqlmesh.core.macros import RuntimeStage
 from sqlmesh.core.model.definition import AuditResult
 from sqlmesh.core.node import IntervalUnit
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
@@ -39,7 +39,7 @@
 from sqlmesh.core.dialect import schema_
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.engine_adapter.shared import InsertOverwriteStrategy, DataObjectType
-from sqlmesh.core.execution_tracker import SeedExecutionTracker
+from sqlmesh.core.execution_tracker import QueryExecutionTracker
 from sqlmesh.core.macros import RuntimeStage
 from sqlmesh.core.model import (
     AuditResult,
@@ -170,19 +170,22 @@ def evaluate(
         Returns:
             The WAP ID of this evaluation if supported, None otherwise.
         """
-        result = self._evaluate_snapshot(
-            start=start,
-            end=end,
-            execution_time=execution_time,
-            snapshot=snapshot,
-            snapshots=snapshots,
-            allow_destructive_snapshots=allow_destructive_snapshots or set(),
-            allow_additive_snapshots=allow_additive_snapshots or set(),
-            deployability_index=deployability_index,
-            batch_index=batch_index,
-            target_table_exists=target_table_exists,
-            **kwargs,
-        )
+        with QueryExecutionTracker.track_execution(
+            f"{snapshot.snapshot_id}_{batch_index}", condition=not snapshot.is_seed
+        ):
+            result = self._evaluate_snapshot(
+                start=start,
+                end=end,
+                execution_time=execution_time,
+                snapshot=snapshot,
+                snapshots=snapshots,
+                allow_destructive_snapshots=allow_destructive_snapshots or set(),
+                allow_additive_snapshots=allow_additive_snapshots or set(),
+                deployability_index=deployability_index,
+                batch_index=batch_index,
+                target_table_exists=target_table_exists,
+                **kwargs,
+            )
         if result is None or isinstance(result, str):
             return result
         raise SQLMeshError(
diff --git a/tests/core/test_execution_tracker.py b/tests/core/test_execution_tracker.py
@@ -1,74 +1,38 @@
-# Tests the sqlmesh.core.execution_tracker module
-# - creates a scenario where executions will take place in multiple threads
-# - generates the scenario with known numbers of rows to be processed
-# - tests that the execution tracker correctly tracks the number of rows processed in both threads
-# - may use mocks, an existing test project, manually created snapshots, or a duckdb database to create the scenario
-
 from __future__ import annotations
 
-import threading
-from queue import Queue
-from typing import List, Optional
+import typing as t
+from concurrent.futures import ThreadPoolExecutor
 
 from sqlmesh.core.execution_tracker import QueryExecutionTracker
 
 
-def test_execution_tracker_thread_isolation_and_aggregation() -> None:
-    """
-    Two worker threads each track executions in their own context. Verify:
-    - isolation across threads
-    - correct aggregation of rows
-    - query metadata is captured
-    - main thread has no active tracking
-    """
-
-    assert not QueryExecutionTracker.is_tracking()
-    assert QueryExecutionTracker.get_execution_stats() is None
-
-    counts_a: List[Optional[int]] = [10, 5, None]
-    counts_b: List[Optional[int]] = [3, 7]
-
-    start_barrier = threading.Barrier(3)  # 2 workers + main
-    results: "Queue[dict]" = Queue()
-
-    def worker(batch_id: str, counts: List[Optional[int]]) -> None:
-        with QueryExecutionTracker.track_execution(batch_id) as ctx:
-            # tracking active in this thread
+def test_execution_tracker_thread_isolation() -> None:
+    def worker(id: str, row_counts: list[int]) -> t.Dict[str, t.Any]:
+        with QueryExecutionTracker.track_execution(id) as ctx:
             assert QueryExecutionTracker.is_tracking()
-            # synchronize start to overlap execution
-            start_barrier.wait()
-            for c in counts:
-                QueryExecutionTracker.record_execution("SELECT 1", c)
 
-            stats = ctx.get_execution_stats()
+            for count in row_counts:
+                QueryExecutionTracker.record_execution("SELECT 1", count)
 
-            assert stats["snapshot_batch"] == batch_id
-            assert stats["query_count"] == len(counts)
-            results.put(stats)
+            assert ctx is not None
+            return ctx.get_execution_stats()
 
-    t1 = threading.Thread(target=worker, args=("batch_A", counts_a))
-    t2 = threading.Thread(target=worker, args=("batch_B", counts_b))
-
-    t1.start()
-    t2.start()
-    # Release workers at the same time
-    start_barrier.wait()
-    t1.join()
-    t2.join()
+    with ThreadPoolExecutor() as executor:
+        futures = [
+            executor.submit(worker, "batch_A", [10, 5]),
+            executor.submit(worker, "batch_B", [3, 7]),
+        ]
+        results = [f.result() for f in futures]
 
     # Main thread has no active tracking context
     assert not QueryExecutionTracker.is_tracking()
     QueryExecutionTracker.record_execution("q", 10)
-    assert QueryExecutionTracker.get_execution_stats() is None
-
-    collected = [results.get_nowait(), results.get_nowait()]
-    # by name since order is non-deterministic
-    by_batch = {s["snapshot_batch"]: s for s in collected}
+    assert QueryExecutionTracker.get_execution_stats("q") is None
 
-    stats_a = by_batch["batch_A"]
-    assert stats_a["total_rows_processed"] == 15  # 10 + 5 + 0 (None)
-    assert stats_a["query_count"] == len(counts_a)
+    # Order of results is not deterministic, so look up by id
+    by_batch = {s["id"]: s for s in results}
 
-    stats_b = by_batch["batch_B"]
-    assert stats_b["total_rows_processed"] == 10  # 3 + 7
-    assert stats_b["query_count"] == len(counts_b)
+    assert by_batch["batch_A"]["total_rows_processed"] == 15
+    assert by_batch["batch_A"]["query_count"] == 2
+    assert by_batch["batch_B"]["total_rows_processed"] == 10
+    assert by_batch["batch_B"]["query_count"] == 2
