Skip to content

Migrate BigQueryInsertJobTrigger to use on_kill() for user-initiated kills #66704

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shahar1 merged 1 commit into from
May 12, 2026
Merged

Migrate BigQueryInsertJobTrigger to use on_kill() for user-initiated kills #66704

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 68 additions & 53 deletions providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,14 +26,14 @@

from airflow.providers.common.compat.sdk import AirflowException
from airflow.providers.google.cloud.hooks.bigquery import BigQueryAsyncHook, BigQueryTableAsyncHook
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_3_PLUS
from airflow.triggers.base import BaseTrigger, TriggerEvent
from airflow.utils.state import TaskInstanceState

if TYPE_CHECKING:
from sqlalchemy.orm.session import Session

if not AIRFLOW_V_3_0_PLUS:
if not AIRFLOW_V_3_3_PLUS:
from sqlalchemy import select

from airflow.models.taskinstance import TaskInstance
Expand Down Expand Up @@ -103,7 +103,20 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
},
)

if not AIRFLOW_V_3_0_PLUS:
async def on_kill(self) -> None:
"""Cancel the BigQuery job when the task is killed by a user action."""
if self.job_id and self.cancel_on_kill:
self.log.info(
"Cancelling BigQuery job. Project ID: %s, Location: %s, Job ID: %s",
self.project_id,
self.location,
self.job_id,
)
hook = self._get_async_hook()
await hook.cancel_job(job_id=self.job_id, project_id=self.project_id, location=self.location)
self.log.info("BigQuery job %s cancelled successfully.", self.job_id)

if not AIRFLOW_V_3_3_PLUS:

@provide_session
def get_task_instance(self, session: Session) -> TaskInstance:
Expand All @@ -125,41 +138,41 @@ def get_task_instance(self, session: Session) -> TaskInstance:
)
return task_instance

async def get_task_state(self):
from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
async def get_task_state(self):
from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance

task_states_response = await sync_to_async(RuntimeTaskInstance.get_task_states)(
dag_id=self.task_instance.dag_id,
task_ids=[self.task_instance.task_id],
run_ids=[self.task_instance.run_id],
map_index=self.task_instance.map_index,
)
try:
task_state = task_states_response[self.task_instance.run_id][self.task_instance.task_id]
except Exception:
raise AirflowException(
"TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
self.task_instance.dag_id,
self.task_instance.task_id,
self.task_instance.run_id,
self.task_instance.map_index,
task_states_response = await sync_to_async(RuntimeTaskInstance.get_task_states)(
dag_id=self.task_instance.dag_id,
task_ids=[self.task_instance.task_id],
run_ids=[self.task_instance.run_id],
map_index=self.task_instance.map_index,
)
return task_state
try:
task_state = task_states_response[self.task_instance.run_id][self.task_instance.task_id]
except Exception:
raise AirflowException(
"TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
self.task_instance.dag_id,
self.task_instance.task_id,
self.task_instance.run_id,
self.task_instance.map_index,
)
return task_state

async def safe_to_cancel(self) -> bool:
"""
Whether it is safe to cancel the external job which is being executed by this trigger.
async def safe_to_cancel(self) -> bool:
"""
Whether it is safe to cancel the external job which is being executed by this trigger.

This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
Because in those cases, we should NOT cancel the external job.
"""
if AIRFLOW_V_3_0_PLUS:
task_state = await self.get_task_state()
else:
# Database query is needed to get the latest state of the task instance.
task_instance = self.get_task_instance() # type: ignore[call-arg]
task_state = task_instance.state
return task_state != TaskInstanceState.DEFERRED
This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
Because in those cases, we should NOT cancel the external job.
"""
if AIRFLOW_V_3_0_PLUS:
task_state = await self.get_task_state()
else:
# Database query is needed to get the latest state of the task instance.
task_instance = self.get_task_instance() # type: ignore[call-arg]
task_state = task_instance.state
return task_state != TaskInstanceState.DEFERRED

async def run(self) -> AsyncIterator[TriggerEvent]:
"""Get current job execution status and yields a TriggerEvent."""
Expand Down Expand Up @@ -196,25 +209,27 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
)
await asyncio.sleep(self.poll_interval)
except asyncio.CancelledError:
if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
self.log.info(
"The job is safe to cancel the as airflow TaskInstance is not in deferred state."
)
self.log.info(
"Cancelling job. Project ID: %s, Location: %s, Job ID: %s",
self.project_id,
self.location,
self.job_id,
)
await hook.cancel_job(job_id=self.job_id, project_id=self.project_id, location=self.location)
else:
self.log.info(
"Trigger may have shutdown. Skipping to cancel job because the airflow "
"task is not cancelled yet: Project ID: %s, Location:%s, Job ID:%s",
self.project_id,
self.location,
self.job_id,
)
# Legacy path for Airflow < 3.3.0
# On Airflow 3.3.0+, on_kill() handles user-initiated kills
if not AIRFLOW_V_3_3_PLUS:
if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
self.log.info(
"Cancelling job (legacy path). Project ID: %s, Location: %s, Job ID: %s",
self.project_id,
self.location,
self.job_id,
)
await hook.cancel_job(
job_id=self.job_id, project_id=self.project_id, location=self.location
)
else:
self.log.info(
"Trigger may have shutdown. Skipping to cancel job because the airflow "
"task is not cancelled yet: Project ID: %s, Location:%s, Job ID:%s",
self.project_id,
self.location,
self.job_id,
)
raise
except Exception as e:
self.log.exception("Exception occurred while checking for query completion")
Expand Down
117 changes: 64 additions & 53 deletions providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@
)
from airflow.triggers.base import TriggerEvent

