sqlmesh 0.209.1.dev17__tar.gz → 0.209.1.dev19__tar.gz

{sqlmesh-0.209.1.dev17 → sqlmesh-0.209.1.dev19}/PKG-INFO

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

{sqlmesh-0.209.1.dev17 → sqlmesh-0.209.1.dev19}/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.dev19'
+__version_tuple__ = version_tuple = (0, 209, 1, 'dev19')

{sqlmesh-0.209.1.dev17 → sqlmesh-0.209.1.dev19}/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-0.209.1.dev17 → sqlmesh-0.209.1.dev19}/sqlmesh/core/plan/builder.py

@@ -155,6 +155,7 @@ class PlanBuilder:
 
         self._backfill_models = backfill_models
         self._end = end or default_end
+        self._default_start = default_start
         self._apply = apply
         self._console = console or get_console()
         self._choices: t.Dict[SnapshotId, SnapshotChangeCategory] = {}
@@ -802,6 +803,25 @@ class PlanBuilder:
                     f"Plan end date: '{time_like_to_str(end)}' cannot be in the future (execution time: '{time_like_to_str(self.execution_time)}')"
                 )
 
+        # Validate model-specific start/end dates
+        if (start := self.start or self._default_start) and (end := self.end):
+            start_ts = to_datetime(start)
+            end_ts = to_datetime(end)
+            if start_ts > end_ts:
+                models_to_check: t.Set[str] = (
+                    set(self._backfill_models or [])
+                    | set(self._context_diff.modified_snapshots.keys())
+                    | {s.name for s in self._context_diff.added}
+                    | set((self._end_override_per_model or {}).keys())
+                )
+                for model_name in models_to_check:
+                    if snapshot := self._model_fqn_to_snapshot.get(model_name):
+                        if snapshot.node.start is None or to_datetime(snapshot.node.start) > end_ts:
+                            raise PlanError(
+                                f"Model '{model_name}': Start date / time '({time_like_to_str(start_ts)})' can't be greater than end date / time '({time_like_to_str(end_ts)})'.\n"
+                                f"Set the `start` attribute in your project config model defaults to avoid this issue."
+                            )
+
     def _ensure_no_broken_references(self) -> None:
         for snapshot in self._context_diff.snapshots.values():
             broken_references = {

{sqlmesh-0.209.1.dev17 → sqlmesh-0.209.1.dev19}/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 → sqlmesh-0.209.1.dev19}/sqlmesh.egg-info/PKG-INFO

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

{sqlmesh-0.209.1.dev17 → sqlmesh-0.209.1.dev19}/tests/core/engine_adapter/test_bigquery.py

