sqlmesh 0.209.1.dev17__py3-none-any.whl → 0.209.1.dev18__py3-none-any.whl

sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.209.1.dev17'
-__version_tuple__ = version_tuple = (0, 209, 1, 'dev17')
+__version__ = version = '0.209.1.dev18'
+__version_tuple__ = version_tuple = (0, 209, 1, 'dev18')

sqlmesh/core/engine_adapter/bigquery.py

@@ -33,6 +33,7 @@ if t.TYPE_CHECKING:
     from google.cloud import bigquery
     from google.cloud.bigquery import StandardSqlDataType
     from google.cloud.bigquery.client import Client as BigQueryClient
+    from google.cloud.bigquery.job import QueryJob
     from google.cloud.bigquery.job.base import _AsyncJob as BigQueryQueryResult
     from google.cloud.bigquery.table import Table as BigQueryTable
 
@@ -186,6 +187,31 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
             )
         ]
 
+    def close(self) -> t.Any:
+        # 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()
+
     def _begin_session(self, properties: SessionProperties) -> None:
         from google.cloud.bigquery import QueryJobConfig
 
@@ -318,7 +344,10 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
         if len(table.parts) == 3 and "." in table.name:
             # The client's `get_table` method can't handle paths with >3 identifiers
             self.execute(exp.select("*").from_(table).limit(0))
-            query_results = self._query_job._query_results
+            query_job = self._query_job
+            assert query_job is not None
+
+            query_results = query_job._query_results
             columns = create_mapping_schema(query_results.schema)
         else:
             bq_table = self._get_table(table)
@@ -717,7 +746,9 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
         self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
     ) -> DF:
         self.execute(query, quote_identifiers=quote_identifiers)
-        return self._query_job.to_dataframe()
+        query_job = self._query_job
+        assert query_job is not None
+        return query_job.to_dataframe()
 
     def _create_column_comments(
         self,
@@ -1021,20 +1052,23 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
             job_config=job_config,
             timeout=self._extra_config.get("job_creation_timeout_seconds"),
         )
+        query_job = self._query_job
+        assert query_job is not None
 
         logger.debug(
             "BigQuery job created: https://console.cloud.google.com/bigquery?project=%s&j=bq:%s:%s",
-            self._query_job.project,
-            self._query_job.location,
-            self._query_job.job_id,
+            query_job.project,
+            query_job.location,
+            query_job.job_id,
         )
 
         results = self._db_call(
-            self._query_job.result,
+            query_job.result,
             timeout=self._extra_config.get("job_execution_timeout_seconds"),  # type: ignore
         )
+
         self._query_data = iter(results) if results.total_rows else iter([])
-        query_results = self._query_job._query_results
+        query_results = query_job._query_results
         self.cursor._set_rowcount(query_results)
         self.cursor._set_description(query_results.schema)
 
@@ -1198,15 +1232,15 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
 
     @_query_data.setter
     def _query_data(self, value: t.Any) -> None:
-        return self._connection_pool.set_attribute("query_data", value)
+        self._connection_pool.set_attribute("query_data", value)
 
     @property
-    def _query_job(self) -> t.Any:
+    def _query_job(self) -> t.Optional[QueryJob]:
         return self._connection_pool.get_attribute("query_job")
 
     @_query_job.setter
     def _query_job(self, value: t.Any) -> None:
-        return self._connection_pool.set_attribute("query_job", value)
+        self._connection_pool.set_attribute("query_job", value)
 
     @property
     def _session_id(self) -> t.Any:
@@ -1214,7 +1248,7 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
 
     @_session_id.setter
     def _session_id(self, value: t.Any) -> None:
-        return self._connection_pool.set_attribute("session_id", value)
+        self._connection_pool.set_attribute("session_id", value)
 
 
 class _ErrorCounter:

sqlmesh/utils/connection_pool.py

@@ -48,6 +48,17 @@ class ConnectionPool(abc.ABC):
             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 @@ class _ThreadLocalBase(_TransactionManagementMixin):
         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 @@ class SingletonConnectionPool(_TransactionManagementMixin):
     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

