PR Feedback 3

Author: VaggelisD

sqlmesh/core/context.py

@@ -446,18 +446,22 @@ def engine_adapter(self) -> EngineAdapter:
             self._engine_adapter = self.connection_config.create_engine_adapter()
         return self._engine_adapter
 
-    def snapshot_evaluator(self, job_id: t.Optional[CorrelationId] = None) -> SnapshotEvaluator:
-        # Cache snapshot evaluators by job_id to avoid old job_ids being attached to future Context operations
-        if job_id not in self._snapshot_evaluators:
-            self._snapshot_evaluators[job_id] = SnapshotEvaluator(
+    def snapshot_evaluator(
+        self, correlation_id: t.Optional[CorrelationId] = None
+    ) -> SnapshotEvaluator:
+        # Cache snapshot evaluators by correlation_id to avoid old correlation_ids being attached to future Context operations
+        if correlation_id not in self._snapshot_evaluators:
+            self._snapshot_evaluators[correlation_id] = SnapshotEvaluator(
                 {
-                    gateway: adapter.with_settings(log_level=logging.INFO, job_id=job_id)
+                    gateway: adapter.with_settings(
+                        log_level=logging.INFO, correlation_id=correlation_id
+                    )
                     for gateway, adapter in self.engine_adapters.items()
                 },
                 ddl_concurrent_tasks=self.concurrent_tasks,
                 selected_gateway=self.selected_gateway,
             )
-        return self._snapshot_evaluators[job_id]
+        return self._snapshot_evaluators[correlation_id]
 
     def execution_context(
         self,
@@ -539,7 +543,7 @@ def scheduler(self, environment: t.Optional[str] = None) -> Scheduler:
         return self.create_scheduler(snapshots)
 
     def create_scheduler(
-        self, snapshots: t.Iterable[Snapshot], job_id: t.Optional[CorrelationId] = None
+        self, snapshots: t.Iterable[Snapshot], correlation_id: t.Optional[CorrelationId] = None
     ) -> Scheduler:
         """Creates the built-in scheduler.
 
@@ -551,7 +555,7 @@ def create_scheduler(
         """
         return Scheduler(
             snapshots,
-            self.snapshot_evaluator(job_id),
+            self.snapshot_evaluator(correlation_id),
             self.state_sync,
             default_catalog=self.default_catalog,
             max_workers=self.concurrent_tasks,
@@ -1592,7 +1596,9 @@ def apply(
             )
             explainer.evaluate(
                 plan.to_evaluatable(),
-                snapshot_evaluator=self.snapshot_evaluator(job_id=CorrelationId.from_plan(plan)),
+                snapshot_evaluator=self.snapshot_evaluator(
+                    correlation_id=CorrelationId.from_plan(plan)
+                ),
             )
             return
 
@@ -2341,9 +2347,12 @@ def close(self) -> None:
         """Releases all resources allocated by this context."""
         for evaluator in self._snapshot_evaluators.values():
             evaluator.close()
+
         if self._state_sync:
             self._state_sync.close()
 
+        self._snapshot_evaluators.clear()
+
     def _run(
         self,
         environment: str,
@@ -2395,7 +2404,9 @@ def _run(
     def _apply(self, plan: Plan, circuit_breaker: t.Optional[t.Callable[[], bool]]) -> None:
         self._scheduler.create_plan_evaluator(self).evaluate(
             plan.to_evaluatable(),
-            snapshot_evaluator=self.snapshot_evaluator(job_id=CorrelationId.from_plan(plan)),
+            snapshot_evaluator=self.snapshot_evaluator(
+                correlation_id=CorrelationId.from_plan(plan)
+            ),
             circuit_breaker=circuit_breaker,
         )
 

sqlmesh/core/engine_adapter/base.py

@@ -123,7 +123,7 @@ def __init__(
         pre_ping: bool = False,
         pretty_sql: bool = False,
         shared_connection: bool = False,
-        job_id: t.Optional[CorrelationId] = None,
+        correlation_id: t.Optional[CorrelationId] = None,
         **kwargs: t.Any,
     ):
         self.dialect = dialect.lower() or self.DIALECT
@@ -145,7 +145,7 @@ def __init__(
         self._pre_ping = pre_ping
         self._pretty_sql = pretty_sql
         self._multithreaded = multithreaded
-        self._job_id = job_id
+        self.correlation_id = correlation_id
 
     def with_settings(self, log_level: int, **kwargs: t.Any) -> EngineAdapter:
         adapter = self.__class__(
@@ -2178,8 +2178,8 @@ def execute(
                 else:
                     sql = t.cast(str, e)
 
-                if self._job_id:
-                    sql = f"/* {self._job_id} */ {sql}"
+                if self.correlation_id:
+                    sql = f"/* {self.correlation_id} */ {sql}"
 
                 self._log_sql(
                     sql,

sqlmesh/utils/__init__.py

@@ -390,9 +390,7 @@ def to_snake_case(name: str) -> str:
 
 class JobType(Enum):
     PLAN = "SQLMESH_PLAN"
-
-
-RUN = "SQLMESH_RUN"
+    RUN = "SQLMESH_RUN"
 
 
 @dataclass(frozen=True)