@@ -1072,3 +1072,81 @@ def test_get_alter_expressions_includes_catalog(
     assert schema.db == "bar"
     assert schema.sql(dialect="bigquery") == "catalog2.bar"
     assert tables == {"bing"}
+
+
+def test_job_cancellation_on_keyboard_interrupt_job_still_running(mocker: MockerFixture):
+    # Create a mock connection
+    connection_mock = mocker.NonCallableMock()
+    cursor_mock = mocker.Mock()
+    cursor_mock.connection = connection_mock
+    connection_mock.cursor.return_value = cursor_mock
+
+    # Mock the query job
+    mock_job = mocker.Mock()
+    mock_job.project = "test-project"
+    mock_job.location = "us-central1"
+    mock_job.job_id = "test-job-123"
+    mock_job.done.return_value = False  # Job is still running
+    mock_job.result.side_effect = KeyboardInterrupt()
+    mock_job._query_results = mocker.Mock()
+    mock_job._query_results.total_rows = 0
+    mock_job._query_results.schema = []
+
+    # Set up the client to return our mock job
+    connection_mock._client.query.return_value = mock_job
+
+    # Create adapter with the mocked connection
+    adapter = BigQueryEngineAdapter(lambda: connection_mock, job_retries=0)
+
+    # Execute a query and expect KeyboardInterrupt
+    with pytest.raises(KeyboardInterrupt):
+        adapter.execute("SELECT 1")
+
+    # Ensure the adapter's closed, so that the job can be aborted
+    adapter.close()
+
+    # Verify the job was created
+    connection_mock._client.query.assert_called_once()
+
+    # Verify job status was checked and cancellation was called
+    mock_job.done.assert_called_once()
+    mock_job.cancel.assert_called_once()
+
+
+def test_job_cancellation_on_keyboard_interrupt_job_already_done(mocker: MockerFixture):
+    # Create a mock connection
+    connection_mock = mocker.NonCallableMock()
+    cursor_mock = mocker.Mock()
+    cursor_mock.connection = connection_mock
+    connection_mock.cursor.return_value = cursor_mock
+
+    # Mock the query job
+    mock_job = mocker.Mock()
+    mock_job.project = "test-project"
+    mock_job.location = "us-central1"
+    mock_job.job_id = "test-job-456"
+    mock_job.done.return_value = True  # Job is already done
+    mock_job.result.side_effect = KeyboardInterrupt()
+    mock_job._query_results = mocker.Mock()
+    mock_job._query_results.total_rows = 0
+    mock_job._query_results.schema = []
+
+    # Set up the client to return our mock job
+    connection_mock._client.query.return_value = mock_job
+
+    # Create adapter with the mocked connection
+    adapter = BigQueryEngineAdapter(lambda: connection_mock, job_retries=0)
+
+    # Execute a query and expect KeyboardInterrupt
+    with pytest.raises(KeyboardInterrupt):
+        adapter.execute("SELECT 1")
+
+    # Ensure the adapter's closed, so that the job can be aborted
+    adapter.close()
+
+    # Verify the job was created
+    connection_mock._client.query.assert_called_once()
+
+    # Verify job status was checked but cancellation was NOT called
+    mock_job.done.assert_called_once()
+    mock_job.cancel.assert_not_called()

{sqlmesh-0.209.1.dev17 → sqlmesh-0.209.1.dev19}/tests/core/test_context.py

@@ -3007,3 +3007,59 @@ SELECT * FROM test_db.uppercase_gateway_table;
     assert len(uppercase_in_yaml_models) == 1, (
         f"External model with uppercase gateway in YAML should be found. Found {len(uppercase_in_yaml_models)} models"
     )
+
+
+def test_plan_no_start_configured():
+    context = Context(config=Config())
+    context.upsert_model(
+        load_sql_based_model(
+            parse(
+                """
+                MODEL(
+                    name db.xvg,
+                    kind INCREMENTAL_BY_TIME_RANGE (
+                        time_column ds
+                    ),
+                    cron '@daily'
+                );
+
+                SELECT id, ds FROM (VALUES
+                    ('1', '2020-01-01'),
+                ) data(id, ds)
+                WHERE ds BETWEEN @start_ds AND @end_ds
+                """
+            )
+        )
+    )
+
+    prod_plan = context.plan(auto_apply=True)
+    assert len(prod_plan.new_snapshots) == 1
+
+    context.upsert_model(
+        load_sql_based_model(
+            parse(
+                """
+                MODEL(
+                    name db.xvg,
+                    kind INCREMENTAL_BY_TIME_RANGE (
+                        time_column ds
+                    ),
+                    cron '@daily',
+                    physical_properties ('some_prop' = 1),
+                );
+
+                SELECT id, ds FROM (VALUES
+                    ('1', '2020-01-01'),
+                ) data(id, ds)
+                WHERE ds BETWEEN @start_ds AND @end_ds
+                """
+            )
+        )
+    )
+
+    # This should raise an error because the model has no start configured and the end time is less than the start time which will be calculated from the intervals
+    with pytest.raises(
+        PlanError,
+        match=r"Model '.*xvg.*': Start date / time .* can't be greater than end date / time .*\.\nSet the `start` attribute in your project config model defaults to avoid this issue",
+    ):
+        context.plan("dev", execution_time="1999-01-05")