ashb commented on code in PR #74222:
URL: https://github.com/apache/airflow/pull/74222#discussion_r4182327750
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -1071,28 +1122,98 @@ def is_premature(self) -> bool:
# is the task still in the retry waiting period?
return self.state == TaskInstanceState.UP_FOR_RETRY and not
self.ready_for_retry()
- def prepare_db_for_next_try(self, session: Session):
- """Archive this attempt and allocate the next attempt's UUID and try
number."""
- from airflow.models.taskinstancehistory import TaskInstanceHistory
+ def retire(self, *, reason: str, session: Session) -> None:
+ """Remove this attempt from the working set while retaining its UUID
and children."""
+ current = session.scalar(
+ select(TaskInstance.working_set).where(TaskInstance.id ==
self.id).with_for_update()
+ )
+ if current is not True:
+ raise ValueError("A retired task instance cannot be retired again")
+ if self.state not in State.finished:
+ self.state = TaskInstanceState.FAILED
+ if self.end_date is None:
+ self.end_date = timezone.utcnow()
+ self.set_duration()
+ self.working_set = None
+ self.archived_reason = reason
+ self.trigger_id = None
+ session.flush()
- TaskInstanceHistory.record_ti(self, session=session)
- session.execute(delete(TaskReschedule).filter_by(ti_id=self.id))
- self.external_executor_id = None
- self.id = uuid7()
- self.try_number += 1
+ @classmethod
+ def delete_attempts(
+ cls,
+ *,
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ map_index: int | None = None,
+ session: Session,
+ ) -> None:
+ """Delete every attempt, current and retired, of a task; all map
indexes if ``map_index`` is None."""
+ statement = delete(cls).where(cls.dag_id == dag_id, cls.run_id ==
run_id, cls.task_id == task_id)
+ if map_index is not None:
+ statement = statement.where(cls.map_index == map_index)
+ session.execute(statement)
+
+ @classmethod
+ def get_last_try_numbers(
+ cls,
+ *,
+ dag_id: str,
+ task_id: str,
+ run_id: str,
+ map_indexes: Collection[int] | None = None,
+ session: Session,
+ ) -> dict[int, int]:
+ """Return the highest try number per map index across current and
historical task instances."""
+ if map_indexes is not None and not map_indexes:
+ return {}
+ statement = (
+ select(cls.map_index, func.max(cls.try_number))
+ .where(cls.dag_id == dag_id, cls.task_id == task_id, cls.run_id ==
run_id)
+ .group_by(cls.map_index)
+ )
+ if map_indexes is not None:
+ statement = statement.where(cls.map_index.in_(map_indexes))
+ return {map_index: last_try for map_index, last_try in
session.execute(statement)}
+
+ def prepare_db_for_next_try(self, session: Session) -> TaskInstance:
Review Comment:
Yes, true. We absolutely can fix this though, just trying to think if this
matters enough for it to be in this PR or can be a follow up (and _how_ we
would fix it without a change to the API)
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
