mik-laj commented on a change in pull request #8550:
URL: https://github.com/apache/airflow/pull/8550#discussion_r416185412
##########
File path: airflow/providers/google/cloud/operators/dataflow.py
##########
@@ -406,6 +406,71 @@ def on_kill(self) -> None:
self.hook.cancel_job(job_id=self.job_id,
project_id=self.project_id)
+class DataflowStartFlexTemplateOperator(BaseOperator):
+ """
+ Starts flex templates with the Dataflow pipeline.
+
+ :param body: The request body
+ :param location: The location of the Dataflow job (for example
europe-west1)
+ :type location: str
+ :param project_id: The ID of the GCP project that owns the job.
+ If set to ``None`` or missing, the default project_id from the GCP
connection is used.
+ :type project_id: Optional[str]
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud
+ Platform.
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ template_fields = ["body", 'location', 'project_id', 'gcp_conn_id']
+
+ @apply_defaults
+ def __init__(
+ self,
+ body: Dict,
+ location: str,
+ project_id: Optional[str] = None,
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None,
+ *args,
+ **kwargs
+ ) -> None:
+ super().__init__(*args, **kwargs)
+ self.body = body
+ self.location = location
+ self.project_id = project_id
+ self.gcp_conn_id = gcp_conn_id
+ self.delegate_to = delegate_to
+ self.job_id = None
+ self.hook: Optional[DataflowHook] = None
+
+ def execute(self, context):
+ self.hook = DataflowHook(
+ gcp_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to,
+ )
+
+ def set_current_job_id(job_id):
+ self.job_id = job_id
+
+ job = self.hook.start_flex_template(
+ body=self.body,
+ location=self.location,
+ project_id=self.project_id,
+ on_new_job_id_callback=set_current_job_id,
+ )
+
+ return job
+
+ def on_kill(self) -> None:
+ self.log.info("On kill.")
+ if self.job_id:
+ self.hook.cancel_job(job_id=self.job_id,
project_id=self.project_id)
Review comment:
Good point. I will change it
----------------------------------------------------------------
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]