{sqlmesh-0.209.1.dev17.dist-info → sqlmesh-0.209.1.dev18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.209.1.dev17
+Version: 0.209.1.dev18
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.209.1.dev17.dist-info → sqlmesh-0.209.1.dev18.dist-info}/RECORD

@@ -1,5 +1,5 @@
 sqlmesh/__init__.py,sha256=QjwaamJIdW7voXoS880yAx7tHm-Im4_mqGA-XJgafos,7631
-sqlmesh/_version.py,sha256=z0GqQTt3meFnDC_Rsl2GKjLHSTISwLNjBVHIt-zAisw,530
+sqlmesh/_version.py,sha256=VXqoS95PazQMdQuI-Ud6xRAzmn1s_sL1yIIWZt5amJI,530
 sqlmesh/magics.py,sha256=bEzYqXHcM0P3CjyUZ1O9arlmXAdLAhoglKbQ29fD9AE,41842
 sqlmesh/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlmesh/cicd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -61,7 +61,7 @@ sqlmesh/core/engine_adapter/_typing.py,sha256=_63pYCThfC39UtPVHY3N0H6qJ1YBT32Oqw
 sqlmesh/core/engine_adapter/athena.py,sha256=AS1w_l_VrFZsc2CIicElg1dtS9y6q5hHHmTg4ECH-xs,26453
 sqlmesh/core/engine_adapter/base.py,sha256=qefboMkuzRn4ctUyuMZiFtL3tSHBZoMzKcpm_90Nvxw,107386
 sqlmesh/core/engine_adapter/base_postgres.py,sha256=DOsOdCJyQAVfA9S0GznuH0J1bonMqYVP_yJLDYcxFYM,6772
-sqlmesh/core/engine_adapter/bigquery.py,sha256=oL3xISn6Mw2jnJareKkeoG3XFCGW4lVof4OorXdQfBA,51468
+sqlmesh/core/engine_adapter/bigquery.py,sha256=JR7YhT4h6POqE6DhNeBxT83zzDyuHR-dBh3fnmavt18,52811
 sqlmesh/core/engine_adapter/clickhouse.py,sha256=LoPdm_s_s-ZxS1XJPA3Et0iOHeJ1LgqTAFbwBdiqE7A,34912
 sqlmesh/core/engine_adapter/databricks.py,sha256=8ZExPIeT3dDh9MTJW8udvmh9LiY-f9HkRAN1Q5Qebhs,14541
 sqlmesh/core/engine_adapter/duckdb.py,sha256=bv1kfcaDtYlAtJBjxxfteG-dIf31-6r-NipfgtNbFKQ,7133
@@ -273,7 +273,7 @@ sqlmesh/utils/aws.py,sha256=krxyhrKtsMJDEjT567UbGZogjJ9JRHkvAV-Cg3FVndM,1062
 sqlmesh/utils/cache.py,sha256=GITQ3sldwG4vlI0wKIDI7yOr1CND1EI3necAmeCEW8g,4743
 sqlmesh/utils/concurrency.py,sha256=EMgEecfPFN9_2Wjz4TUksJngECGse7fnKVcLswkmhAk,9274
 sqlmesh/utils/config.py,sha256=-SjIn7Z0HHnMtjPavQIvFcSaeg86XLsGd3sGIwOKrTE,1941
-sqlmesh/utils/connection_pool.py,sha256=To1zr9d5MQOZ5im09jLQKO4DFmfzH6s_q75cEa1y0PE,11390
+sqlmesh/utils/connection_pool.py,sha256=-IYbML9HlDN75uhCu-Y3wdsuhlbscE4T89GDI03F0vY,12252
 sqlmesh/utils/conversions.py,sha256=Ul8u6TzlKzFz_lVrUzFSzmJBKHUXOYe2e-xmkMXV084,406
 sqlmesh/utils/cron.py,sha256=eGwn4iUeiRoQzwcd9eS2TZkut8nR4yWud77N7xQ9CQ0,1829
 sqlmesh/utils/dag.py,sha256=L_jwKf81wCySmt1TFs4vEKsrIHSOPTBhptlmv1bsSSg,7025
