Move all tracking into snapshot evaluator, remove seed tracker class

Author: treysp

sqlmesh/core/engine_adapter/base.py

@@ -40,7 +40,7 @@
 )
 from sqlmesh.core.model.kind import TimeColumn
 from sqlmesh.core.schema_diff import SchemaDiffer
-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, random_id
 from sqlmesh.utils.connection_pool import ConnectionPool, create_connection_pool
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_time_column
@@ -2283,7 +2283,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:
@@ -2292,7 +2296,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(

sqlmesh/core/engine_adapter/bigquery.py

@@ -20,7 +20,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 SchemaDiffer
 from sqlmesh.utils import optional_import
@@ -1091,7 +1091,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

sqlmesh/core/execution_tracker.py

@@ -5,6 +5,7 @@
 from contextlib import contextmanager
 from threading import local
 from dataclasses import dataclass, field
+from sqlmesh.utils.errors import SQLMeshError
 
 
 @dataclass
@@ -27,7 +28,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,28 +46,41 @@ 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
+
+        if snapshot_id_batch in cls._contexts:
+            raise SQLMeshError(
+                f"Snapshot ID batch {snapshot_id_batch} execution has already been tracked. Each snapshot should only be tracked once."
+            )
+
+        context = QueryExecutionContext(id=snapshot_id_batch)
         cls._thread_local.context = context
+        cls._contexts[snapshot_id_batch] = context
         try:
             yield context
         finally:
@@ -76,67 +89,12 @@ def track_execution(cls, snapshot_name_batch: str) -> t.Iterator[QueryExecutionC
 
     @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()
-        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)
+    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
-
-    @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)

sqlmesh/core/scheduler.py

@@ -7,7 +7,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
@@ -427,69 +427,59 @@ def evaluate_node(node: SchedulingUnit) -> None:
                 return
             snapshot = self.snapshots_by_name[snapshot_name]
 
-            with QueryExecutionTracker.track_execution(
-                f"{snapshot.name}_{batch_idx}"
-            ) as execution_context:
-                self.console.start_snapshot_evaluation_progress(snapshot)
-
-                execution_start_ts = now_timestamp()
-                evaluation_duration_ms: t.Optional[int] = None
-
-                audit_results: t.List[AuditResult] = []
-                try:
-                    assert execution_time  # mypy
-                    assert deployability_index  # mypy
-
-                    if audit_only:
-                        audit_results = self._audit_snapshot(
-                            snapshot=snapshot,
-                            environment_naming_info=environment_naming_info,
-                            deployability_index=deployability_index,
-                            snapshots=self.snapshots_by_name,
-                            start=start,
-                            end=end,
-                            execution_time=execution_time,
-                        )
-                    else:
-                        audit_results = self.evaluate(
-                            snapshot=snapshot,
-                            environment_naming_info=environment_naming_info,
-                            start=start,
-                            end=end,
-                            execution_time=execution_time,
-                            deployability_index=deployability_index,
-                            batch_index=batch_idx,
-                        )
-
-                    evaluation_duration_ms = now_timestamp() - execution_start_ts
-                finally:
-                    num_audits = len(audit_results)
-                    num_audits_failed = sum(1 for result in audit_results if result.count)
-
-                    rows_processed = None
-                    if snapshot.is_seed:
-                        # seed stats are tracked in SeedStrategy.create by model name, not snapshot name
-                        seed_stats = SeedExecutionTracker.get_and_clear_seed_stats(
-                            snapshot.model.name
-                        )
-                        rows_processed = (
-                            seed_stats.get("total_rows_processed") if seed_stats else None
-                        )
-                    else:
-                        rows_processed = (
-                            execution_context.total_rows_processed if execution_context else None
-                        )
-
-                    self.console.update_snapshot_evaluation_progress(
-                        snapshot,
-                        batched_intervals[snapshot][batch_idx],
-                        batch_idx,
-                        evaluation_duration_ms,
-                        num_audits - num_audits_failed,
-                        num_audits_failed,
-                        rows_processed=rows_processed,
+            self.console.start_snapshot_evaluation_progress(snapshot)
+
+            execution_start_ts = now_timestamp()
+            evaluation_duration_ms: t.Optional[int] = None
+
+            audit_results: t.List[AuditResult] = []
+            try:
+                assert execution_time  # mypy
+                assert deployability_index  # mypy
+
+                if audit_only:
+                    audit_results = self._audit_snapshot(
+                        snapshot=snapshot,
+                        environment_naming_info=environment_naming_info,
+                        deployability_index=deployability_index,
+                        snapshots=self.snapshots_by_name,
+                        start=start,
+                        end=end,
+                        execution_time=execution_time,
+                    )
+                else:
+                    audit_results = self.evaluate(
+                        snapshot=snapshot,
+                        environment_naming_info=environment_naming_info,
+                        start=start,
+                        end=end,
+                        execution_time=execution_time,
+                        deployability_index=deployability_index,
+                        batch_index=batch_idx,
                     )
 
+                evaluation_duration_ms = now_timestamp() - execution_start_ts
+            finally:
+                num_audits = len(audit_results)
+                num_audits_failed = sum(1 for result in audit_results if result.count)
+
+                execution_stats = QueryExecutionTracker.get_execution_stats(
+                    f"{snapshot.snapshot_id}_{batch_idx}"
+                )
+                rows_processed = (
+                    execution_stats["total_rows_processed"] if execution_stats else None
+                )
+
+                self.console.update_snapshot_evaluation_progress(
+                    snapshot,
+                    batched_intervals[snapshot][batch_idx],
+                    batch_idx,
+                    evaluation_duration_ms,
+                    num_audits - num_audits_failed,
+                    num_audits_failed,
+                    rows_processed=rows_processed,
+                )
+
         try:
             with self.snapshot_evaluator.concurrent_context():
                 errors, skipped_intervals = concurrent_apply_to_dag(
@@ -529,9 +519,6 @@ def evaluate_node(node: SchedulingUnit) -> None:
 
             self.state_sync.recycle()
 
-            # Clean up any remaining seed execution stats
-            SeedExecutionTracker.clear_all_seed_stats()
-
     def _dag(self, batches: SnapshotToIntervals) -> DAG[SchedulingUnit]:
         """Builds a DAG of snapshot intervals to be evaluated.