You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2025/02/20 12:35:40 UTC
(airflow) branch main updated: Refactor deferrable mode for BeamRunPythonPipelineOperator and BeamRunJavaPipelineOperator (#46678)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ee68ddf605c Refactor deferrable mode for BeamRunPythonPipelineOperator and BeamRunJavaPipelineOperator (#46678)
ee68ddf605c is described below
commit ee68ddf605cdcc457af2cb56f8443571757fe32f
Author: Maksim <ma...@google.com>
AuthorDate: Thu Feb 20 04:35:30 2025 -0800
Refactor deferrable mode for BeamRunPythonPipelineOperator and BeamRunJavaPipelineOperator (#46678)
---
.../airflow/providers/apache/beam/hooks/beam.py | 22 +-
.../providers/apache/beam/operators/beam.py | 295 ++++++++++-----------
.../airflow/providers/apache/beam/triggers/beam.py | 104 +-------
.../beam/tests/unit/apache/beam/hooks/test_beam.py | 32 +--
.../tests/unit/apache/beam/operators/test_beam.py | 18 +-
.../tests/unit/apache/beam/triggers/test_beam.py | 49 ----
.../cloud/dataflow/example_dataflow_native_java.py | 34 ++-
.../dataflow/example_dataflow_native_python.py | 32 ++-
8 files changed, 225 insertions(+), 361 deletions(-)
diff --git a/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py b/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
index 20b83df4631..08102aabfca 100644
--- a/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
+++ b/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
@@ -95,7 +95,7 @@ def process_fd(
fd,
log: logging.Logger,
process_line_callback: Callable[[str], None] | None = None,
- check_job_status_callback: Callable[[], bool | None] | None = None,
+ is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None,
):
"""
Print output to logs.
@@ -117,7 +117,7 @@ def process_fd(
if process_line_callback:
process_line_callback(line)
func_log(line.rstrip("\n"))
- if check_job_status_callback and check_job_status_callback():
+ if is_dataflow_job_id_exist_callback and is_dataflow_job_id_exist_callback():
return
@@ -126,7 +126,7 @@ def run_beam_command(
log: logging.Logger,
process_line_callback: Callable[[str], None] | None = None,
working_directory: str | None = None,
- check_job_status_callback: Callable[[], bool | None] | None = None,
+ is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None,
) -> None:
"""
Run pipeline command in subprocess.
@@ -158,8 +158,8 @@ def run_beam_command(
continue
for readable_fd in readable_fds:
- process_fd(proc, readable_fd, log, process_line_callback, check_job_status_callback)
- if check_job_status_callback and check_job_status_callback():
+ process_fd(proc, readable_fd, log, process_line_callback, is_dataflow_job_id_exist_callback)
+ if is_dataflow_job_id_exist_callback and is_dataflow_job_id_exist_callback():
return
if proc.poll() is not None:
@@ -167,7 +167,7 @@ def run_beam_command(
# Corner case: check if more output was created between the last read and the process termination
for readable_fd in reads:
- process_fd(proc, readable_fd, log, process_line_callback, check_job_status_callback)
+ process_fd(proc, readable_fd, log, process_line_callback, is_dataflow_job_id_exist_callback)
log.info("Process exited with return code: %s", proc.returncode)
@@ -198,7 +198,7 @@ class BeamHook(BaseHook):
command_prefix: list[str],
process_line_callback: Callable[[str], None] | None = None,
working_directory: str | None = None,
- check_job_status_callback: Callable[[], bool | None] | None = None,
+ is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None,
) -> None:
cmd = [*command_prefix, f"--runner={self.runner}"]
if variables:
@@ -208,7 +208,7 @@ class BeamHook(BaseHook):
process_line_callback=process_line_callback,
working_directory=working_directory,
log=self.log,
- check_job_status_callback=check_job_status_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
)
def start_python_pipeline(
@@ -220,7 +220,7 @@ class BeamHook(BaseHook):
py_requirements: list[str] | None = None,
py_system_site_packages: bool = False,
process_line_callback: Callable[[str], None] | None = None,
- check_job_status_callback: Callable[[], bool | None] | None = None,
+ is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None,
):
"""
Start Apache Beam python pipeline.
@@ -289,7 +289,7 @@ class BeamHook(BaseHook):
variables=variables,
command_prefix=command_prefix,
process_line_callback=process_line_callback,
- check_job_status_callback=check_job_status_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
)
def start_java_pipeline(
@@ -298,6 +298,7 @@ class BeamHook(BaseHook):
jar: str,
job_class: str | None = None,
process_line_callback: Callable[[str], None] | None = None,
+ is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None,
) -> None:
"""
Start Apache Beam Java pipeline.
@@ -316,6 +317,7 @@ class BeamHook(BaseHook):
variables=variables,
command_prefix=command_prefix,
process_line_callback=process_line_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
)
def start_go_pipeline(
diff --git a/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py b/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py
index 812c74cd954..38a573edd89 100644
--- a/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py
+++ b/providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py
@@ -37,11 +37,15 @@ from airflow.providers.apache.beam.hooks.beam import BeamHook, BeamRunnerType
from airflow.providers.apache.beam.triggers.beam import BeamJavaPipelineTrigger, BeamPythonPipelineTrigger
from airflow.providers.google.cloud.hooks.dataflow import (
DataflowHook,
+ DataflowJobStatus,
process_line_and_extract_dataflow_job_id_callback,
)
from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
from airflow.providers.google.cloud.links.dataflow import DataflowJobLink
from airflow.providers.google.cloud.operators.dataflow import CheckJobRunning, DataflowConfiguration
+from airflow.providers.google.cloud.triggers.dataflow import (
+ DataflowJobStatusTrigger,
+)
from airflow.utils.helpers import convert_camel_to_snake, exactly_one
from airflow.version import version
@@ -67,7 +71,7 @@ class BeamDataflowMixin(metaclass=ABCMeta):
self,
pipeline_options: dict,
job_name_variable_key: str | None = None,
- ) -> tuple[str, dict, Callable[[str], None], Callable[[], bool | None]]:
+ ) -> tuple[str, dict, Callable[[str], None], Callable[[], bool]]:
self.dataflow_hook = self.__set_dataflow_hook()
self.dataflow_config.project_id = self.dataflow_config.project_id or self.dataflow_hook.project_id
dataflow_job_name = self.__get_dataflow_job_name()
@@ -75,8 +79,8 @@ class BeamDataflowMixin(metaclass=ABCMeta):
pipeline_options, dataflow_job_name, job_name_variable_key
)
process_line_callback = self.__get_dataflow_process_callback()
- check_job_status_callback = self.__check_dataflow_job_status_callback()
- return dataflow_job_name, pipeline_options, process_line_callback, check_job_status_callback
+ is_dataflow_job_id_exist_callback = self.__is_dataflow_job_id_exist_callback()
+ return dataflow_job_name, pipeline_options, process_line_callback, is_dataflow_job_id_exist_callback
def __set_dataflow_hook(self) -> DataflowHook:
self.dataflow_hook = DataflowHook(
@@ -125,18 +129,11 @@ class BeamDataflowMixin(metaclass=ABCMeta):
on_new_job_id_callback=set_current_dataflow_job_id
)
- def __check_dataflow_job_status_callback(self) -> Callable[[], bool | None]:
- def check_dataflow_job_status() -> bool | None:
- if self.dataflow_job_id and self.dataflow_hook:
- return self.dataflow_hook.is_job_done(
- location=self.dataflow_config.location,
- project_id=self.dataflow_config.project_id,
- job_id=self.dataflow_job_id,
- )
- else:
- return None
+ def __is_dataflow_job_id_exist_callback(self) -> Callable[[], bool]:
+ def is_dataflow_job_id_exist() -> bool:
+ return True if self.dataflow_job_id else False
- return check_dataflow_job_status
+ return is_dataflow_job_id_exist
class BeamBasePipelineOperator(BaseOperator, BeamDataflowMixin, ABC):
@@ -219,11 +216,11 @@ class BeamBasePipelineOperator(BaseOperator, BeamDataflowMixin, ABC):
self,
format_pipeline_options: bool = False,
job_name_variable_key: str | None = None,
- ) -> tuple[bool, str | None, dict, Callable[[str], None] | None, Callable[[], bool | None] | None]:
+ ) -> tuple[bool, str | None, dict, Callable[[str], None] | None, Callable[[], bool] | None]:
self.beam_hook = BeamHook(runner=self.runner)
pipeline_options = self.default_pipeline_options.copy()
process_line_callback: Callable[[str], None] | None = None
- check_job_status_callback: Callable[[], bool | None] | None = None
+ is_dataflow_job_id_exist_callback: Callable[[], bool] | None = None
is_dataflow = self.runner.lower() == BeamRunnerType.DataflowRunner.lower()
dataflow_job_name: str | None = None
if is_dataflow:
@@ -231,7 +228,7 @@ class BeamBasePipelineOperator(BaseOperator, BeamDataflowMixin, ABC):
dataflow_job_name,
pipeline_options,
process_line_callback,
- check_job_status_callback,
+ is_dataflow_job_id_exist_callback,
) = self._set_dataflow(
pipeline_options=pipeline_options,
job_name_variable_key=job_name_variable_key,
@@ -249,7 +246,7 @@ class BeamBasePipelineOperator(BaseOperator, BeamDataflowMixin, ABC):
dataflow_job_name,
snake_case_pipeline_options,
process_line_callback,
- check_job_status_callback,
+ is_dataflow_job_id_exist_callback,
)
return (
@@ -257,7 +254,7 @@ class BeamBasePipelineOperator(BaseOperator, BeamDataflowMixin, ABC):
dataflow_job_name,
pipeline_options,
process_line_callback,
- check_job_status_callback,
+ is_dataflow_job_id_exist_callback,
)
def execute_complete(self, context: Context, event: dict[str, Any]):
@@ -274,18 +271,6 @@ class BeamBasePipelineOperator(BaseOperator, BeamDataflowMixin, ABC):
self.task_id,
event["message"],
)
- self.dataflow_job_id = event["dataflow_job_id"]
- self.project_id = event["project_id"]
- self.location = event["location"]
-
- DataflowJobLink.persist(
- self,
- context,
- self.project_id,
- self.location,
- self.dataflow_job_id,
- )
- return {"dataflow_job_id": self.dataflow_job_id}
class BeamRunPythonPipelineOperator(BeamBasePipelineOperator):
@@ -379,18 +364,26 @@ class BeamRunPythonPipelineOperator(BeamBasePipelineOperator):
self.dataflow_job_name,
self.snake_case_pipeline_options,
self.process_line_callback,
- self.check_job_status_callback,
+ self.is_dataflow_job_id_exist_callback,
) = self._init_pipeline_options(format_pipeline_options=True, job_name_variable_key="job_name")
if not self.beam_hook:
raise AirflowException("Beam hook is not defined.")
- # Check deferrable parameter passed to the operator
- # to determine type of run - asynchronous or synchronous
- if self.deferrable:
- self.execute_async(context)
- else:
- return self.execute_sync(context)
- def execute_sync(self, context: Context):
+ if self.deferrable and not self.is_dataflow:
+ self.defer(
+ trigger=BeamPythonPipelineTrigger(
+ variables=self.snake_case_pipeline_options,
+ py_file=self.py_file,
+ py_options=self.py_options,
+ py_interpreter=self.py_interpreter,
+ py_requirements=self.py_requirements,
+ py_system_site_packages=self.py_system_site_packages,
+ runner=self.runner,
+ gcp_conn_id=self.gcp_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
with ExitStack() as exit_stack:
if self.py_file.lower().startswith("gs://"):
gcs_hook = GCSHook(gcp_conn_id=self.gcp_conn_id)
@@ -404,25 +397,7 @@ class BeamRunPythonPipelineOperator(BeamBasePipelineOperator):
self.snake_case_pipeline_options["requirements_file"] = tmp_req_file.name
if self.is_dataflow and self.dataflow_hook:
- with self.dataflow_hook.provide_authorized_gcloud():
- self.beam_hook.start_python_pipeline(
- variables=self.snake_case_pipeline_options,
- py_file=self.py_file,
- py_options=self.py_options,
- py_interpreter=self.py_interpreter,
- py_requirements=self.py_requirements,
- py_system_site_packages=self.py_system_site_packages,
- process_line_callback=self.process_line_callback,
- check_job_status_callback=self.check_job_status_callback,
- )
- DataflowJobLink.persist(
- self,
- context,
- self.dataflow_config.project_id,
- self.dataflow_config.location,
- self.dataflow_job_id,
- )
- return {"dataflow_job_id": self.dataflow_job_id}
+ return self.execute_on_dataflow(context)
else:
self.beam_hook.start_python_pipeline(
variables=self.snake_case_pipeline_options,
@@ -431,41 +406,48 @@ class BeamRunPythonPipelineOperator(BeamBasePipelineOperator):
py_interpreter=self.py_interpreter,
py_requirements=self.py_requirements,
py_system_site_packages=self.py_system_site_packages,
- process_line_callback=self.process_line_callback,
)
- def execute_async(self, context: Context):
- if self.is_dataflow and self.dataflow_hook:
- with self.dataflow_hook.provide_authorized_gcloud():
- self.defer(
- trigger=BeamPythonPipelineTrigger(
- variables=self.snake_case_pipeline_options,
- py_file=self.py_file,
- py_options=self.py_options,
- py_interpreter=self.py_interpreter,
- py_requirements=self.py_requirements,
- py_system_site_packages=self.py_system_site_packages,
- runner=self.runner,
- gcp_conn_id=self.gcp_conn_id,
- project_id=self.dataflow_config.project_id,
- location=self.dataflow_config.location,
- ),
- method_name="execute_complete",
- )
- else:
+ def execute_on_dataflow(self, context: Context):
+ """Execute the Apache Beam Pipeline on Dataflow runner."""
+ if not self.dataflow_hook:
+ self.dataflow_hook = self.__set_dataflow_hook()
+ with self.dataflow_hook.provide_authorized_gcloud():
+ self.beam_hook.start_python_pipeline(
+ variables=self.snake_case_pipeline_options,
+ py_file=self.py_file,
+ py_options=self.py_options,
+ py_interpreter=self.py_interpreter,
+ py_requirements=self.py_requirements,
+ py_system_site_packages=self.py_system_site_packages,
+ process_line_callback=self.process_line_callback,
+ is_dataflow_job_id_exist_callback=self.is_dataflow_job_id_exist_callback,
+ )
+ DataflowJobLink.persist(
+ self,
+ context,
+ self.dataflow_config.project_id,
+ self.dataflow_config.location,
+ self.dataflow_job_id,
+ )
+ if self.deferrable:
self.defer(
- trigger=BeamPythonPipelineTrigger(
- variables=self.snake_case_pipeline_options,
- py_file=self.py_file,
- py_options=self.py_options,
- py_interpreter=self.py_interpreter,
- py_requirements=self.py_requirements,
- py_system_site_packages=self.py_system_site_packages,
- runner=self.runner,
+ trigger=DataflowJobStatusTrigger(
+ job_id=self.dataflow_job_id,
+ expected_statuses={DataflowJobStatus.JOB_STATE_DONE},
+ project_id=self.dataflow_config.project_id,
+ location=self.dataflow_config.location,
gcp_conn_id=self.gcp_conn_id,
),
method_name="execute_complete",
)
+ self.dataflow_hook.wait_for_done(
+ job_name=self.dataflow_job_name,
+ location=self.dataflow_config.location,
+ job_id=self.dataflow_job_id,
+ project_id=self.dataflow_config.project_id,
+ )
+ return {"dataflow_job_id": self.dataflow_job_id}
def on_kill(self) -> None:
if self.dataflow_hook and self.dataflow_job_id:
@@ -552,17 +534,23 @@ class BeamRunJavaPipelineOperator(BeamBasePipelineOperator):
self.dataflow_job_name,
self.pipeline_options,
self.process_line_callback,
- _,
+ self.is_dataflow_job_id_exist_callback,
) = self._init_pipeline_options()
if not self.beam_hook:
raise AirflowException("Beam hook is not defined.")
- if self.deferrable:
- self.execute_async(context)
- else:
- return self.execute_sync(context)
- def execute_sync(self, context: Context):
- """Execute the Apache Beam Pipeline."""
+ if self.deferrable and not self.is_dataflow:
+ self.defer(
+ trigger=BeamJavaPipelineTrigger(
+ variables=self.pipeline_options,
+ jar=self.jar,
+ job_class=self.job_class,
+ runner=self.runner,
+ gcp_conn_id=self.gcp_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
with ExitStack() as exit_stack:
if self.jar.lower().startswith("gs://"):
gcs_hook = GCSHook(self.gcp_conn_id)
@@ -570,86 +558,69 @@ class BeamRunJavaPipelineOperator(BeamBasePipelineOperator):
self.jar = tmp_gcs_file.name
if self.is_dataflow and self.dataflow_hook:
- is_running = self.dataflow_config.check_if_running == CheckJobRunning.WaitForRun
- while is_running and self.dataflow_config.check_if_running == CheckJobRunning.WaitForRun:
- # The reason for disable=no-value-for-parameter is that project_id parameter is
- # required but here is not passed, moreover it cannot be passed here.
- # This method is wrapped by @_fallback_to_project_id_from_variables decorator which
- # fallback project_id value from variables and raise error if project_id is
- # defined both in variables and as parameter (here is already defined in variables)
- is_running = self.dataflow_hook.is_job_dataflow_running(
- name=self.dataflow_config.job_name,
- variables=self.pipeline_options,
- location=self.dataflow_config.location,
- )
-
- if not is_running:
- self.pipeline_options["jobName"] = self.dataflow_job_name
- with self.dataflow_hook.provide_authorized_gcloud():
- self.beam_hook.start_java_pipeline(
- variables=self.pipeline_options,
- jar=self.jar,
- job_class=self.job_class,
- process_line_callback=self.process_line_callback,
- )
- if self.dataflow_job_name and self.dataflow_config.location:
- multiple_jobs = self.dataflow_config.multiple_jobs or False
- DataflowJobLink.persist(
- self,
- context,
- self.dataflow_config.project_id,
- self.dataflow_config.location,
- self.dataflow_job_id,
- )
- self.dataflow_hook.wait_for_done(
- job_name=self.dataflow_job_name,
- location=self.dataflow_config.location,
- job_id=self.dataflow_job_id,
- multiple_jobs=multiple_jobs,
- project_id=self.dataflow_config.project_id,
- )
- return {"dataflow_job_id": self.dataflow_job_id}
+ return self.execute_on_dataflow(context)
else:
self.beam_hook.start_java_pipeline(
variables=self.pipeline_options,
jar=self.jar,
job_class=self.job_class,
- process_line_callback=self.process_line_callback,
)
- def execute_async(self, context: Context):
- if self.is_dataflow and self.dataflow_hook:
+ def execute_on_dataflow(self, context: Context):
+ """Execute the Apache Beam Pipeline on Dataflow runner."""
+ if not self.dataflow_hook:
+ self.dataflow_hook = self.__set_dataflow_hook()
+ is_running = self.dataflow_config.check_if_running == CheckJobRunning.WaitForRun
+ while is_running and self.dataflow_config.check_if_running == CheckJobRunning.WaitForRun:
+ # The reason for disable=no-value-for-parameter is that project_id parameter is
+ # required but here is not passed, moreover it cannot be passed here.
+ # This method is wrapped by @_fallback_to_project_id_from_variables decorator which
+ # fallback project_id value from variables and raise error if project_id is
+ # defined both in variables and as parameter (here is already defined in variables)
+ is_running = self.dataflow_hook.is_job_dataflow_running(
+ name=self.dataflow_config.job_name,
+ variables=self.pipeline_options,
+ location=self.dataflow_config.location,
+ )
+
+ if not is_running:
+ self.pipeline_options["jobName"] = self.dataflow_job_name
with self.dataflow_hook.provide_authorized_gcloud():
- self.pipeline_options["jobName"] = self.dataflow_job_name
- self.defer(
- trigger=BeamJavaPipelineTrigger(
- variables=self.pipeline_options,
- jar=self.jar,
- job_class=self.job_class,
- runner=self.runner,
- check_if_running=self.dataflow_config.check_if_running == CheckJobRunning.WaitForRun,
- project_id=self.dataflow_config.project_id,
- location=self.dataflow_config.location,
- job_name=self.dataflow_job_name,
- gcp_conn_id=self.gcp_conn_id,
- impersonation_chain=self.dataflow_config.impersonation_chain,
- poll_sleep=self.dataflow_config.poll_sleep,
- cancel_timeout=self.dataflow_config.cancel_timeout,
- ),
- method_name="execute_complete",
- )
- else:
- self.defer(
- trigger=BeamJavaPipelineTrigger(
+ self.beam_hook.start_java_pipeline(
variables=self.pipeline_options,
jar=self.jar,
job_class=self.job_class,
- runner=self.runner,
- check_if_running=self.dataflow_config.check_if_running == CheckJobRunning.WaitForRun,
- gcp_conn_id=self.gcp_conn_id,
- ),
- method_name="execute_complete",
- )
+ process_line_callback=self.process_line_callback,
+ is_dataflow_job_id_exist_callback=self.is_dataflow_job_id_exist_callback,
+ )
+ if self.dataflow_job_name and self.dataflow_config.location:
+ DataflowJobLink.persist(
+ self,
+ context,
+ self.dataflow_config.project_id,
+ self.dataflow_config.location,
+ self.dataflow_job_id,
+ )
+ if self.deferrable:
+ self.defer(
+ trigger=DataflowJobStatusTrigger(
+ job_id=self.dataflow_job_id,
+ expected_statuses={DataflowJobStatus.JOB_STATE_DONE},
+ project_id=self.dataflow_config.project_id,
+ location=self.dataflow_config.location,
+ gcp_conn_id=self.gcp_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+ multiple_jobs = self.dataflow_config.multiple_jobs or False
+ self.dataflow_hook.wait_for_done(
+ job_name=self.dataflow_job_name,
+ location=self.dataflow_config.location,
+ job_id=self.dataflow_job_id,
+ multiple_jobs=multiple_jobs,
+ project_id=self.dataflow_config.project_id,
+ )
+ return {"dataflow_job_id": self.dataflow_job_id}
def on_kill(self) -> None:
if self.dataflow_hook and self.dataflow_job_id:
diff --git a/providers/apache/beam/src/airflow/providers/apache/beam/triggers/beam.py b/providers/apache/beam/src/airflow/providers/apache/beam/triggers/beam.py
index 1c77089b9f2..d4a6d7c9ac3 100644
--- a/providers/apache/beam/src/airflow/providers/apache/beam/triggers/beam.py
+++ b/providers/apache/beam/src/airflow/providers/apache/beam/triggers/beam.py
@@ -18,16 +18,10 @@ from __future__ import annotations
import asyncio
import contextlib
-from collections.abc import AsyncIterator, Sequence
-from typing import IO, Any, Callable
+from collections.abc import AsyncIterator
+from typing import IO, Any
-from google.cloud.dataflow_v1beta3 import ListJobsRequest
-
-from airflow.providers.apache.beam.hooks.beam import BeamAsyncHook, BeamRunnerType
-from airflow.providers.google.cloud.hooks.dataflow import (
- AsyncDataflowHook,
- process_line_and_extract_dataflow_job_id_callback,
-)
+from airflow.providers.apache.beam.hooks.beam import BeamAsyncHook
from airflow.providers.google.cloud.hooks.gcs import GCSHook
from airflow.triggers.base import BaseTrigger, TriggerEvent
@@ -39,18 +33,6 @@ class BeamPipelineBaseTrigger(BaseTrigger):
def _get_async_hook(*args, **kwargs) -> BeamAsyncHook:
return BeamAsyncHook(*args, **kwargs)
- @staticmethod
- def _get_sync_dataflow_hook(**kwargs) -> AsyncDataflowHook:
- return AsyncDataflowHook(**kwargs)
-
- def _get_dataflow_process_callback(self) -> Callable[[str], None]:
- def set_current_dataflow_job_id(job_id):
- self.dataflow_job_id = job_id
-
- return process_line_and_extract_dataflow_job_id_callback(
- on_new_job_id_callback=set_current_dataflow_job_id
- )
-
class BeamPythonPipelineTrigger(BeamPipelineBaseTrigger):
"""
@@ -70,8 +52,6 @@ class BeamPythonPipelineTrigger(BeamPipelineBaseTrigger):
:param py_system_site_packages: Whether to include system_site_packages in your virtualenv.
See virtualenv documentation for more information.
This option is only relevant if the ``py_requirements`` parameter is not None.
- :param project_id: Optional, the Google Cloud project ID in which to start a job.
- :param location: Optional, Job location.
:param runner: Runner on which pipeline will be run. By default, "DirectRunner" is being used.
Other possible options: DataflowRunner, SparkRunner, FlinkRunner, PortableRunner.
See: :class:`~providers.apache.beam.hooks.beam.BeamRunnerType`
@@ -87,8 +67,6 @@ class BeamPythonPipelineTrigger(BeamPipelineBaseTrigger):
py_interpreter: str = "python3",
py_requirements: list[str] | None = None,
py_system_site_packages: bool = False,
- project_id: str | None = None,
- location: str | None = None,
runner: str = "DirectRunner",
gcp_conn_id: str = "google_cloud_default",
):
@@ -99,9 +77,6 @@ class BeamPythonPipelineTrigger(BeamPipelineBaseTrigger):
self.py_interpreter = py_interpreter
self.py_requirements = py_requirements
self.py_system_site_packages = py_system_site_packages
- self.dataflow_job_id: str | None = None
- self.project_id = project_id
- self.location = location
self.runner = runner
self.gcp_conn_id = gcp_conn_id
@@ -116,8 +91,6 @@ class BeamPythonPipelineTrigger(BeamPipelineBaseTrigger):
"py_interpreter": self.py_interpreter,
"py_requirements": self.py_requirements,
"py_system_site_packages": self.py_system_site_packages,
- "project_id": self.project_id,
- "location": self.location,
"runner": self.runner,
"gcp_conn_id": self.gcp_conn_id,
},
@@ -126,7 +99,6 @@ class BeamPythonPipelineTrigger(BeamPipelineBaseTrigger):
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
"""Get current pipeline status and yields a TriggerEvent."""
hook = self._get_async_hook(runner=self.runner)
- is_dataflow = self.runner.lower() == BeamRunnerType.DataflowRunner.lower()
try:
# Get the current running event loop to manage I/O operations asynchronously
@@ -152,7 +124,6 @@ class BeamPythonPipelineTrigger(BeamPipelineBaseTrigger):
py_interpreter=self.py_interpreter,
py_requirements=self.py_requirements,
py_system_site_packages=self.py_system_site_packages,
- process_line_callback=self._get_dataflow_process_callback() if is_dataflow else None,
)
except Exception as e:
self.log.exception("Exception occurred while checking for pipeline state")
@@ -163,9 +134,6 @@ class BeamPythonPipelineTrigger(BeamPipelineBaseTrigger):
{
"status": "success",
"message": "Pipeline has finished SUCCESSFULLY",
- "dataflow_job_id": self.dataflow_job_id,
- "project_id": self.project_id,
- "location": self.location,
}
)
else:
@@ -184,23 +152,7 @@ class BeamJavaPipelineTrigger(BeamPipelineBaseTrigger):
Other possible options: DataflowRunner, SparkRunner, FlinkRunner, PortableRunner.
See: :class:`~providers.apache.beam.hooks.beam.BeamRunnerType`
See: https://beam.apache.org/documentation/runners/capability-matrix/
- :param check_if_running: Optional. Before running job, validate that a previous run is not in process.
- :param project_id: Optional. The Google Cloud project ID in which to start a job.
- :param location: Optional. Job location.
- :param job_name: Optional. The 'jobName' to use when executing the Dataflow job.
:param gcp_conn_id: Optional. The connection ID to use connecting to Google Cloud.
- :param impersonation_chain: Optional. GCP service account to impersonate using short-term
- credentials, or chained list of accounts required to get the access_token
- of the last account in the list, which will be impersonated in the request.
- If set as a string, the account must grant the originating account
- the Service Account Token Creator IAM role.
- 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 poll_sleep: Optional. The time in seconds to sleep between polling GCP for the dataflow job status.
- Default value is 10s.
- :param cancel_timeout: Optional. How long (in seconds) operator should wait for the pipeline to be
- successfully cancelled when task is being killed. Default value is 300s.
"""
def __init__(
@@ -209,29 +161,14 @@ class BeamJavaPipelineTrigger(BeamPipelineBaseTrigger):
jar: str,
job_class: str | None = None,
runner: str = "DirectRunner",
- check_if_running: bool = False,
- project_id: str | None = None,
- location: str | None = None,
- job_name: str | None = None,
gcp_conn_id: str = "google_cloud_default",
- impersonation_chain: str | Sequence[str] | None = None,
- poll_sleep: int = 10,
- cancel_timeout: int | None = None,
):
super().__init__()
self.variables = variables
self.jar = jar
self.job_class = job_class
self.runner = runner
- self.check_if_running = check_if_running
- self.project_id = project_id
- self.location = location
- self.job_name = job_name
self.gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
- self.poll_sleep = poll_sleep
- self.cancel_timeout = cancel_timeout
- self.dataflow_job_id: str | None = None
def serialize(self) -> tuple[str, dict[str, Any]]:
"""Serialize BeamJavaPipelineTrigger arguments and classpath."""
@@ -242,45 +179,14 @@ class BeamJavaPipelineTrigger(BeamPipelineBaseTrigger):
"jar": self.jar,
"job_class": self.job_class,
"runner": self.runner,
- "check_if_running": self.check_if_running,
- "project_id": self.project_id,
- "location": self.location,
- "job_name": self.job_name,
"gcp_conn_id": self.gcp_conn_id,
- "impersonation_chain": self.impersonation_chain,
- "poll_sleep": self.poll_sleep,
- "cancel_timeout": self.cancel_timeout,
},
)
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
"""Get current Java pipeline status and yields a TriggerEvent."""
hook = self._get_async_hook(runner=self.runner)
- is_dataflow = self.runner.lower() == BeamRunnerType.DataflowRunner.lower()
-
return_code = 0
- if self.check_if_running:
- dataflow_hook = self._get_sync_dataflow_hook(
- gcp_conn_id=self.gcp_conn_id,
- poll_sleep=self.poll_sleep,
- impersonation_chain=self.impersonation_chain,
- cancel_timeout=self.cancel_timeout,
- )
- is_running = True
- while is_running:
- try:
- jobs = await dataflow_hook.list_jobs(
- project_id=self.project_id,
- location=self.location,
- jobs_filter=ListJobsRequest.Filter.ACTIVE,
- )
- is_running = bool([job async for job in jobs if job.name == self.job_name])
- except Exception as e:
- self.log.exception("Exception occurred while requesting jobs with name %s", self.job_name)
- yield TriggerEvent({"status": "error", "message": str(e)})
- return
- if is_running:
- await asyncio.sleep(self.poll_sleep)
try:
# Get the current running event loop to manage I/O operations asynchronously
loop = asyncio.get_running_loop()
@@ -302,7 +208,6 @@ class BeamJavaPipelineTrigger(BeamPipelineBaseTrigger):
variables=self.variables,
jar=self.jar,
job_class=self.job_class,
- process_line_callback=self._get_dataflow_process_callback() if is_dataflow else None,
)
except Exception as e:
self.log.exception("Exception occurred while starting the Java pipeline")
@@ -313,9 +218,6 @@ class BeamJavaPipelineTrigger(BeamPipelineBaseTrigger):
{
"status": "success",
"message": "Pipeline has finished SUCCESSFULLY",
- "dataflow_job_id": self.dataflow_job_id,
- "project_id": self.project_id,
- "location": self.location,
}
)
else:
diff --git a/providers/apache/beam/tests/unit/apache/beam/hooks/test_beam.py b/providers/apache/beam/tests/unit/apache/beam/hooks/test_beam.py
index 61d2d2e4735..9807e167e18 100644
--- a/providers/apache/beam/tests/unit/apache/beam/hooks/test_beam.py
+++ b/providers/apache/beam/tests/unit/apache/beam/hooks/test_beam.py
@@ -82,14 +82,14 @@ class TestBeamHook:
def test_start_python_pipeline(self, mock_check_output, mock_runner):
hook = BeamHook(runner=DEFAULT_RUNNER)
process_line_callback = MagicMock()
- check_job_status_callback = MagicMock()
+ is_dataflow_job_id_exist_callback = MagicMock()
hook.start_python_pipeline(
variables=copy.deepcopy(BEAM_VARIABLES_PY),
py_file=PY_FILE,
py_options=PY_OPTIONS,
process_line_callback=process_line_callback,
- check_job_status_callback=check_job_status_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
)
expected_cmd = [
@@ -105,7 +105,7 @@ class TestBeamHook:
process_line_callback=process_line_callback,
working_directory=None,
log=ANY,
- check_job_status_callback=check_job_status_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
)
@mock.patch("airflow.providers.apache.beam.hooks.beam.subprocess.check_output", return_value=b"2.35.0")
@@ -126,7 +126,7 @@ class TestBeamHook:
py_requirements=None,
py_system_site_packages=False,
process_line_callback=MagicMock(),
- check_job_status_callback=MagicMock(),
+ is_dataflow_job_id_exist_callback=MagicMock(),
)
@pytest.mark.parametrize(
@@ -145,7 +145,7 @@ class TestBeamHook:
):
hook = BeamHook(runner=DEFAULT_RUNNER)
process_line_callback = MagicMock()
- check_job_status_callback = MagicMock()
+ is_dataflow_job_id_exist_callback = MagicMock()
hook.start_python_pipeline(
variables=copy.deepcopy(BEAM_VARIABLES_PY),
@@ -153,7 +153,7 @@ class TestBeamHook:
py_options=PY_OPTIONS,
py_interpreter=py_interpreter,
process_line_callback=process_line_callback,
- check_job_status_callback=check_job_status_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
)
expected_cmd = [
@@ -169,7 +169,7 @@ class TestBeamHook:
process_line_callback=process_line_callback,
working_directory=None,
log=ANY,
- check_job_status_callback=check_job_status_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
)
@pytest.mark.parametrize(
@@ -194,7 +194,7 @@ class TestBeamHook:
hook = BeamHook(runner=DEFAULT_RUNNER)
mock_virtualenv.return_value = "/dummy_dir/bin/python"
process_line_callback = MagicMock()
- check_job_status_callback = MagicMock()
+ is_dataflow_job_id_exist_callback = MagicMock()
hook.start_python_pipeline(
variables=copy.deepcopy(BEAM_VARIABLES_PY),
@@ -203,7 +203,7 @@ class TestBeamHook:
py_requirements=current_py_requirements,
py_system_site_packages=current_py_system_site_packages,
process_line_callback=process_line_callback,
- check_job_status_callback=check_job_status_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
)
expected_cmd = [
@@ -217,7 +217,7 @@ class TestBeamHook:
mock_runner.assert_called_once_with(
cmd=expected_cmd,
process_line_callback=process_line_callback,
- check_job_status_callback=check_job_status_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
working_directory=None,
log=ANY,
)
@@ -236,7 +236,7 @@ class TestBeamHook:
hook = BeamHook(runner=DEFAULT_RUNNER)
wait_for_done = mock_runner.return_value.wait_for_done
process_line_callback = MagicMock()
- check_job_status_callback = MagicMock()
+ is_dataflow_job_id_exist_callback = MagicMock()
with pytest.raises(AirflowException, match=r"Invalid method invocation\."):
hook.start_python_pipeline(
@@ -245,7 +245,7 @@ class TestBeamHook:
py_options=PY_OPTIONS,
py_requirements=[],
process_line_callback=process_line_callback,
- check_job_status_callback=check_job_status_callback,
+ is_dataflow_job_id_exist_callback=is_dataflow_job_id_exist_callback,
)
mock_runner.assert_not_called()
@@ -275,7 +275,7 @@ class TestBeamHook:
process_line_callback=process_line_callback,
working_directory=None,
log=ANY,
- check_job_status_callback=None,
+ is_dataflow_job_id_exist_callback=None,
)
@mock.patch(BEAM_STRING.format("run_beam_command"))
@@ -304,7 +304,7 @@ class TestBeamHook:
process_line_callback=process_line_callback,
working_directory=None,
log=ANY,
- check_job_status_callback=None,
+ is_dataflow_job_id_exist_callback=None,
)
@mock.patch(BEAM_STRING.format("shutil.which"))
@@ -335,7 +335,7 @@ class TestBeamHook:
process_line_callback=process_line_callback,
working_directory=go_workspace,
log=ANY,
- check_job_status_callback=None,
+ is_dataflow_job_id_exist_callback=None,
)
@mock.patch(BEAM_STRING.format("shutil.which"))
@@ -381,7 +381,7 @@ class TestBeamHook:
process_line_callback=process_line_callback,
working_directory=None,
log=ANY,
- check_job_status_callback=None,
+ is_dataflow_job_id_exist_callback=None,
)
diff --git a/providers/apache/beam/tests/unit/apache/beam/operators/test_beam.py b/providers/apache/beam/tests/unit/apache/beam/operators/test_beam.py
index a74e9002784..fc567f011ce 100644
--- a/providers/apache/beam/tests/unit/apache/beam/operators/test_beam.py
+++ b/providers/apache/beam/tests/unit/apache/beam/operators/test_beam.py
@@ -109,9 +109,6 @@ class TestBeamBasePipelineOperator:
event={
"status": "success",
"message": "Pipeline has finished SUCCESSFULLY",
- "dataflow_job_id": "test_dataflow_job_id",
- "project_id": "test_project_id",
- "location": "test_location",
},
)
assert f"{TASK_ID} completed with response Pipeline has finished SUCCESSFULLY" in caplog.text
@@ -192,7 +189,6 @@ class TestBeamRunPythonPipelineOperator:
py_interpreter=PY_INTERPRETER,
py_requirements=None,
py_system_site_packages=False,
- process_line_callback=None,
)
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
@@ -252,7 +248,7 @@ class TestBeamRunPythonPipelineOperator:
py_requirements=None,
py_system_site_packages=False,
process_line_callback=mock.ANY,
- check_job_status_callback=mock.ANY,
+ is_dataflow_job_id_exist_callback=mock.ANY,
)
dataflow_hook_mock.return_value.provide_authorized_gcloud.assert_called_once_with()
@@ -411,7 +407,6 @@ class TestBeamRunJavaPipelineOperator:
variables={**default_options, **pipeline_options},
jar=gcs_provide_file.return_value.__enter__.return_value.name,
job_class=JOB_CLASS,
- process_line_callback=None,
)
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
@@ -463,6 +458,7 @@ class TestBeamRunJavaPipelineOperator:
jar=gcs_provide_file.return_value.__enter__.return_value.name,
job_class=JOB_CLASS,
process_line_callback=mock.ANY,
+ is_dataflow_job_id_exist_callback=mock.ANY,
)
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
@@ -960,7 +956,8 @@ class TestBeamRunPythonPipelineOperatorAsync:
@mock.patch(BEAM_OPERATOR_PATH.format("BeamHook"))
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
- def test_exec_dataflow_runner(self, dataflow_hook_mock, beam_hook_mock):
+ @mock.patch(BEAM_OPERATOR_PATH.format("GCSHook"))
+ def test_exec_dataflow_runner(self, gcs_hook_mock, dataflow_hook_mock, beam_hook_mock):
"""
Test DataflowHook is created and the right args are passed to
start_python_dataflow when executing Dataflow runner.
@@ -984,7 +981,7 @@ class TestBeamRunPythonPipelineOperatorAsync:
cancel_timeout=dataflow_config.cancel_timeout,
wait_until_finished=dataflow_config.wait_until_finished,
)
- beam_hook_mock.return_value.start_python_pipeline.assert_not_called()
+ beam_hook_mock.return_value.start_python_pipeline.assert_called_once()
dataflow_hook_mock.return_value.provide_authorized_gcloud.assert_called_once_with()
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowJobLink.persist"))
@@ -1066,7 +1063,8 @@ class TestBeamRunJavaPipelineOperatorAsync:
@mock.patch(BEAM_OPERATOR_PATH.format("BeamHook"))
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
- def test_exec_dataflow_runner(self, dataflow_hook_mock, beam_hook_mock):
+ @mock.patch(BEAM_OPERATOR_PATH.format("GCSHook"))
+ def test_exec_dataflow_runner(self, gcs_hook_mock, dataflow_hook_mock, beam_hook_mock):
"""
Test DataflowHook is created and the right args are passed to
start_java_pipeline when executing Dataflow runner.
@@ -1075,6 +1073,7 @@ class TestBeamRunJavaPipelineOperatorAsync:
op = BeamRunJavaPipelineOperator(
runner="DataflowRunner", dataflow_config=dataflow_config, **self.default_op_kwargs
)
+ dataflow_hook_mock.return_value.is_job_dataflow_running.return_value = False
magic_mock = mock.MagicMock()
with pytest.raises(TaskDeferred):
op.execute(context=magic_mock)
@@ -1095,6 +1094,7 @@ class TestBeamRunJavaPipelineOperatorAsync:
@mock.patch(BEAM_OPERATOR_PATH.format("GCSHook"))
@mock.patch(BEAM_OPERATOR_PATH.format("DataflowHook"))
def test_on_kill_dataflow_runner(self, dataflow_hook_mock, _, __, ___):
+ dataflow_hook_mock.return_value.is_job_dataflow_running.return_value = False
dataflow_cancel_job = dataflow_hook_mock.return_value.cancel_job
op = BeamRunJavaPipelineOperator(runner="DataflowRunner", **self.default_op_kwargs)
with pytest.raises(TaskDeferred):
diff --git a/providers/apache/beam/tests/unit/apache/beam/triggers/test_beam.py b/providers/apache/beam/tests/unit/apache/beam/triggers/test_beam.py
index c1e1f65ce6b..5803f01dbd1 100644
--- a/providers/apache/beam/tests/unit/apache/beam/triggers/test_beam.py
+++ b/providers/apache/beam/tests/unit/apache/beam/triggers/test_beam.py
@@ -28,13 +28,7 @@ HOOK_STATUS_STR_JAVA = "airflow.providers.apache.beam.hooks.beam.BeamAsyncHook.s
CLASSPATH_PYTHON = "airflow.providers.apache.beam.triggers.beam.BeamPythonPipelineTrigger"
CLASSPATH_JAVA = "airflow.providers.apache.beam.triggers.beam.BeamJavaPipelineTrigger"
-TASK_ID = "test_task"
-LOCATION = "test-location"
-INSTANCE_NAME = "airflow-test-instance"
-INSTANCE = {"type": "BASIC", "displayName": INSTANCE_NAME}
-PROJECT_ID = "test_project_id"
TEST_GCP_CONN_ID = "test_gcp_conn_id"
-TEST_IMPERSONATION_CHAIN = "A,B,C"
TEST_VARIABLES = {"output": "gs://bucket_test/output", "labels": {"airflow-version": "v2-7-0-dev0"}}
TEST_PY_FILE = "apache_beam.examples.wordcount"
TEST_PY_OPTIONS: list[str] = []
@@ -46,10 +40,6 @@ TEST_JAR_FILE = "example.jar"
TEST_GCS_JAR_FILE = "gs://my-bucket/example/test.jar"
TEST_GCS_PY_FILE = "gs://my-bucket/my-object.py"
TEST_JOB_CLASS = "TestClass"
-TEST_CHECK_IF_RUNNING = False
-TEST_JOB_NAME = "test_job_name"
-TEST_POLL_SLEEP = 10
-TEST_CANCEL_TIMEOUT = 300
@pytest.fixture
@@ -61,8 +51,6 @@ def python_trigger():
py_interpreter=TEST_PY_INTERPRETER,
py_requirements=TEST_PY_REQUIREMENTS,
py_system_site_packages=TEST_PY_PACKAGES,
- project_id=PROJECT_ID,
- location=LOCATION,
runner=TEST_RUNNER,
gcp_conn_id=TEST_GCP_CONN_ID,
)
@@ -75,14 +63,7 @@ def java_trigger():
jar=TEST_JAR_FILE,
job_class=TEST_JOB_CLASS,
runner=TEST_RUNNER,
- check_if_running=TEST_CHECK_IF_RUNNING,
- project_id=PROJECT_ID,
- location=LOCATION,
- job_name=TEST_JOB_NAME,
gcp_conn_id=TEST_GCP_CONN_ID,
- impersonation_chain=TEST_IMPERSONATION_CHAIN,
- poll_sleep=TEST_POLL_SLEEP,
- cancel_timeout=TEST_CANCEL_TIMEOUT,
)
@@ -101,8 +82,6 @@ class TestBeamPythonPipelineTrigger:
"py_interpreter": TEST_PY_INTERPRETER,
"py_requirements": TEST_PY_REQUIREMENTS,
"py_system_site_packages": TEST_PY_PACKAGES,
- "project_id": PROJECT_ID,
- "location": LOCATION,
"runner": TEST_RUNNER,
"gcp_conn_id": TEST_GCP_CONN_ID,
}
@@ -123,9 +102,6 @@ class TestBeamPythonPipelineTrigger:
{
"status": "success",
"message": "Pipeline has finished SUCCESSFULLY",
- "dataflow_job_id": None,
- "project_id": PROJECT_ID,
- "location": LOCATION,
}
)
== actual
@@ -183,14 +159,7 @@ class TestBeamJavaPipelineTrigger:
"jar": TEST_JAR_FILE,
"job_class": TEST_JOB_CLASS,
"runner": TEST_RUNNER,
- "check_if_running": TEST_CHECK_IF_RUNNING,
- "project_id": PROJECT_ID,
- "location": LOCATION,
- "job_name": TEST_JOB_NAME,
"gcp_conn_id": TEST_GCP_CONN_ID,
- "impersonation_chain": TEST_IMPERSONATION_CHAIN,
- "poll_sleep": TEST_POLL_SLEEP,
- "cancel_timeout": TEST_CANCEL_TIMEOUT,
}
@pytest.mark.asyncio
@@ -209,9 +178,6 @@ class TestBeamJavaPipelineTrigger:
{
"status": "success",
"message": "Pipeline has finished SUCCESSFULLY",
- "dataflow_job_id": None,
- "project_id": PROJECT_ID,
- "location": LOCATION,
}
)
== actual
@@ -243,21 +209,6 @@ class TestBeamJavaPipelineTrigger:
actual = await generator.asend(None)
assert TriggerEvent({"status": "error", "message": "Test exception"}) == actual
- @pytest.mark.asyncio
- @mock.patch("airflow.providers.google.cloud.hooks.dataflow.AsyncDataflowHook.list_jobs")
- async def test_beam_trigger_exception_list_jobs_should_execute_successfully(
- self, mock_list_jobs, java_trigger
- ):
- """
- Test that BeamJavaPipelineTrigger fires the correct event in case of an error.
- """
- mock_list_jobs.side_effect = Exception("Test exception")
-
- java_trigger.check_if_running = True
- generator = java_trigger.run()
- actual = await generator.asend(None)
- assert TriggerEvent({"status": "error", "message": "Test exception"}) == actual
-
@pytest.mark.asyncio
@mock.patch("airflow.providers.apache.beam.triggers.beam.GCSHook")
async def test_beam_trigger_gcs_provide_file_should_execute_successfully(self, gcs_hook, java_trigger):
diff --git a/providers/google/tests/system/google/cloud/dataflow/example_dataflow_native_java.py b/providers/google/tests/system/google/cloud/dataflow/example_dataflow_native_java.py
index 66d0ce0995c..70efb000a6a 100644
--- a/providers/google/tests/system/google/cloud/dataflow/example_dataflow_native_java.py
+++ b/providers/google/tests/system/google/cloud/dataflow/example_dataflow_native_java.py
@@ -80,8 +80,8 @@ with DAG(
)
# [START howto_operator_start_java_job_local_jar]
- start_java_job_local = BeamRunJavaPipelineOperator(
- task_id="start_java_job_local",
+ start_java_job_direct = BeamRunJavaPipelineOperator(
+ task_id="start_java_job_direct",
jar=LOCAL_JAR,
pipeline_options={
"output": GCS_OUTPUT,
@@ -95,10 +95,25 @@ with DAG(
)
# [END howto_operator_start_java_job_local_jar]
+ start_java_job_direct_deferrable = BeamRunJavaPipelineOperator(
+ task_id="start_java_job_direct_deferrable",
+ jar=LOCAL_JAR,
+ pipeline_options={
+ "output": GCS_OUTPUT,
+ },
+ job_class="org.apache.beam.examples.WordCount",
+ dataflow_config={
+ "check_if_running": CheckJobRunning.WaitForRun,
+ "location": LOCATION,
+ "poll_sleep": 10,
+ },
+ deferrable=True,
+ )
+
# [START howto_operator_start_java_job_jar_on_gcs]
- start_java_job = BeamRunJavaPipelineOperator(
+ start_java_job_dataflow = BeamRunJavaPipelineOperator(
runner=BeamRunnerType.DataflowRunner,
- task_id="start_java_job",
+ task_id="start_java_job_dataflow",
jar=GCS_JAR,
pipeline_options={
"output": GCS_OUTPUT,
@@ -115,9 +130,9 @@ with DAG(
# [END howto_operator_start_java_job_jar_on_gcs]
# [START howto_operator_start_java_job_jar_on_gcs_deferrable]
- start_java_deferrable = BeamRunJavaPipelineOperator(
+ start_java_job_dataflow_deferrable = BeamRunJavaPipelineOperator(
runner=BeamRunnerType.DataflowRunner,
- task_id="start_java_job_deferrable",
+ task_id="start_java_job_dataflow_deferrable",
jar=GCS_JAR,
pipeline_options={
"output": GCS_OUTPUT,
@@ -143,7 +158,12 @@ with DAG(
create_bucket
>> download_file
# TEST BODY
- >> [start_java_job_local, start_java_job, start_java_deferrable]
+ >> [
+ start_java_job_direct,
+ start_java_job_direct_deferrable,
+ start_java_job_dataflow,
+ start_java_job_dataflow_deferrable,
+ ]
# TEST TEARDOWN
>> delete_bucket
)
diff --git a/providers/google/tests/system/google/cloud/dataflow/example_dataflow_native_python.py b/providers/google/tests/system/google/cloud/dataflow/example_dataflow_native_python.py
index cc90081bb14..29f34a562fd 100644
--- a/providers/google/tests/system/google/cloud/dataflow/example_dataflow_native_python.py
+++ b/providers/google/tests/system/google/cloud/dataflow/example_dataflow_native_python.py
@@ -60,9 +60,9 @@ with DAG(
create_bucket = GCSCreateBucketOperator(task_id="create_bucket", bucket_name=BUCKET_NAME)
# [START howto_operator_start_python_job]
- start_python_job = BeamRunPythonPipelineOperator(
+ start_python_job_dataflow = BeamRunPythonPipelineOperator(
runner=BeamRunnerType.DataflowRunner,
- task_id="start_python_job",
+ task_id="start_python_job_dataflow",
py_file=GCS_PYTHON_SCRIPT,
py_options=[],
pipeline_options={
@@ -75,8 +75,8 @@ with DAG(
)
# [END howto_operator_start_python_job]
- start_python_job_local = BeamRunPythonPipelineOperator(
- task_id="start_python_job_local",
+ start_python_job_direct = BeamRunPythonPipelineOperator(
+ task_id="start_python_job_direct",
py_file="apache_beam.examples.wordcount",
py_options=["-m"],
pipeline_options={
@@ -87,9 +87,22 @@ with DAG(
py_system_site_packages=False,
)
- start_python_deferrable = BeamRunPythonPipelineOperator(
+ start_python_job_direct_deferrable = BeamRunPythonPipelineOperator(
+ task_id="start_python_job_direct_deferrable",
+ py_file="apache_beam.examples.wordcount",
+ py_options=["-m"],
+ pipeline_options={
+ "output": GCS_OUTPUT,
+ },
+ py_requirements=["apache-beam[gcp]==2.59.0"],
+ py_interpreter="python3",
+ py_system_site_packages=False,
+ deferrable=True,
+ )
+
+ start_python_job_dataflow_deferrable = BeamRunPythonPipelineOperator(
runner=BeamRunnerType.DataflowRunner,
- task_id="start_python_job_deferrable",
+ task_id="start_python_job_dataflow_deferrable",
py_file=GCS_PYTHON_SCRIPT,
py_options=[],
pipeline_options={
@@ -118,7 +131,12 @@ with DAG(
# TEST SETUP
create_bucket
# TEST BODY
- >> [start_python_job, start_python_job_local, start_python_deferrable]
+ >> [
+ start_python_job_dataflow,
+ start_python_job_direct,
+ start_python_job_direct_deferrable,
+ start_python_job_dataflow_deferrable,
+ ]
>> stop_dataflow_job
# TEST TEARDOWN
>> delete_bucket