Re: [PR] Add revoke_task implementation to EdgeExecutor for task queued timeout support [airflow]

["dheerajturaga (via GitHub)" <gi...@apache.org>]

dheerajturaga commented on code in PR #56240:
URL: https://github.com/apache/airflow/pull/56240#discussion_r2393428867


##########
providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py:
##########
@@ -399,6 +399,36 @@ def end(self) -> None:
     def terminate(self):
         """Terminate the executor is not doing anything."""
 
+    @provide_session
+    def revoke_task(self, *, ti: TaskInstance, session: Session = NEW_SESSION):
+        """
+        Revoke a task instance from the executor.
+
+        This method removes the task from the executor's internal state and deletes
+        the corresponding EdgeJobModel record to prevent edge workers from picking it up.
+
+        :param ti: Task instance to revoke
+        :param session: Database session
+        """
+        # Remove from executor's internal state
+        self.running.discard(ti.key)
+        self.queued_tasks.pop(ti.key, None)
+        if ti.key in self.last_reported_state:
+            del self.last_reported_state[ti.key]
+
+        # Delete the job from the database to prevent edge workers from picking it up
+        session.execute(
+            delete(EdgeJobModel).where(
+                EdgeJobModel.dag_id == ti.dag_id,
+                EdgeJobModel.task_id == ti.task_id,
+                EdgeJobModel.run_id == ti.run_id,
+                EdgeJobModel.map_index == ti.map_index,
+                EdgeJobModel.try_number == ti.try_number,
+            )
+        )
+        session.commit()

Review Comment:
   Done!



-- 
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-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org