Repository: apache/airflow
File: airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py (Python)

Code before:
    elif isinstance(ti_patch_payload, TIDeferredStatePayload):
        # Calculate timeout if it was passed
        timeout = None
        if ti_patch_payload.trigger_timeout is not None:
            timeout = timezone.utcnow() + ti_patch_payload.trigger_timeout

        trigger_kwargs = ti_patch_payload.trigger_kwargs
        if not isinstance(trigger_kwargs, str):
            # If it's passed as a string, assume the client encrypted it, otherwise assume it doesn't need to
            # be. Just JSON serialize it
            trigger_kwargs = json.dumps(trigger_kwargs)

        trigger_row = Trigger(
            classpath=ti_patch_payload.classpath,
            kwargs={},
        )
        trigger_row.encrypted_kwargs = trigger_kwargs
        session.add(trigger_row)
        session.flush()

        # TODO: HANDLE execution timeout later as it requires a call to the DB
        # either get it from the serialised DAG or get it from the API

        query = update(TI).where(TI.id == ti_id_str)

        # Store next_kwargs directly (already serialized by worker)
        query = query.values(
            state=TaskInstanceState.DEFERRED,
            trigger_id=trigger_row.id,
            next_method=ti_patch_payload.next_method,
            next_kwargs=ti_patch_payload.next_kwargs,
            trigger_timeout=timeout,
        )
        updated_state = TaskInstanceState.DEFERRED
    elif isinstance(ti_patch_payload, TIRescheduleStatePayload):
        # Quick check for poke_interval isn't immediately over MySQL's TIMESTAMP limit.
        # This check is only rudimentary to catch trivial user errors, e.g. mistakenly
        # set the value to milliseconds instead of seconds. There's another check when
        # we actually try to reschedule to ensure database coherence.
        if get_dialect_name(session) == "mysql":
            # As documented in https://dev.mysql.com/doc/refman/5.7/en/datetime.html.
            _MYSQL_TIMESTAMP_MAX = timezone.datetime(2038, 1, 19, 3, 14, 7)
            if ti_patch_payload.reschedule_date > _MYSQL_TIMESTAMP_MAX:
                # Set a task to failed in case any unexpected exception happened during task state update
                log.error(
                    "Cannot reschedule task past MySQL limit. Setting the task to failed.",
                    payload=ti_patch_payload,
                    mysql_timestamp_max=_MYSQL_TIMESTAMP_MAX,
                )
                data = ti_patch_payload.model_dump(exclude={"reschedule_date"}, exclude_unset=True)
                query = update(TI).where(TI.id == ti_id_str).values(data)

Review comment (suggestion):
We should probably change the type on `next_kwargs` to JsonValue

Code after:

        trigger_kwargs = ti_patch_payload.trigger_kwargs
        if not isinstance(trigger_kwargs, str):
            # If it's passed as a string, assume the client encrypted it, otherwise assume it doesn't need to
            # be. Just JSON serialize it
            trigger_kwargs = json.dumps(trigger_kwargs)

        trigger_row = Trigger(
            classpath=ti_patch_payload.classpath,
            kwargs={},
        )
        trigger_row.encrypted_kwargs = trigger_kwargs
        session.add(trigger_row)
        session.flush()

        # TODO: HANDLE execution timeout later as it requires a call to the DB
        # either get it from the serialised DAG or get it from the API

        query = update(TI).where(TI.id == ti_id_str)

        # Store next_kwargs directly (already serialized by worker)
        query = query.values(
            state=TaskInstanceState.DEFERRED,
            trigger_id=trigger_row.id,
            next_method=ti_patch_payload.next_method,
            next_kwargs=ti_patch_payload.next_kwargs,
            trigger_timeout=timeout,
        )
        updated_state = TaskInstanceState.DEFERRED
    elif isinstance(ti_patch_payload, TIRescheduleStatePayload):
        # Quick check for poke_interval isn't immediately over MySQL's TIMESTAMP limit.
        # This check is only rudimentary to catch trivial user errors, e.g. mistakenly
        # set the value to milliseconds instead of seconds. There's another check when
        # we actually try to reschedule to ensure database coherence.
        if get_dialect_name(session) == "mysql":
            # As documented in https://dev.mysql.com/doc/refman/5.7/en/datetime.html.
            _MYSQL_TIMESTAMP_MAX = timezone.datetime(2038, 1, 19, 3, 14, 7)
            if ti_patch_payload.reschedule_date > _MYSQL_TIMESTAMP_MAX:
                # Set a task to failed in case any unexpected exception happened during task state update
                log.error(
                    "Cannot reschedule task past MySQL limit. Setting the task to failed.",
                    payload=ti_patch_payload,
                    mysql_timestamp_max=_MYSQL_TIMESTAMP_MAX,
                )
                data = ti_patch_payload.model_dump(exclude={"reschedule_date"}, exclude_unset=True)
                query = update(TI).where(TI.id == ti_id_str).values(data)
                if session.bind is not None:
                    query = TI.duration_expression_update(timezone.utcnow(), query, session.bind)
                query = query.values(state=TaskInstanceState.FAILED)
                # We skip fail_fast handling in this error case to avoid fetching the TI object while the row
                # is still locked from the earlier with_for_update() query, which might cause deadlock issues