@@ -291,7 +291,7 @@ sqlmesh/utils/pydantic.py,sha256=o_NsXbIpDqNpUA1Uc5xF0ZzoXQYYB0DfHwdRxBwXPNk,120
 sqlmesh/utils/rich.py,sha256=cwQ5nJ6sgz64xHtoh6_ec7ReV5YpsOGhMtUJnwoRfEI,3549
 sqlmesh/utils/windows.py,sha256=r7Yin4WEwbvH7sKA3hOJtRsclmscNuBeCklCfVPvtQs,611
 sqlmesh/utils/yaml.py,sha256=hvo5lmybYLcIx4p6CQwufS8X4MbwS60lEBnCmYJCXmU,2336
-sqlmesh-0.209.1.dev17.dist-info/licenses/LICENSE,sha256=OlMefUjgWJdULtf84BLW0AZZcY8DwdgQqb_1j2862j8,11346
+sqlmesh-0.209.1.dev18.dist-info/licenses/LICENSE,sha256=OlMefUjgWJdULtf84BLW0AZZcY8DwdgQqb_1j2862j8,11346
 sqlmesh_dbt/__init__.py,sha256=awYS5y5mz-1NUmx6i5h5NSTJ7tidRl9NC0FAnFWSF6U,350
 sqlmesh_dbt/cli.py,sha256=Ekn77WLMmyyr31aoDoWrZp5DS0rfiddy7pNohBK8G0s,2115
 sqlmesh_dbt/console.py,sha256=1B0B5_R1DlrEkyWVDl8G78-B6_K-ZtLPsQEhqPCfLEU,188
@@ -413,8 +413,8 @@ web/server/api/endpoints/models.py,sha256=kwj0s7uve3iZSMfmjkoPVMFMeY1sD0peTeyrWf
 web/server/api/endpoints/modules.py,sha256=8hqqgonGay_mJmpCw0IdbjsPhWlQH2VLdKAqha-myac,468
 web/server/api/endpoints/plan.py,sha256=bbbY50W_2MsZSTxOHWMKz0tbIm75nsRSlPy8GI2fg9Q,9306
 web/server/api/endpoints/table_diff.py,sha256=SCy3-pr5x6bEEiOaZGCDbiHQhpgxwmnz3Acqa2Ie6jI,2090
-sqlmesh-0.209.1.dev17.dist-info/METADATA,sha256=yXmWoSCIcCeFPW6WEIG63FJ7vYd5f9GqngQu497tCEw,26686
-sqlmesh-0.209.1.dev17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sqlmesh-0.209.1.dev17.dist-info/entry_points.txt,sha256=sHAf6tQczIM8xZoduN4qaUjV7QEPVUUW_LCT8EDUMv4,155
-sqlmesh-0.209.1.dev17.dist-info/top_level.txt,sha256=RQ-33FPe2IgL0rgossAfJkCRtqslz9b7wFARqiWLC5Q,24
-sqlmesh-0.209.1.dev17.dist-info/RECORD,,
+sqlmesh-0.209.1.dev18.dist-info/METADATA,sha256=DhsqsAAfCLUo_qjBGTON1HbyNiXLjtm-eXvr1DlrUso,26686
+sqlmesh-0.209.1.dev18.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sqlmesh-0.209.1.dev18.dist-info/entry_points.txt,sha256=sHAf6tQczIM8xZoduN4qaUjV7QEPVUUW_LCT8EDUMv4,155
+sqlmesh-0.209.1.dev18.dist-info/top_level.txt,sha256=RQ-33FPe2IgL0rgossAfJkCRtqslz9b7wFARqiWLC5Q,24
+sqlmesh-0.209.1.dev18.dist-info/RECORD,,