from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS

TEST_CONN_ID = "bq_default"
TEST_JOB_ID = "1234"
TEST_GCP_PROJECT_ID = "test-project"
Expand Down Expand Up @@ -234,72 +236,81 @@ async def test_bigquery_op_trigger_exception(self, mock_job_status, caplog, inse
assert TriggerEvent({"status": "error", "message": "Test exception"}) == actual

@pytest.mark.asyncio
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.cancel_job")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
@pytest.mark.skipif(AIRFLOW_V_3_3_PLUS, reason="on_kill() handles cancellation for Airflow 3.3.0+")
@pytest.mark.parametrize("is_safe_to_cancel", [True, False])
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.BigQueryInsertJobTrigger._get_async_hook")
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.BigQueryInsertJobTrigger.safe_to_cancel")
async def test_bigquery_insert_job_trigger_cancellation(
self, mock_get_task_instance, mock_get_job_status, mock_cancel_job, caplog, insert_job_trigger
async def test_insert_job_trigger_run_cancelled(
self, mock_safe_to_cancel, mock_get_async_hook, insert_job_trigger, is_safe_to_cancel
):
"""
Test that BigQueryInsertJobTrigger handles cancellation correctly, logs the appropriate message,
and conditionally cancels the job based on the `cancel_on_kill` attribute.
"""
mock_get_task_instance.return_value = True
"""Test CancelledError handling for Airflow < 3.3.0."""
mock_safe_to_cancel.return_value = is_safe_to_cancel
mock_hook = mock_get_async_hook.return_value
mock_hook.get_job_status = AsyncMock()
mock_hook.get_job_status.side_effect = asyncio.CancelledError
mock_hook.cancel_job = AsyncMock()

async_gen = insert_job_trigger.run()
try:
await async_gen.asend(None)
except (asyncio.CancelledError, StopAsyncIteration):
pass
except Exception as e:
pytest.fail(f"Unexpected exception raised: {e}")

if insert_job_trigger.cancel_on_kill and is_safe_to_cancel:
mock_hook.cancel_job.assert_awaited_once_with(
job_id=insert_job_trigger.job_id,
project_id=insert_job_trigger.project_id,
location=insert_job_trigger.location,
)
else:
mock_hook.cancel_job.assert_not_awaited()

await async_gen.aclose()

@pytest.mark.asyncio
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.BigQueryInsertJobTrigger._get_async_hook")
async def test_on_kill_cancels_job(self, mock_get_async_hook, insert_job_trigger):
"""Test that on_kill cancels the BigQuery job."""
mock_hook = mock_get_async_hook.return_value
mock_hook.cancel_job = AsyncMock()
insert_job_trigger.job_id = TEST_JOB_ID
insert_job_trigger.cancel_on_kill = True
insert_job_trigger.job_id = "1234"

mock_get_job_status.side_effect = [
{"status": "running", "message": "Job is still running"},
asyncio.CancelledError(),
]
await insert_job_trigger.on_kill()

mock_cancel_job.return_value = asyncio.Future()
mock_cancel_job.return_value.set_result(None)
mock_hook.cancel_job.assert_awaited_once_with(
job_id=TEST_JOB_ID,
project_id=TEST_GCP_PROJECT_ID,
location=TEST_LOCATION,
)

caplog.set_level(logging.INFO)
@pytest.mark.asyncio
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.BigQueryInsertJobTrigger._get_async_hook")
async def test_on_kill_respects_cancel_on_kill_false(self, mock_get_async_hook, insert_job_trigger):
"""Test that on_kill does not cancel the job when cancel_on_kill is False."""
mock_hook = mock_get_async_hook.return_value
mock_hook.cancel_job = AsyncMock()
insert_job_trigger.job_id = TEST_JOB_ID
insert_job_trigger.cancel_on_kill = False

with pytest.raises(asyncio.CancelledError):
async for _ in insert_job_trigger.run():
pass
await insert_job_trigger.on_kill()

assert (
"Task was killed" in caplog.text
or "Bigquery job status is running. Sleeping for 4.0 seconds." in caplog.text
), "Expected messages about task status or cancellation not found in log."
mock_cancel_job.assert_awaited_once()
mock_hook.cancel_job.assert_not_called()

@pytest.mark.asyncio
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.cancel_job")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.BigQueryInsertJobTrigger.safe_to_cancel")
async def test_bigquery_insert_job_trigger_cancellation_unsafe_cancellation(
self, mock_safe_to_cancel, mock_get_job_status, mock_cancel_job, caplog, insert_job_trigger
):
"""
Test that BigQueryInsertJobTrigger logs the appropriate message and does not cancel the job
if safe_to_cancel returns False even when the task is cancelled.
"""
mock_safe_to_cancel.return_value = False
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.BigQueryInsertJobTrigger._get_async_hook")
async def test_on_kill_no_job_id_does_not_cancel(self, mock_get_async_hook, insert_job_trigger):
"""Test that on_kill does not attempt to cancel when job_id is not set."""
mock_hook = mock_get_async_hook.return_value
mock_hook.cancel_job = AsyncMock()
insert_job_trigger.job_id = None
insert_job_trigger.cancel_on_kill = True
insert_job_trigger.job_id = "1234"

# Simulate the initial job status as running
mock_get_job_status.side_effect = [
{"status": "running", "message": "Job is still running"},
asyncio.CancelledError(),
{"status": "running", "message": "Job is still running after cancellation"},
]

caplog.set_level(logging.INFO)
await insert_job_trigger.on_kill()

with pytest.raises(asyncio.CancelledError):
async for _ in insert_job_trigger.run():
pass

assert "Skipping to cancel job" in caplog.text, (
"Expected message about skipping cancellation not found in log."
)
assert mock_get_job_status.call_count == 2, "Job status should be checked multiple times"
mock_hook.cancel_job.assert_not_called()


class TestBigQueryGetDataTrigger:
Expand Down
Loading