Skip to content

Refactor BaseSensorOperator to support Reschedule mode with Task SDK #45580

Closed
@kaxil

Description

@kaxil

BaseSensorOperator has some logic to use max_tries from Context as well as access DB directly. This won't work with Task SDK and as such DAGs like example_sensor_decorator will fail.

if self.reschedule:
ti = context["ti"]
max_tries: int = ti.max_tries or 0
retries: int = self.retries or 0
# If reschedule, use the start date of the first try (first try can be either the very
# first execution of the task, or the first execution after the task was cleared).
# If the first try's record was not saved due to the Exception occurred and the following
# transaction rollback, the next available attempt should be taken
# to prevent falling in the endless rescheduling
first_try_number = max_tries - retries + 1
with create_session() as session:
start_date = session.scalar(
select(TaskReschedule)
.where(
TaskReschedule.dag_id == ti.dag_id,
TaskReschedule.task_id == ti.task_id,
TaskReschedule.run_id == ti.run_id,
TaskReschedule.map_index == ti.map_index,
TaskReschedule.try_number >= first_try_number,
)
.order_by(TaskReschedule.id.asc())
.with_only_columns(TaskReschedule.start_date)
.limit(1)
)

Image Image

The max_tries is easier since we can pass it we can in the following code and re-generate the API client.

class TaskInstance(BaseModel):
"""Schema for TaskInstance with minimal required fields needed for Executors and Task SDK."""
id: uuid.UUID
task_id: str
dag_id: str
run_id: str
try_number: int
map_index: int | None = None

Replacing direct DB access in that BaseSensorOperator, on the other hand, needs to be figured out!

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions