TobKed commented on a change in pull request #12814:
URL: https://github.com/apache/airflow/pull/12814#discussion_r567436289
##########
File path: airflow/providers/google/cloud/operators/dataflow.py
##########
@@ -43,6 +48,214 @@ class CheckJobRunning(Enum):
WaitForRun = 3
+class DataflowConfiguration(metaclass=ABCMeta):
+ """Abstract class for Dataflow configuration to be passed to Beam
operators"""
+
+ template_fields = ["job_name", "location"]
+
+ def __init__(
+ self,
+ *,
+ job_name: Optional[str] = "{{task.task_id}}",
+ append_job_name: bool = True,
+ project_id: Optional[str] = None,
+ location: Optional[str] = DEFAULT_DATAFLOW_LOCATION,
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None,
+ poll_sleep: int = 10,
+ impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
+ drain_pipeline: bool = False,
+ cancel_timeout: Optional[int] = 5 * 60,
+ wait_until_finished: Optional[bool] = None,
+ ) -> None:
+ self.job_name = job_name
+ self.append_job_name = append_job_name
+ self.project_id = project_id
+ self.location = location
+ self.gcp_conn_id = gcp_conn_id
+ self.delegate_to = delegate_to
+ self.poll_sleep = poll_sleep
+ self.impersonation_chain = impersonation_chain
+ self.drain_pipeline = drain_pipeline
+ self.cancel_timeout = cancel_timeout
+ self.wait_until_finished = wait_until_finished
+
+
+class DataflowPythonConfiguration(DataflowConfiguration):
+ """
+ Dataflow configuration that can be passed to
+
:py:class:`~airflow.providers.apache.beam.operators.beam.BeamRunPythonPipelineOperator`
+
+ :param job_name: The 'jobName' to use when executing the DataFlow job
+ (templated). This ends up being set in the pipeline options, so any
entry
+ with key ``'jobName'`` or ``'job_name'``in ``options`` will be
overwritten.
+ :type job_name: str
+ :param append_job_name: True if unique suffix has to be appended to job
name.
+ :type append_job_name: bool
+ :param project_id: Optional, the Google Cloud project ID in which to start
a job.
+ If set to None or missing, the default project_id from the Google
Cloud connection is used.
+ :type project_id: str
+ :param location: Job location.
+ :type location: str
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate using domain-wide
delegation of authority,
+ if any. For this to work, the service account making the request must
have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ :param poll_sleep: The time in seconds to sleep between polling Google
+ Cloud Platform for the dataflow job status while the job is in the
+ JOB_STATE_RUNNING state.
+ :type poll_sleep: int
+ :param impersonation_chain: Optional 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).
+ :type impersonation_chain: Union[str, Sequence[str]]
+ :param drain_pipeline: Optional, set to True if want to stop streaming job
by draining it
+ instead of canceling during during killing task instance. See:
+ https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
+ :type drain_pipeline: bool
+ :param cancel_timeout: How long (in seconds) operator should wait for the
pipeline to be
+ successfully cancelled when task is being killed.
+ :type cancel_timeout: Optional[int]
+ :param wait_until_finished: (Optional)
+ If True, wait for the end of pipeline execution before exiting.
+ If False, only submits job.
+ If None, default behavior.
+
+ The default behavior depends on the type of pipeline:
+
+ * for the streaming pipeline, wait for jobs to start,
+ * for the batch pipeline, wait for the jobs to complete.
+
+ .. warning::
+
+ You cannot call ``PipelineResult.wait_until_finish`` method in
your pipeline code for the operator
+ to work properly. i. e. you must use asynchronous execution.
Otherwise, your pipeline will
+ always wait until finished. For more information, look at:
+ `Asynchronous execution
+
<https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#python_10>`__
+
+ The process of starting the Dataflow job in Airflow consists of two
steps:
+
+ * running a subprocess and reading the stderr/stderr log for the job
id.
+ * loop waiting for the end of the job ID from the previous step.
+ This loop checks the status of the job.
+
+ Step two is started just after step one has finished, so if you have
wait_until_finished in your
+ pipeline code, step two will not start until the process stops. When
this process stops,
+ steps two will run, but it will only execute one iteration as the job
will be in a terminal state.
+
+ If you in your pipeline do not call the wait_for_pipeline method but
pass wait_until_finish=True
+ to the operator, the second loop will wait for the job's terminal
state.
+
+ If you in your pipeline do not call the wait_for_pipeline method, and
pass wait_until_finish=False
+ to the operator, the second loop will check once is job not in
terminal state and exit the loop.
+ :type wait_until_finished: Optional[bool]
+ """
+
+
+class DataflowJavaConfiguration(DataflowConfiguration):
+ """
+ Dataflow configuration that can be passed to
+
:py:class:`~airflow.providers.apache.beam.operators.beam.BeamRunJavaPipelineOperator`
+
+ :param job_name: The 'jobName' to use when executing the DataFlow job
+ (templated). This ends up being set in the pipeline options, so any
entry
+ with key ``'jobName'`` or ``'job_name'``in ``options`` will be
overwritten.
+ :type job_name: str
+ :param append_job_name: True if unique suffix has to be appended to job
name.
+ :type append_job_name: bool
+ :param project_id: Optional, the Google Cloud project ID in which to start
a job.
+ If set to None or missing, the default project_id from the Google
Cloud connection is used.
+ :type project_id: str
+ :param location: Job location.
+ :type location: str
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate using domain-wide
delegation of authority,
+ if any. For this to work, the service account making the request must
have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ :param poll_sleep: The time in seconds to sleep between polling Google
+ Cloud Platform for the dataflow job status while the job is in the
+ JOB_STATE_RUNNING state.
+ :type poll_sleep: int
+ :param impersonation_chain: Optional 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).
+ :type impersonation_chain: Union[str, Sequence[str]]
+ :param drain_pipeline: Optional, set to True if want to stop streaming job
by draining it
+ instead of canceling during during killing task instance. See:
+ https://cloud.google.com/dataflow/docs/guides/stopping-a-pipeline
+ :type drain_pipeline: bool
+ :param cancel_timeout: How long (in seconds) operator should wait for the
pipeline to be
+ successfully cancelled when task is being killed.
+ :type cancel_timeout: Optional[int]
+ :param wait_until_finished: (Optional)
+ If True, wait for the end of pipeline execution before exiting.
+ If False, only submits job.
+ If None, default behavior.
+
+ The default behavior depends on the type of pipeline:
+
+ * for the streaming pipeline, wait for jobs to start,
+ * for the batch pipeline, wait for the jobs to complete.
+
+ .. warning::
+
+ You cannot call ``PipelineResult.wait_until_finish`` method in
your pipeline code for the operator
+ to work properly. i. e. you must use asynchronous execution.
Otherwise, your pipeline will
+ always wait until finished. For more information, look at:
+ `Asynchronous execution
+
<https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#python_10>`__
+
+ The process of starting the Dataflow job in Airflow consists of two
steps:
+
+ * running a subprocess and reading the stderr/stderr log for the job
id.
+ * loop waiting for the end of the job ID from the previous step.
+ This loop checks the status of the job.
+
+ Step two is started just after step one has finished, so if you have
wait_until_finished in your
+ pipeline code, step two will not start until the process stops. When
this process stops,
+ steps two will run, but it will only execute one iteration as the job
will be in a terminal state.
+
+ If you in your pipeline do not call the wait_for_pipeline method but
pass wait_until_finish=True
+ to the operator, the second loop will wait for the job's terminal
state.
+
+ If you in your pipeline do not call the wait_for_pipeline method, and
pass wait_until_finish=False
+ to the operator, the second loop will check once is job not in
terminal state and exit the loop.
+ :type wait_until_finished: Optional[bool]
+ :param multiple_jobs: If pipeline creates multiple jobs then monitor all
jobs
+ :type multiple_jobs: boolean
+ :param check_if_running: before running job, validate that a previous run
is not in process
+ :type check_if_running: CheckJobRunning(IgnoreJob = do not check if
running, FinishIfRunning=
+ if job is running finish with nothing, WaitForRun= wait until job
finished and the run job)
+ ``jar``, ``options``, and ``job_name`` are templated so you can use
variables in them.
+ """
+
+ def __init__(
+ self,
+ *,
+ multiple_jobs: Optional[bool] = None,
+ check_if_running: CheckJobRunning = CheckJobRunning.WaitForRun,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.multiple_jobs = multiple_jobs
+ self.check_if_running = check_if_running
+
+
# pylint: disable=too-many-instance-attributes
class DataflowCreateJavaJobOperator(BaseOperator):
Review comment:
I added deprecation warnings in docstrings and logs.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
