shahar1 commented on code in PR #74564:
URL: https://github.com/apache/airflow/pull/74564#discussion_r4302845632
##########
providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py:
##########
@@ -296,6 +303,14 @@ def execute_complete(self, context: Context, event:
dict[str, Any]):
return self.destination_cloud_storage_uris
+ def on_kill(self) -> None:
+ if self.job_id and self.cancel_on_kill:
Review Comment:
**nit** — `self.job_id` can be truthy while `self.hook` is still `None`:
`job_id` is a constructor argument, so a kill that lands before `execute`
assigns `self.hook`, or one during `execute_complete` after a deferral resumes
in a fresh process, raises `AttributeError` here. The `# type:
ignore[union-attr]` is suppressing exactly that. Checking the hook closes that,
and the `# type: ignore[union-attr]` on the next line can then go:
```suggestion
if self.hook and self.job_id and self.cancel_on_kill:
```
(`BigQueryInsertJobOperator.on_kill` has the same shape, so this is
optional.)
##########
providers/google/tests/unit/google/cloud/transfers/test_bigquery_to_gcs.py:
##########
@@ -191,6 +191,48 @@ def test_execute_deferrable_mode(self, mock_hook):
nowait=True,
)
+
@mock.patch("airflow.providers.google.cloud.transfers.bigquery_to_gcs.BigQueryHook")
+ def test_execute_deferrable_mode_passes_cancel_on_kill_to_trigger(self,
mock_hook):
+ mock_hook.return_value.split_tablename.return_value = (PROJECT_ID,
TEST_DATASET, TEST_TABLE_ID)
+ mock_hook.return_value.insert_job.return_value =
MagicMock(job_id="real_job_id", error_result=False)
+
+ operator = BigQueryToGCSOperator(
+ task_id=TASK_ID,
+
source_project_dataset_table=f"{PROJECT_ID}:{TEST_DATASET}.{TEST_TABLE_ID}",
+ destination_cloud_storage_uris=["gs://some-bucket/some-file.txt"],
+ project_id=JOB_PROJECT_ID,
+ deferrable=True,
+ cancel_on_kill=False,
+ )
+ with pytest.raises(TaskDeferred) as exc:
+ operator.execute(context=mock.MagicMock())
+
+ assert exc.value.trigger.cancel_on_kill is False
+
+
@mock.patch("airflow.providers.google.cloud.transfers.bigquery_to_gcs.BigQueryHook")
+ def test_on_kill(self, mock_hook):
Review Comment:
**nit** — AGENTS.md (Testing Standards): "Use `@pytest.mark.parametrize` for
multiple similar inputs — consolidate tests that only differ in input/expected
values into a single parametrized test."
This test flips `operator.cancel_on_kill` between two `on_kill()` calls. A
`@pytest.mark.parametrize("cancel_on_kill", [True, False])` keeps each case
independent and makes a failure point at the right one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
