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) |
|
) |
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!
