sunank200 commented on code in PR #39230:
URL: https://github.com/apache/airflow/pull/39230#discussion_r1578101930
##########
airflow/providers/google/cloud/triggers/dataproc.py:
##########
@@ -91,20 +93,28 @@ def serialize(self):
"gcp_conn_id": self.gcp_conn_id,
"impersonation_chain": self.impersonation_chain,
"polling_interval_seconds": self.polling_interval_seconds,
+ "cancel_on_kill": self.cancel_on_kill,
},
)
async def run(self):
- while True:
- job = await self.get_async_hook().get_job(
- project_id=self.project_id, region=self.region,
job_id=self.job_id
- )
- state = job.status.state
- self.log.info("Dataproc job: %s is in state: %s", self.job_id,
state)
- if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED,
JobStatus.State.ERROR):
- break
- await asyncio.sleep(self.polling_interval_seconds)
- yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job":
job})
+ try:
+ while True:
+ job = await self.get_async_hook().get_job(
+ project_id=self.project_id, region=self.region,
job_id=self.job_id
+ )
+ state = job.status.state
+ self.log.info("Dataproc job: %s is in state: %s", self.job_id,
state)
+ if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED,
JobStatus.State.ERROR):
+ break
+ await asyncio.sleep(self.polling_interval_seconds)
+ yield TriggerEvent({"job_id": self.job_id, "job_state": state,
"job": job})
+ except asyncio.CancelledError:
+ self.log.info("Task got cancelled.")
+ if self.job_id and self.cancel_on_kill:
+ await self.get_async_hook().cancel_job(
Review Comment:
Yes. I tested this
##########
airflow/providers/google/cloud/triggers/dataproc.py:
##########
@@ -91,20 +93,28 @@ def serialize(self):
"gcp_conn_id": self.gcp_conn_id,
"impersonation_chain": self.impersonation_chain,
"polling_interval_seconds": self.polling_interval_seconds,
+ "cancel_on_kill": self.cancel_on_kill,
},
)
async def run(self):
- while True:
- job = await self.get_async_hook().get_job(
- project_id=self.project_id, region=self.region,
job_id=self.job_id
- )
- state = job.status.state
- self.log.info("Dataproc job: %s is in state: %s", self.job_id,
state)
- if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED,
JobStatus.State.ERROR):
- break
- await asyncio.sleep(self.polling_interval_seconds)
- yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job":
job})
+ try:
+ while True:
+ job = await self.get_async_hook().get_job(
+ project_id=self.project_id, region=self.region,
job_id=self.job_id
+ )
+ state = job.status.state
+ self.log.info("Dataproc job: %s is in state: %s", self.job_id,
state)
+ if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED,
JobStatus.State.ERROR):
+ break
+ await asyncio.sleep(self.polling_interval_seconds)
+ yield TriggerEvent({"job_id": self.job_id, "job_state": state,
"job": job})
+ except asyncio.CancelledError:
+ self.log.info("Task got cancelled.")
+ if self.job_id and self.cancel_on_kill:
+ await self.get_async_hook().cancel_job(
Review Comment:
Yes i tested this
--
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.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
