Refactor to fetch all jobs across all threads

Author: georgesittas

sqlmesh/core/engine_adapter/bigquery.py

@@ -188,22 +188,27 @@ def query_factory() -> Query:
         ]
 
     def close(self) -> t.Any:
-        query_job = self._query_job
-        if not query_job:
-            return super().close()
-
-        # Cancel the last submitted query job if it's still pending, to avoid it becoming orphan (e.g., if interrupted)
-        try:
-            if not self._db_call(query_job.done):
-                self._db_call(query_job.cancel)
-        except Exception as ex:
-            logger.debug(
-                "Failed to cancel BigQuery job: https://console.cloud.google.com/bigquery?project=%s&j=bq:%s:%s. %s",
-                query_job.project,
-                query_job.location,
-                query_job.job_id,
-                str(ex),
-            )
+        # Cancel all pending query jobs across all threads
+        all_query_jobs = self._connection_pool.get_all_attributes("query_job")
+        for query_job in all_query_jobs:
+            if query_job:
+                try:
+                    if not self._db_call(query_job.done):
+                        self._db_call(query_job.cancel)
+                        logger.debug(
+                            "Cancelled BigQuery job: https://console.cloud.google.com/bigquery?project=%s&j=bq:%s:%s",
+                            query_job.project,
+                            query_job.location,
+                            query_job.job_id,
+                        )
+                except Exception as ex:
+                    logger.debug(
+                        "Failed to cancel BigQuery job: https://console.cloud.google.com/bigquery?project=%s&j=bq:%s:%s. %s",
+                        query_job.project,
+                        query_job.location,
+                        query_job.job_id,
+                        str(ex),
+                    )
 
         return super().close()
 

sqlmesh/utils/connection_pool.py

@@ -48,6 +48,17 @@ def set_attribute(self, key: str, value: t.Any) -> None:
             value: Attribute value.
         """
 
+    @abc.abstractmethod
+    def get_all_attributes(self, key: str) -> t.List[t.Any]:
+        """Returns all attributes with the given key across all connections/threads.
+
+        Args:
+            key: Attribute key.
+
+        Returns:
+            List of attribute values from all connections/threads.
+        """
+
     @abc.abstractmethod
     def begin(self) -> None:
         """Starts a new transaction."""
@@ -142,6 +153,14 @@ def set_attribute(self, key: str, value: t.Any) -> None:
         thread_id = get_ident()
         self._thread_attributes[thread_id][key] = value
 
+    def get_all_attributes(self, key: str) -> t.List[t.Any]:
+        """Returns all attributes with the given key across all threads."""
+        return [
+            thread_attrs[key]
+            for thread_attrs in self._thread_attributes.values()
+            if key in thread_attrs
+        ]
+
     def begin(self) -> None:
         self._do_begin()
         with self._thread_transactions_lock:
@@ -282,6 +301,11 @@ def get_attribute(self, key: str) -> t.Optional[t.Any]:
     def set_attribute(self, key: str, value: t.Any) -> None:
         self._attributes[key] = value
 
+    def get_all_attributes(self, key: str) -> t.List[t.Any]:
+        """Returns all attributes with the given key (single-threaded pool has at most one)."""
+        value = self._attributes.get(key)
+        return [value] if value is not None else []
+
     def begin(self) -> None:
         self._do_begin()
         self._is_transaction_active = True