Add multithread context manager to the snapshot evaluator

Author: izeigerman

sqlmesh/core/config.py

@@ -11,7 +11,7 @@
     import duckdb
     from sqlmesh.core.engine_adapter import EngineAdapter
     local_config = Config(
-        engine_config_factory=duckdb.connect,
+        engine_connection_factory=duckdb.connect,
         engine_dialect="duckdb"
     )
     # End config.py
@@ -274,7 +274,7 @@ class Config(PydanticModel):
     engine_dialect: str = "duckdb"
     scheduler_backend: SchedulerBackend = BuiltInSchedulerBackend()
     notification_targets: t.List[NotificationTarget] = []
-    dialect: t.Optional[str] = None
+    dialect: str = ""
     physical_schema: str = ""
     snapshot_ttl: str = ""
     ignore_patterns: t.List[str] = []

sqlmesh/core/scheduler.py

@@ -136,7 +136,7 @@ def run(
             # We have to run all batches per snapshot to mark it as completed
             self.console.start_snapshot_progress(snapshot.name, len(intervals))
 
-        with ThreadPoolExecutor() as snapshot_pool, ThreadPoolExecutor(
+        with self.snapshot_evaluator.multithreaded_context(), ThreadPoolExecutor() as snapshot_pool, ThreadPoolExecutor(
             max_workers=self.max_workers
         ) as batch_pool:
             while True:
@@ -168,8 +168,6 @@ def run(
         else:
             self.console.complete_snapshot_progress()
 
-        self.snapshot_evaluator.recycle()
-
         return self.failed
 
     def interval_params(

sqlmesh/core/snapshot_evaluator.py

@@ -23,6 +23,7 @@
 
 import logging
 import typing as t
+from contextlib import contextmanager
 
 from sqlglot import exp, select
 
@@ -137,12 +138,12 @@ def promote(
             target_snapshots: Snapshots to promote.
             environment: The target environment.
         """
-        concurrent_apply_to_snapshots(
-            target_snapshots,
-            lambda s: self._promote_snapshot(s, environment),
-            self.ddl_concurrent_tasks,
-        )
-        self.recycle()
+        with self.multithreaded_context():
+            concurrent_apply_to_snapshots(
+                target_snapshots,
+                lambda s: self._promote_snapshot(s, environment),
+                self.ddl_concurrent_tasks,
+            )
 
     def demote(
         self, target_snapshots: t.Iterable[SnapshotInfoLike], environment: str
@@ -153,12 +154,12 @@ def demote(
             target_snapshots: Snapshots to demote.
             environment: The target environment.
         """
-        concurrent_apply_to_snapshots(
-            target_snapshots,
-            lambda s: self._demote_snapshot(s, environment),
-            self.ddl_concurrent_tasks,
-        )
-        self.recycle()
+        with self.multithreaded_context():
+            concurrent_apply_to_snapshots(
+                target_snapshots,
+                lambda s: self._demote_snapshot(s, environment),
+                self.ddl_concurrent_tasks,
+            )
 
     def create(
         self,
@@ -170,26 +171,26 @@ def create(
         Args:
             target_snapshots: Target snapshost.
         """
-        concurrent_apply_to_snapshots(
-            target_snapshots,
-            lambda s: self._create_snapshot(s, snapshots),
-            self.ddl_concurrent_tasks,
-        )
-        self.recycle()
+        with self.multithreaded_context():
+            concurrent_apply_to_snapshots(
+                target_snapshots,
+                lambda s: self._create_snapshot(s, snapshots),
+                self.ddl_concurrent_tasks,
+            )
 
     def cleanup(self, target_snapshots: t.Iterable[SnapshotInfoLike]) -> None:
         """Cleans up the given snapshots by removing its table
 
         Args:
             target_snapshots: Snapshots to cleanup.
         """
-        concurrent_apply_to_snapshots(
-            target_snapshots,
-            self._cleanup_snapshot,
-            self.ddl_concurrent_tasks,
-            reverse_order=True,
-        )
-        self.recycle()
+        with self.multithreaded_context():
+            concurrent_apply_to_snapshots(
+                target_snapshots,
+                self._cleanup_snapshot,
+                self.ddl_concurrent_tasks,
+                reverse_order=True,
+            )
 
     def audit(
         self,
@@ -235,6 +236,13 @@ def audit(
             results.append(AuditResult(audit=audit, count=count, query=query))
         return results
 
+    @contextmanager
+    def multithreaded_context(self) -> t.Generator[None, None, None]:
+        try:
+            yield
+        finally:
+            self.recycle()
+
     def recycle(self) -> None:
         """Closes all open connections and releases all allocated resources associated with any thread
         except the calling one."""