Cancel the Dataflow job when a user kills the deferred task

providers/google/src/airflow/providers/google/cloud/operators/dataflow.py

@@ -298,6 +298,9 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
  https://cloud.google.com/dataflow/docs/templates/executing-templates
 
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  """
 
  template_fields: Sequence[str] = (
@@ -334,11 +337,13 @@ def __init__(
  append_job_name: bool = True,
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
  expected_terminal_state: str | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
 
  self.template = template
+ self.cancel_on_kill = cancel_on_kill
  self.job_name = job_name
  self.options = options or {}
  self.dataflow_default_options = dataflow_default_options or {}
@@ -437,6 +442,7 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -453,7 +459,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
  return job_id
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.log.info("Cancelling job %s", self.job_name)
  self.hook.cancel_job(
@@ -525,6 +534,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: The expected final status of the operator on which the corresponding
  Airflow task succeeds. When not specified, it will be determined by the hook.
  :param append_job_name: True if unique suffix has to be appended to job name.
@@ -550,11 +562,13 @@ def __init__(
  append_job_name: bool = True,
  expected_terminal_state: str | None = None,
  poll_sleep: int = 10,
+ cancel_on_kill: bool = True,
  *args,
  **kwargs,
  ) -> None:
  super().__init__(*args, **kwargs)
  self.body = body
+ self.cancel_on_kill = cancel_on_kill
  self.location = location
  self.project_id = project_id
  self.gcp_conn_id = gcp_conn_id
@@ -633,6 +647,8 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -658,7 +674,10 @@ def execute_complete(self, context: Context, event: dict) -> dict[str, str]:
  return job
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.hook.cancel_job(
  job_id=self.job.get("id"),
@@ -689,10 +708,13 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
  :param append_job_name: Optional. Set to True if a unique suffix has to be appended to the `job_name`.
  Defaults to True.
  :param drain_pipeline: Optional. Set to True if you want to stop a streaming pipeline job by draining it
- instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs
-or in the deferrable mode. Defaults to False.
+ instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs.
+ Defaults to False.
  For more info see: https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
  :param deferrable: Optional. Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: Optional. The expected terminal state of the Dataflow job at which the
  operator task is set to succeed. Defaults to 'JOB_STATE_DONE' for the batch jobs and 'JOB_STATE_RUNNING'
  for the streaming jobs.
@@ -748,10 +770,12 @@ def __init__(
  jinja_variables: dict[str, str] | None = None,
  options: dict[str, Any] | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
  self.job_name = job_name
+ self.cancel_on_kill = cancel_on_kill
  self.yaml_pipeline_file = yaml_pipeline_file
  self.region = region
  self.project_id = project_id
@@ -793,6 +817,8 @@ def execute(self, context: Context) -> dict[str, Any]:
  cancel_timeout=self.cancel_timeout,
  expected_terminal_state=self.expected_terminal_state,
  impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -818,10 +844,13 @@ def on_kill(self):
  """
  Cancel the dataflow job if a task instance gets killed.
 
- This method will not be called if a task instance is killed in a deferred
- state.
+ This method is not called for a task instance killed in a deferred state;
+ in that case the trigger cancels the job instead, honoring cancel_on_kill
+ and drain_pipeline.
  """
  self.log.info("On kill called.")
+ if not self.cancel_on_kill:
+ return
  if self.job_id:
  self.hook.cancel_job(
  job_id=self.job_id,

providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py

@@ -22,6 +22,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from asgiref.sync import sync_to_async
 from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.dataflow_v1beta3 import JobState
 from google.cloud.dataflow_v1beta3.types import (
@@ -33,7 +34,11 @@
  MetricUpdate,
 )
 
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+ AsyncDataflowHook,
+ DataflowHook,
+ DataflowJobStatus,
+)
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -62,6 +67,10 @@ class TemplateJobStartTrigger(BaseTrigger):
  account from the list granting this role to the originating account (templated).
  :param cancel_timeout: Optional. How long (in seconds) operator should wait for the pipeline to be
  successfully cancelled when task is being killed.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark-failed, clear or mark-succeeded).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling, matching the operator's behaviour before deferral.
  """
 
  def __init__(
@@ -73,6 +82,8 @@ def __init__(
  poll_sleep: int = 10,
  impersonation_chain: str | Sequence[str] | None = None,
  cancel_timeout: int | None = 5 * 60,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -82,6 +93,8 @@ def __init__(
  self.poll_sleep = poll_sleep
  self.impersonation_chain = impersonation_chain
  self.cancel_timeout = cancel_timeout
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -95,9 +108,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "poll_sleep": self.poll_sleep,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job status or yield certain Events.
@@ -300,6 +352,10 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
  If set as a sequence, the identities from the list must grant
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark failed, clear, or mark success).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling when the task is killed.
  """
 
  def __init__(
@@ -312,6 +368,8 @@ def __init__(
  cancel_timeout: int | None = 5 * 60,
  expected_terminal_state: str | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -322,6 +380,8 @@ def __init__(
  self.cancel_timeout = cancel_timeout
  self.expected_terminal_state = expected_terminal_state
  self.impersonation_chain = impersonation_chain
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -336,9 +396,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "expected_terminal_state": self.expected_terminal_state,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job and yield events depending on the job's type and state.

providers/google/src/airflow/providers/google/cloud/operators/dataflow.py

@@ -298,6 +298,9 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
  https://cloud.google.com/dataflow/docs/templates/executing-templates
 
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  """
 
  template_fields: Sequence[str] = (
@@ -334,11 +337,13 @@ def __init__(
  append_job_name: bool = True,
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
  expected_terminal_state: str | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
 
  self.template = template
+ self.cancel_on_kill = cancel_on_kill
  self.job_name = job_name
  self.options = options or {}
  self.dataflow_default_options = dataflow_default_options or {}
@@ -437,6 +442,7 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -453,7 +459,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
  return job_id
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.log.info("Cancelling job %s", self.job_name)
  self.hook.cancel_job(
@@ -525,6 +534,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: The expected final status of the operator on which the corresponding
  Airflow task succeeds. When not specified, it will be determined by the hook.
  :param append_job_name: True if unique suffix has to be appended to job name.
@@ -550,11 +562,13 @@ def __init__(
  append_job_name: bool = True,
  expected_terminal_state: str | None = None,
  poll_sleep: int = 10,
+ cancel_on_kill: bool = True,
  *args,
  **kwargs,
  ) -> None:
  super().__init__(*args, **kwargs)
  self.body = body
+ self.cancel_on_kill = cancel_on_kill
  self.location = location
  self.project_id = project_id
  self.gcp_conn_id = gcp_conn_id
@@ -633,6 +647,8 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -658,7 +674,10 @@ def execute_complete(self, context: Context, event: dict) -> dict[str, str]:
  return job
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.hook.cancel_job(
  job_id=self.job.get("id"),
@@ -689,10 +708,13 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
  :param append_job_name: Optional. Set to True if a unique suffix has to be appended to the `job_name`.
  Defaults to True.
  :param drain_pipeline: Optional. Set to True if you want to stop a streaming pipeline job by draining it
- instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs
-or in the deferrable mode. Defaults to False.
+ instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs.
+ Defaults to False.
  For more info see: https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
  :param deferrable: Optional. Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: Optional. The expected terminal state of the Dataflow job at which the
  operator task is set to succeed. Defaults to 'JOB_STATE_DONE' for the batch jobs and 'JOB_STATE_RUNNING'
  for the streaming jobs.
@@ -748,10 +770,12 @@ def __init__(
  jinja_variables: dict[str, str] | None = None,
  options: dict[str, Any] | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
  self.job_name = job_name
+ self.cancel_on_kill = cancel_on_kill
  self.yaml_pipeline_file = yaml_pipeline_file
  self.region = region
  self.project_id = project_id
@@ -793,6 +817,8 @@ def execute(self, context: Context) -> dict[str, Any]:
  cancel_timeout=self.cancel_timeout,
  expected_terminal_state=self.expected_terminal_state,
  impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -818,10 +844,13 @@ def on_kill(self):
  """
  Cancel the dataflow job if a task instance gets killed.
 
- This method will not be called if a task instance is killed in a deferred
- state.
+ This method is not called for a task instance killed in a deferred state;
+ in that case the trigger cancels the job instead, honoring cancel_on_kill
+ and drain_pipeline.
  """
  self.log.info("On kill called.")
+ if not self.cancel_on_kill:
+ return
  if self.job_id:
  self.hook.cancel_job(
  job_id=self.job_id,

providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py

@@ -22,6 +22,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from asgiref.sync import sync_to_async
 from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.dataflow_v1beta3 import JobState
 from google.cloud.dataflow_v1beta3.types import (
@@ -33,7 +34,11 @@
  MetricUpdate,
 )
 
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+ AsyncDataflowHook,
+ DataflowHook,
+ DataflowJobStatus,
+)
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -62,6 +67,10 @@ class TemplateJobStartTrigger(BaseTrigger):
  account from the list granting this role to the originating account (templated).
  :param cancel_timeout: Optional. How long (in seconds) operator should wait for the pipeline to be
  successfully cancelled when task is being killed.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark-failed, clear or mark-succeeded).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling, matching the operator's behaviour before deferral.
  """
 
  def __init__(
@@ -73,6 +82,8 @@ def __init__(
  poll_sleep: int = 10,
  impersonation_chain: str | Sequence[str] | None = None,
  cancel_timeout: int | None = 5 * 60,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -82,6 +93,8 @@ def __init__(
  self.poll_sleep = poll_sleep
  self.impersonation_chain = impersonation_chain
  self.cancel_timeout = cancel_timeout
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -95,9 +108,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "poll_sleep": self.poll_sleep,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job status or yield certain Events.
@@ -300,6 +352,10 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
  If set as a sequence, the identities from the list must grant
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark failed, clear, or mark success).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling when the task is killed.
  """
 
  def __init__(
@@ -312,6 +368,8 @@ def __init__(
  cancel_timeout: int | None = 5 * 60,
  expected_terminal_state: str | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -322,6 +380,8 @@ def __init__(
  self.cancel_timeout = cancel_timeout
  self.expected_terminal_state = expected_terminal_state
  self.impersonation_chain = impersonation_chain
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -336,9 +396,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "expected_terminal_state": self.expected_terminal_state,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job and yield events depending on the job's type and state.

providers/google/src/airflow/providers/google/cloud/operators/dataflow.py

@@ -298,6 +298,9 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
  https://cloud.google.com/dataflow/docs/templates/executing-templates
 
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  """
 
  template_fields: Sequence[str] = (
@@ -334,11 +337,13 @@ def __init__(
  append_job_name: bool = True,
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
  expected_terminal_state: str | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
 
  self.template = template
+ self.cancel_on_kill = cancel_on_kill
  self.job_name = job_name
  self.options = options or {}
  self.dataflow_default_options = dataflow_default_options or {}
@@ -437,6 +442,7 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -453,7 +459,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
  return job_id
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.log.info("Cancelling job %s", self.job_name)
  self.hook.cancel_job(
@@ -525,6 +534,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: The expected final status of the operator on which the corresponding
  Airflow task succeeds. When not specified, it will be determined by the hook.
  :param append_job_name: True if unique suffix has to be appended to job name.
@@ -550,11 +562,13 @@ def __init__(
  append_job_name: bool = True,
  expected_terminal_state: str | None = None,
  poll_sleep: int = 10,
+ cancel_on_kill: bool = True,
  *args,
  **kwargs,
  ) -> None:
  super().__init__(*args, **kwargs)
  self.body = body
+ self.cancel_on_kill = cancel_on_kill
  self.location = location
  self.project_id = project_id
  self.gcp_conn_id = gcp_conn_id
@@ -633,6 +647,8 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -658,7 +674,10 @@ def execute_complete(self, context: Context, event: dict) -> dict[str, str]:
  return job
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.hook.cancel_job(
  job_id=self.job.get("id"),
@@ -689,10 +708,13 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
  :param append_job_name: Optional. Set to True if a unique suffix has to be appended to the `job_name`.
  Defaults to True.
  :param drain_pipeline: Optional. Set to True if you want to stop a streaming pipeline job by draining it
- instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs
-or in the deferrable mode. Defaults to False.
+ instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs.
+ Defaults to False.
  For more info see: https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
  :param deferrable: Optional. Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: Optional. The expected terminal state of the Dataflow job at which the
  operator task is set to succeed. Defaults to 'JOB_STATE_DONE' for the batch jobs and 'JOB_STATE_RUNNING'
  for the streaming jobs.
@@ -748,10 +770,12 @@ def __init__(
  jinja_variables: dict[str, str] | None = None,
  options: dict[str, Any] | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
  self.job_name = job_name
+ self.cancel_on_kill = cancel_on_kill
  self.yaml_pipeline_file = yaml_pipeline_file
  self.region = region
  self.project_id = project_id
@@ -793,6 +817,8 @@ def execute(self, context: Context) -> dict[str, Any]:
  cancel_timeout=self.cancel_timeout,
  expected_terminal_state=self.expected_terminal_state,
  impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -818,10 +844,13 @@ def on_kill(self):
  """
  Cancel the dataflow job if a task instance gets killed.
 
- This method will not be called if a task instance is killed in a deferred
- state.
+ This method is not called for a task instance killed in a deferred state;
+ in that case the trigger cancels the job instead, honoring cancel_on_kill
+ and drain_pipeline.
  """
  self.log.info("On kill called.")
+ if not self.cancel_on_kill:
+ return
  if self.job_id:
  self.hook.cancel_job(
  job_id=self.job_id,

providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py

@@ -22,6 +22,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from asgiref.sync import sync_to_async
 from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.dataflow_v1beta3 import JobState
 from google.cloud.dataflow_v1beta3.types import (
@@ -33,7 +34,11 @@
  MetricUpdate,
 )
 
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+ AsyncDataflowHook,
+ DataflowHook,
+ DataflowJobStatus,
+)
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -62,6 +67,10 @@ class TemplateJobStartTrigger(BaseTrigger):
  account from the list granting this role to the originating account (templated).
  :param cancel_timeout: Optional. How long (in seconds) operator should wait for the pipeline to be
  successfully cancelled when task is being killed.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark-failed, clear or mark-succeeded).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling, matching the operator's behaviour before deferral.
  """
 
  def __init__(
@@ -73,6 +82,8 @@ def __init__(
  poll_sleep: int = 10,
  impersonation_chain: str | Sequence[str] | None = None,
  cancel_timeout: int | None = 5 * 60,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -82,6 +93,8 @@ def __init__(
  self.poll_sleep = poll_sleep
  self.impersonation_chain = impersonation_chain
  self.cancel_timeout = cancel_timeout
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -95,9 +108,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "poll_sleep": self.poll_sleep,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job status or yield certain Events.
@@ -300,6 +352,10 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
  If set as a sequence, the identities from the list must grant
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark failed, clear, or mark success).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling when the task is killed.
  """
 
  def __init__(
@@ -312,6 +368,8 @@ def __init__(
  cancel_timeout: int | None = 5 * 60,
  expected_terminal_state: str | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -322,6 +380,8 @@ def __init__(
  self.cancel_timeout = cancel_timeout
  self.expected_terminal_state = expected_terminal_state
  self.impersonation_chain = impersonation_chain
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -336,9 +396,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "expected_terminal_state": self.expected_terminal_state,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job and yield events depending on the job's type and state.

providers/google/src/airflow/providers/google/cloud/operators/dataflow.py

@@ -298,6 +298,9 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
  https://cloud.google.com/dataflow/docs/templates/executing-templates
 
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  """
 
  template_fields: Sequence[str] = (
@@ -334,11 +337,13 @@ def __init__(
  append_job_name: bool = True,
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
  expected_terminal_state: str | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
 
  self.template = template
+ self.cancel_on_kill = cancel_on_kill
  self.job_name = job_name
  self.options = options or {}
  self.dataflow_default_options = dataflow_default_options or {}
@@ -437,6 +442,7 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -453,7 +459,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
  return job_id
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.log.info("Cancelling job %s", self.job_name)
  self.hook.cancel_job(
@@ -525,6 +534,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: The expected final status of the operator on which the corresponding
  Airflow task succeeds. When not specified, it will be determined by the hook.
  :param append_job_name: True if unique suffix has to be appended to job name.
@@ -550,11 +562,13 @@ def __init__(
  append_job_name: bool = True,
  expected_terminal_state: str | None = None,
  poll_sleep: int = 10,
+ cancel_on_kill: bool = True,
  *args,
  **kwargs,
  ) -> None:
  super().__init__(*args, **kwargs)
  self.body = body
+ self.cancel_on_kill = cancel_on_kill
  self.location = location
  self.project_id = project_id
  self.gcp_conn_id = gcp_conn_id
@@ -633,6 +647,8 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -658,7 +674,10 @@ def execute_complete(self, context: Context, event: dict) -> dict[str, str]:
  return job
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.hook.cancel_job(
  job_id=self.job.get("id"),
@@ -689,10 +708,13 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
  :param append_job_name: Optional. Set to True if a unique suffix has to be appended to the `job_name`.
  Defaults to True.
  :param drain_pipeline: Optional. Set to True if you want to stop a streaming pipeline job by draining it
- instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs
-or in the deferrable mode. Defaults to False.
+ instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs.
+ Defaults to False.
  For more info see: https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
  :param deferrable: Optional. Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: Optional. The expected terminal state of the Dataflow job at which the
  operator task is set to succeed. Defaults to 'JOB_STATE_DONE' for the batch jobs and 'JOB_STATE_RUNNING'
  for the streaming jobs.
@@ -748,10 +770,12 @@ def __init__(
  jinja_variables: dict[str, str] | None = None,
  options: dict[str, Any] | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
  self.job_name = job_name
+ self.cancel_on_kill = cancel_on_kill
  self.yaml_pipeline_file = yaml_pipeline_file
  self.region = region
  self.project_id = project_id
@@ -793,6 +817,8 @@ def execute(self, context: Context) -> dict[str, Any]:
  cancel_timeout=self.cancel_timeout,
  expected_terminal_state=self.expected_terminal_state,
  impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -818,10 +844,13 @@ def on_kill(self):
  """
  Cancel the dataflow job if a task instance gets killed.
 
- This method will not be called if a task instance is killed in a deferred
- state.
+ This method is not called for a task instance killed in a deferred state;
+ in that case the trigger cancels the job instead, honoring cancel_on_kill
+ and drain_pipeline.
  """
  self.log.info("On kill called.")
+ if not self.cancel_on_kill:
+ return
  if self.job_id:
  self.hook.cancel_job(
  job_id=self.job_id,

providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py

@@ -22,6 +22,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from asgiref.sync import sync_to_async
 from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.dataflow_v1beta3 import JobState
 from google.cloud.dataflow_v1beta3.types import (
@@ -33,7 +34,11 @@
  MetricUpdate,
 )
 
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+ AsyncDataflowHook,
+ DataflowHook,
+ DataflowJobStatus,
+)
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -62,6 +67,10 @@ class TemplateJobStartTrigger(BaseTrigger):
  account from the list granting this role to the originating account (templated).
  :param cancel_timeout: Optional. How long (in seconds) operator should wait for the pipeline to be
  successfully cancelled when task is being killed.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark-failed, clear or mark-succeeded).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling, matching the operator's behaviour before deferral.
  """
 
  def __init__(
@@ -73,6 +82,8 @@ def __init__(
  poll_sleep: int = 10,
  impersonation_chain: str | Sequence[str] | None = None,
  cancel_timeout: int | None = 5 * 60,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -82,6 +93,8 @@ def __init__(
  self.poll_sleep = poll_sleep
  self.impersonation_chain = impersonation_chain
  self.cancel_timeout = cancel_timeout
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -95,9 +108,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "poll_sleep": self.poll_sleep,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job status or yield certain Events.
@@ -300,6 +352,10 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
  If set as a sequence, the identities from the list must grant
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark failed, clear, or mark success).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling when the task is killed.
  """
 
  def __init__(
@@ -312,6 +368,8 @@ def __init__(
  cancel_timeout: int | None = 5 * 60,
  expected_terminal_state: str | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -322,6 +380,8 @@ def __init__(
  self.cancel_timeout = cancel_timeout
  self.expected_terminal_state = expected_terminal_state
  self.impersonation_chain = impersonation_chain
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -336,9 +396,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "expected_terminal_state": self.expected_terminal_state,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job and yield events depending on the job's type and state.

providers/google/src/airflow/providers/google/cloud/operators/dataflow.py

@@ -298,6 +298,9 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
  https://cloud.google.com/dataflow/docs/templates/executing-templates
 
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  """
 
  template_fields: Sequence[str] = (
@@ -334,11 +337,13 @@ def __init__(
  append_job_name: bool = True,
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
  expected_terminal_state: str | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
 
  self.template = template
+ self.cancel_on_kill = cancel_on_kill
  self.job_name = job_name
  self.options = options or {}
  self.dataflow_default_options = dataflow_default_options or {}
@@ -437,6 +442,7 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -453,7 +459,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
  return job_id
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.log.info("Cancelling job %s", self.job_name)
  self.hook.cancel_job(
@@ -525,6 +534,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: The expected final status of the operator on which the corresponding
  Airflow task succeeds. When not specified, it will be determined by the hook.
  :param append_job_name: True if unique suffix has to be appended to job name.
@@ -550,11 +562,13 @@ def __init__(
  append_job_name: bool = True,
  expected_terminal_state: str | None = None,
  poll_sleep: int = 10,
+ cancel_on_kill: bool = True,
  *args,
  **kwargs,
  ) -> None:
  super().__init__(*args, **kwargs)
  self.body = body
+ self.cancel_on_kill = cancel_on_kill
  self.location = location
  self.project_id = project_id
  self.gcp_conn_id = gcp_conn_id
@@ -633,6 +647,8 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -658,7 +674,10 @@ def execute_complete(self, context: Context, event: dict) -> dict[str, str]:
  return job
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.hook.cancel_job(
  job_id=self.job.get("id"),
@@ -689,10 +708,13 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
  :param append_job_name: Optional. Set to True if a unique suffix has to be appended to the `job_name`.
  Defaults to True.
  :param drain_pipeline: Optional. Set to True if you want to stop a streaming pipeline job by draining it
- instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs
-or in the deferrable mode. Defaults to False.
+ instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs.
+ Defaults to False.
  For more info see: https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
  :param deferrable: Optional. Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: Optional. The expected terminal state of the Dataflow job at which the
  operator task is set to succeed. Defaults to 'JOB_STATE_DONE' for the batch jobs and 'JOB_STATE_RUNNING'
  for the streaming jobs.
@@ -748,10 +770,12 @@ def __init__(
  jinja_variables: dict[str, str] | None = None,
  options: dict[str, Any] | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
  self.job_name = job_name
+ self.cancel_on_kill = cancel_on_kill
  self.yaml_pipeline_file = yaml_pipeline_file
  self.region = region
  self.project_id = project_id
@@ -793,6 +817,8 @@ def execute(self, context: Context) -> dict[str, Any]:
  cancel_timeout=self.cancel_timeout,
  expected_terminal_state=self.expected_terminal_state,
  impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -818,10 +844,13 @@ def on_kill(self):
  """
  Cancel the dataflow job if a task instance gets killed.
 
- This method will not be called if a task instance is killed in a deferred
- state.
+ This method is not called for a task instance killed in a deferred state;
+ in that case the trigger cancels the job instead, honoring cancel_on_kill
+ and drain_pipeline.
  """
  self.log.info("On kill called.")
+ if not self.cancel_on_kill:
+ return
  if self.job_id:
  self.hook.cancel_job(
  job_id=self.job_id,

providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py

@@ -22,6 +22,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from asgiref.sync import sync_to_async
 from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.dataflow_v1beta3 import JobState
 from google.cloud.dataflow_v1beta3.types import (
@@ -33,7 +34,11 @@
  MetricUpdate,
 )
 
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+ AsyncDataflowHook,
+ DataflowHook,
+ DataflowJobStatus,
+)
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -62,6 +67,10 @@ class TemplateJobStartTrigger(BaseTrigger):
  account from the list granting this role to the originating account (templated).
  :param cancel_timeout: Optional. How long (in seconds) operator should wait for the pipeline to be
  successfully cancelled when task is being killed.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark-failed, clear or mark-succeeded).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling, matching the operator's behaviour before deferral.
  """
 
  def __init__(
@@ -73,6 +82,8 @@ def __init__(
  poll_sleep: int = 10,
  impersonation_chain: str | Sequence[str] | None = None,
  cancel_timeout: int | None = 5 * 60,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -82,6 +93,8 @@ def __init__(
  self.poll_sleep = poll_sleep
  self.impersonation_chain = impersonation_chain
  self.cancel_timeout = cancel_timeout
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -95,9 +108,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "poll_sleep": self.poll_sleep,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job status or yield certain Events.
@@ -300,6 +352,10 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
  If set as a sequence, the identities from the list must grant
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark failed, clear, or mark success).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling when the task is killed.
  """
 
  def __init__(
@@ -312,6 +368,8 @@ def __init__(
  cancel_timeout: int | None = 5 * 60,
  expected_terminal_state: str | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -322,6 +380,8 @@ def __init__(
  self.cancel_timeout = cancel_timeout
  self.expected_terminal_state = expected_terminal_state
  self.impersonation_chain = impersonation_chain
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -336,9 +396,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "expected_terminal_state": self.expected_terminal_state,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job and yield events depending on the job's type and state.

providers/google/src/airflow/providers/google/cloud/operators/dataflow.py

@@ -298,6 +298,9 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
  https://cloud.google.com/dataflow/docs/templates/executing-templates
 
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  """
 
  template_fields: Sequence[str] = (
@@ -334,11 +337,13 @@ def __init__(
  append_job_name: bool = True,
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
  expected_terminal_state: str | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
 
  self.template = template
+ self.cancel_on_kill = cancel_on_kill
  self.job_name = job_name
  self.options = options or {}
  self.dataflow_default_options = dataflow_default_options or {}
@@ -437,6 +442,7 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -453,7 +459,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
  return job_id
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.log.info("Cancelling job %s", self.job_name)
  self.hook.cancel_job(
@@ -525,6 +534,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: The expected final status of the operator on which the corresponding
  Airflow task succeeds. When not specified, it will be determined by the hook.
  :param append_job_name: True if unique suffix has to be appended to job name.
@@ -550,11 +562,13 @@ def __init__(
  append_job_name: bool = True,
  expected_terminal_state: str | None = None,
  poll_sleep: int = 10,
+ cancel_on_kill: bool = True,
  *args,
  **kwargs,
  ) -> None:
  super().__init__(*args, **kwargs)
  self.body = body
+ self.cancel_on_kill = cancel_on_kill
  self.location = location
  self.project_id = project_id
  self.gcp_conn_id = gcp_conn_id
@@ -633,6 +647,8 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -658,7 +674,10 @@ def execute_complete(self, context: Context, event: dict) -> dict[str, str]:
  return job
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.hook.cancel_job(
  job_id=self.job.get("id"),
@@ -689,10 +708,13 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
  :param append_job_name: Optional. Set to True if a unique suffix has to be appended to the `job_name`.
  Defaults to True.
  :param drain_pipeline: Optional. Set to True if you want to stop a streaming pipeline job by draining it
- instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs
-or in the deferrable mode. Defaults to False.
+ instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs.
+ Defaults to False.
  For more info see: https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
  :param deferrable: Optional. Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: Optional. The expected terminal state of the Dataflow job at which the
  operator task is set to succeed. Defaults to 'JOB_STATE_DONE' for the batch jobs and 'JOB_STATE_RUNNING'
  for the streaming jobs.
@@ -748,10 +770,12 @@ def __init__(
  jinja_variables: dict[str, str] | None = None,
  options: dict[str, Any] | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
  self.job_name = job_name
+ self.cancel_on_kill = cancel_on_kill
  self.yaml_pipeline_file = yaml_pipeline_file
  self.region = region
  self.project_id = project_id
@@ -793,6 +817,8 @@ def execute(self, context: Context) -> dict[str, Any]:
  cancel_timeout=self.cancel_timeout,
  expected_terminal_state=self.expected_terminal_state,
  impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -818,10 +844,13 @@ def on_kill(self):
  """
  Cancel the dataflow job if a task instance gets killed.
 
- This method will not be called if a task instance is killed in a deferred
- state.
+ This method is not called for a task instance killed in a deferred state;
+ in that case the trigger cancels the job instead, honoring cancel_on_kill
+ and drain_pipeline.
  """
  self.log.info("On kill called.")
+ if not self.cancel_on_kill:
+ return
  if self.job_id:
  self.hook.cancel_job(
  job_id=self.job_id,

providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py

@@ -22,6 +22,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from asgiref.sync import sync_to_async
 from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.dataflow_v1beta3 import JobState
 from google.cloud.dataflow_v1beta3.types import (
@@ -33,7 +34,11 @@
  MetricUpdate,
 )
 
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+ AsyncDataflowHook,
+ DataflowHook,
+ DataflowJobStatus,
+)
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -62,6 +67,10 @@ class TemplateJobStartTrigger(BaseTrigger):
  account from the list granting this role to the originating account (templated).
  :param cancel_timeout: Optional. How long (in seconds) operator should wait for the pipeline to be
  successfully cancelled when task is being killed.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark-failed, clear or mark-succeeded).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling, matching the operator's behaviour before deferral.
  """
 
  def __init__(
@@ -73,6 +82,8 @@ def __init__(
  poll_sleep: int = 10,
  impersonation_chain: str | Sequence[str] | None = None,
  cancel_timeout: int | None = 5 * 60,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -82,6 +93,8 @@ def __init__(
  self.poll_sleep = poll_sleep
  self.impersonation_chain = impersonation_chain
  self.cancel_timeout = cancel_timeout
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -95,9 +108,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "poll_sleep": self.poll_sleep,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job status or yield certain Events.
@@ -300,6 +352,10 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
  If set as a sequence, the identities from the list must grant
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark failed, clear, or mark success).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling when the task is killed.
  """
 
  def __init__(
@@ -312,6 +368,8 @@ def __init__(
  cancel_timeout: int | None = 5 * 60,
  expected_terminal_state: str | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -322,6 +380,8 @@ def __init__(
  self.cancel_timeout = cancel_timeout
  self.expected_terminal_state = expected_terminal_state
  self.impersonation_chain = impersonation_chain
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -336,9 +396,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "expected_terminal_state": self.expected_terminal_state,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job and yield events depending on the job's type and state.

providers/google/src/airflow/providers/google/cloud/operators/dataflow.py

@@ -298,6 +298,9 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
  https://cloud.google.com/dataflow/docs/templates/executing-templates
 
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  """
 
  template_fields: Sequence[str] = (
@@ -334,11 +337,13 @@ def __init__(
  append_job_name: bool = True,
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
  expected_terminal_state: str | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
 
  self.template = template
+ self.cancel_on_kill = cancel_on_kill
  self.job_name = job_name
  self.options = options or {}
  self.dataflow_default_options = dataflow_default_options or {}
@@ -437,6 +442,7 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -453,7 +459,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
  return job_id
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.log.info("Cancelling job %s", self.job_name)
  self.hook.cancel_job(
@@ -525,6 +534,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: The expected final status of the operator on which the corresponding
  Airflow task succeeds. When not specified, it will be determined by the hook.
  :param append_job_name: True if unique suffix has to be appended to job name.
@@ -550,11 +562,13 @@ def __init__(
  append_job_name: bool = True,
  expected_terminal_state: str | None = None,
  poll_sleep: int = 10,
+ cancel_on_kill: bool = True,
  *args,
  **kwargs,
  ) -> None:
  super().__init__(*args, **kwargs)
  self.body = body
+ self.cancel_on_kill = cancel_on_kill
  self.location = location
  self.project_id = project_id
  self.gcp_conn_id = gcp_conn_id
@@ -633,6 +647,8 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -658,7 +674,10 @@ def execute_complete(self, context: Context, event: dict) -> dict[str, str]:
  return job
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.hook.cancel_job(
  job_id=self.job.get("id"),
@@ -689,10 +708,13 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
  :param append_job_name: Optional. Set to True if a unique suffix has to be appended to the `job_name`.
  Defaults to True.
  :param drain_pipeline: Optional. Set to True if you want to stop a streaming pipeline job by draining it
- instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs
-or in the deferrable mode. Defaults to False.
+ instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs.
+ Defaults to False.
  For more info see: https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
  :param deferrable: Optional. Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: Optional. The expected terminal state of the Dataflow job at which the
  operator task is set to succeed. Defaults to 'JOB_STATE_DONE' for the batch jobs and 'JOB_STATE_RUNNING'
  for the streaming jobs.
@@ -748,10 +770,12 @@ def __init__(
  jinja_variables: dict[str, str] | None = None,
  options: dict[str, Any] | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
  self.job_name = job_name
+ self.cancel_on_kill = cancel_on_kill
  self.yaml_pipeline_file = yaml_pipeline_file
  self.region = region
  self.project_id = project_id
@@ -793,6 +817,8 @@ def execute(self, context: Context) -> dict[str, Any]:
  cancel_timeout=self.cancel_timeout,
  expected_terminal_state=self.expected_terminal_state,
  impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -818,10 +844,13 @@ def on_kill(self):
  """
  Cancel the dataflow job if a task instance gets killed.
 
- This method will not be called if a task instance is killed in a deferred
- state.
+ This method is not called for a task instance killed in a deferred state;
+ in that case the trigger cancels the job instead, honoring cancel_on_kill
+ and drain_pipeline.
  """
  self.log.info("On kill called.")
+ if not self.cancel_on_kill:
+ return
  if self.job_id:
  self.hook.cancel_job(
  job_id=self.job_id,

providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py

@@ -22,6 +22,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from asgiref.sync import sync_to_async
 from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.dataflow_v1beta3 import JobState
 from google.cloud.dataflow_v1beta3.types import (
@@ -33,7 +34,11 @@
  MetricUpdate,
 )
 
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+ AsyncDataflowHook,
+ DataflowHook,
+ DataflowJobStatus,
+)
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -62,6 +67,10 @@ class TemplateJobStartTrigger(BaseTrigger):
  account from the list granting this role to the originating account (templated).
  :param cancel_timeout: Optional. How long (in seconds) operator should wait for the pipeline to be
  successfully cancelled when task is being killed.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark-failed, clear or mark-succeeded).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling, matching the operator's behaviour before deferral.
  """
 
  def __init__(
@@ -73,6 +82,8 @@ def __init__(
  poll_sleep: int = 10,
  impersonation_chain: str | Sequence[str] | None = None,
  cancel_timeout: int | None = 5 * 60,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -82,6 +93,8 @@ def __init__(
  self.poll_sleep = poll_sleep
  self.impersonation_chain = impersonation_chain
  self.cancel_timeout = cancel_timeout
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -95,9 +108,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "poll_sleep": self.poll_sleep,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job status or yield certain Events.
@@ -300,6 +352,10 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
  If set as a sequence, the identities from the list must grant
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark failed, clear, or mark success).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling when the task is killed.
  """
 
  def __init__(
@@ -312,6 +368,8 @@ def __init__(
  cancel_timeout: int | None = 5 * 60,
  expected_terminal_state: str | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -322,6 +380,8 @@ def __init__(
  self.cancel_timeout = cancel_timeout
  self.expected_terminal_state = expected_terminal_state
  self.impersonation_chain = impersonation_chain
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -336,9 +396,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "expected_terminal_state": self.expected_terminal_state,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job and yield events depending on the job's type and state.

providers/google/src/airflow/providers/google/cloud/operators/dataflow.py

@@ -298,6 +298,9 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
  https://cloud.google.com/dataflow/docs/templates/executing-templates
 
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  """
 
  template_fields: Sequence[str] = (
@@ -334,11 +337,13 @@ def __init__(
  append_job_name: bool = True,
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
  expected_terminal_state: str | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
 
  self.template = template
+ self.cancel_on_kill = cancel_on_kill
  self.job_name = job_name
  self.options = options or {}
  self.dataflow_default_options = dataflow_default_options or {}
@@ -437,6 +442,7 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -453,7 +459,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
  return job_id
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.log.info("Cancelling job %s", self.job_name)
  self.hook.cancel_job(
@@ -525,6 +534,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
  :param deferrable: Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: The expected final status of the operator on which the corresponding
  Airflow task succeeds. When not specified, it will be determined by the hook.
  :param append_job_name: True if unique suffix has to be appended to job name.
@@ -550,11 +562,13 @@ def __init__(
  append_job_name: bool = True,
  expected_terminal_state: str | None = None,
  poll_sleep: int = 10,
+ cancel_on_kill: bool = True,
  *args,
  **kwargs,
  ) -> None:
  super().__init__(*args, **kwargs)
  self.body = body
+ self.cancel_on_kill = cancel_on_kill
  self.location = location
  self.project_id = project_id
  self.gcp_conn_id = gcp_conn_id
@@ -633,6 +647,8 @@ def set_current_job(current_job):
  poll_sleep=self.poll_sleep,
  impersonation_chain=self.impersonation_chain,
  cancel_timeout=self.cancel_timeout,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -658,7 +674,10 @@ def execute_complete(self, context: Context, event: dict) -> dict[str, str]:
  return job
 
  def on_kill(self) -> None:
+ """Cancel the running job; a kill of a deferred task cancels through the trigger instead."""
  self.log.info("On kill.")
+ if not self.cancel_on_kill:
+ return
  if self.job is not None:
  self.hook.cancel_job(
  job_id=self.job.get("id"),
@@ -689,10 +708,13 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
  :param append_job_name: Optional. Set to True if a unique suffix has to be appended to the `job_name`.
  Defaults to True.
  :param drain_pipeline: Optional. Set to True if you want to stop a streaming pipeline job by draining it
- instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs
-or in the deferrable mode. Defaults to False.
+ instead of canceling when killing the task instance. Note that this does not work for batch pipeline jobs.
+ Defaults to False.
  For more info see: https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
  :param deferrable: Optional. Run operator in the deferrable mode.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the task is killed,
+ both while the operator is running and, for a deferred task, while it waits in the
+ triggerer.
  :param expected_terminal_state: Optional. The expected terminal state of the Dataflow job at which the
  operator task is set to succeed. Defaults to 'JOB_STATE_DONE' for the batch jobs and 'JOB_STATE_RUNNING'
  for the streaming jobs.
@@ -748,10 +770,12 @@ def __init__(
  jinja_variables: dict[str, str] | None = None,
  options: dict[str, Any] | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
  self.job_name = job_name
+ self.cancel_on_kill = cancel_on_kill
  self.yaml_pipeline_file = yaml_pipeline_file
  self.region = region
  self.project_id = project_id
@@ -793,6 +817,8 @@ def execute(self, context: Context) -> dict[str, Any]:
  cancel_timeout=self.cancel_timeout,
  expected_terminal_state=self.expected_terminal_state,
  impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_on_kill=self.cancel_on_kill,
  ),
  method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
  )
@@ -818,10 +844,13 @@ def on_kill(self):
  """
  Cancel the dataflow job if a task instance gets killed.
 
- This method will not be called if a task instance is killed in a deferred
- state.
+ This method is not called for a task instance killed in a deferred state;
+ in that case the trigger cancels the job instead, honoring cancel_on_kill
+ and drain_pipeline.
  """
  self.log.info("On kill called.")
+ if not self.cancel_on_kill:
+ return
  if self.job_id:
  self.hook.cancel_job(
  job_id=self.job_id,

providers/google/src/airflow/providers/google/cloud/triggers/dataflow.py

@@ -22,6 +22,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from asgiref.sync import sync_to_async
 from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.dataflow_v1beta3 import JobState
 from google.cloud.dataflow_v1beta3.types import (
@@ -33,7 +34,11 @@
  MetricUpdate,
 )
 
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
+from airflow.providers.google.cloud.hooks.dataflow import (
+ AsyncDataflowHook,
+ DataflowHook,
+ DataflowJobStatus,
+)
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -62,6 +67,10 @@ class TemplateJobStartTrigger(BaseTrigger):
  account from the list granting this role to the originating account (templated).
  :param cancel_timeout: Optional. How long (in seconds) operator should wait for the pipeline to be
  successfully cancelled when task is being killed.
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark-failed, clear or mark-succeeded).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling, matching the operator's behaviour before deferral.
  """
 
  def __init__(
@@ -73,6 +82,8 @@ def __init__(
  poll_sleep: int = 10,
  impersonation_chain: str | Sequence[str] | None = None,
  cancel_timeout: int | None = 5 * 60,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -82,6 +93,8 @@ def __init__(
  self.poll_sleep = poll_sleep
  self.impersonation_chain = impersonation_chain
  self.cancel_timeout = cancel_timeout
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -95,9 +108,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "poll_sleep": self.poll_sleep,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job status or yield certain Events.
@@ -300,6 +352,10 @@ class DataflowStartYamlJobTrigger(BaseTrigger):
  If set as a sequence, the identities from the list must grant
  Service Account Token Creator IAM role to the directly preceding identity, with first
  account from the list granting this role to the originating account (templated).
+ :param cancel_on_kill: If True (default), cancel the Dataflow job when the user acts on the
+ deferred task (mark failed, clear, or mark success).
+ :param drain_pipeline: Optional. Set to True if you want a streaming job to be stopped by
+ draining it instead of cancelling when the task is killed.
  """
 
  def __init__(
@@ -312,6 +368,8 @@ def __init__(
  cancel_timeout: int | None = 5 * 60,
  expected_terminal_state: str | None = None,
  impersonation_chain: str | Sequence[str] | None = None,
+ cancel_on_kill: bool = True,
+ drain_pipeline: bool = False,
  ):
  super().__init__()
  self.project_id = project_id
@@ -322,6 +380,8 @@ def __init__(
  self.cancel_timeout = cancel_timeout
  self.expected_terminal_state = expected_terminal_state
  self.impersonation_chain = impersonation_chain
+ self.cancel_on_kill = cancel_on_kill
+ self.drain_pipeline = drain_pipeline
 
  def serialize(self) -> tuple[str, dict[str, Any]]:
  """Serialize class arguments and classpath."""
@@ -336,9 +396,48 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
  "expected_terminal_state": self.expected_terminal_state,
  "impersonation_chain": self.impersonation_chain,
  "cancel_timeout": self.cancel_timeout,
+ "cancel_on_kill": self.cancel_on_kill,
+ "drain_pipeline": self.drain_pipeline,
  },
  )
 
+ async def on_kill(self) -> None:
+ """Stop the Dataflow job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or not self.job_id or not self.project_id:
+ return
+ self.log.info(
+ "Stopping Dataflow job. Project ID: %s, Location: %s, Job ID: %s, drain: %s",
+ self.project_id,
+ self.location,
+ self.job_id,
+ self.drain_pipeline,
+ )
+ try:
+ # Build the synchronous hook and cancel inside the worker thread: the hook resolves the
+ # connection eagerly during construction, which must not run in the triggerer's event loop.
+ await sync_to_async(self._stop_job)()
+ self.log.info("Dataflow job %s stopped.", self.job_id)
+ except Exception:
+ self.log.exception(
+ "Failed to stop Dataflow job %s. The job may still be running.",
+ self.job_id,
+ )
+
+ def _stop_job(self) -> None:
+ """Cancel or drain the Dataflow job through the synchronous hook (runs off the event loop)."""
+ hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ drain_pipeline=self.drain_pipeline,
+ cancel_timeout=self.cancel_timeout,
+ poll_sleep=self.poll_sleep,
+ )
+ hook.cancel_job(
+ job_id=self.job_id,
+ project_id=self.project_id,
+ location=self.location,
+ )
+
  async def run(self):
  """
  Fetch job and yield events depending on the job's type and state.