Repository: apache/airflow
File: airflow-core/src/airflow/models/trigger.py (Python)

Code before:
            locked_query = with_row_locks(query.limit(remaining_capacity), session, skip_locked=True)
            result.extend(session.execute(locked_query).all())

        return result


@singledispatch
def handle_event_submit(event: TriggerEvent, *, task_instance: TaskInstance, session: Session) -> None:
    """
    Handle the submit event for a given task instance.

    This function sets the next method and next kwargs of the task instance,
    as well as its state to scheduled. It also adds the event's payload
    into the kwargs for the task.

    :param task_instance: The task instance to handle the submit event for.
    :param session: The session to be used for the database callback sink.
    """
    from airflow.sdk.serde import serialize
    from airflow.utils.state import TaskInstanceState

    # Get the next kwargs of the task instance, or an empty dictionary if it doesn't exist
    next_kwargs = task_instance.next_kwargs or {}

    # Update the next kwargs of the task instance
    next_kwargs["event"] = serialize(event.payload)
    task_instance.next_kwargs = next_kwargs

    # Remove ourselves as its trigger
    task_instance.trigger_id = None

    # Set the state of the task instance to scheduled
    task_instance.state = TaskInstanceState.SCHEDULED
    task_instance.scheduled_dttm = timezone.utcnow()
    session.flush()


@handle_event_submit.register
def _(event: BaseTaskEndEvent, *, task_instance: TaskInstance, session: Session) -> None:
    """
    Submit event for the given task instance.

    Marks the task with the state `task_instance_state` and optionally pushes xcom if applicable.

    :param task_instance: The task instance to be submitted.
    :param session: The session to be used for the database callback sink.
    """
    from airflow.callbacks.callback_requests import TaskCallbackRequest
    from airflow.callbacks.database_callback_sink import DatabaseCallbackSink
    from airflow.utils.state import TaskInstanceState


Review comment (bug):
If I read the path correctly, it now serializes `event.payload` with serde before adding it to `next_kwargs`. However, `next_kwargs` might already contain data in the old BaseSerialization format (for tasks deferred before upgrade)?

Example scenario I am thinking that is worth validating:

1. Task deferred before upgrade has next_kwargs with keys in BaseSerialization format
2. Trigger fires after upgrade, adds "event" key in serde format to the dict
3. Worker tries to deserialize - gets mixed format dict
4. Some keys deserialize correctly, others fail or return wrong types

---

maybe you might want to do the same compat shim? Unsure but worth validating

```py
    try:
        next_kwargs = deserialize(next_kwargs_raw) if next_kwargs_raw else {}
    except (ImportError, KeyError, AttributeError, TypeError):
        from airflow.serialization.serialized_objects import BaseSerialization
        next_kwargs = BaseSerialization.deserialize(next_kwargs_raw) if next_kwargs_raw else {}
    
```

Code after:
            locked_query = with_row_locks(query.limit(remaining_capacity), session, skip_locked=True)
            result.extend(session.execute(locked_query).all())

        return result


@singledispatch
def handle_event_submit(event: TriggerEvent, *, task_instance: TaskInstance, session: Session) -> None:
    """
    Handle the submit event for a given task instance.

    This function sets the next method and next kwargs of the task instance,
    as well as its state to scheduled. It also adds the event's payload
    into the kwargs for the task.

    :param task_instance: The task instance to handle the submit event for.
    :param session: The session to be used for the database callback sink.
    """
    from airflow.sdk.serde import deserialize, serialize
    from airflow.utils.state import TaskInstanceState

    next_kwargs_raw = task_instance.next_kwargs or {}

    # deserialize first to provide a compat layer if there are mixed serialized (BaseSerialisation and serde) data
    # which can happen if a deferred task resumes after upgrade
    try:
        next_kwargs = deserialize(next_kwargs_raw)
    except (ImportError, KeyError, AttributeError, TypeError):
        from airflow.serialization.serialized_objects import BaseSerialization

        next_kwargs = BaseSerialization.deserialize(next_kwargs_raw)

    # Add event to the plain dict, then serialize everything together. This ensures that the event is properly
    # nested inside __var__ in the final serde serialized structure.
    if TYPE_CHECKING:
        assert isinstance(next_kwargs, dict)
    next_kwargs["event"] = event.payload

    # re-serialize the entire dict using serde to ensure consistent structure
    task_instance.next_kwargs = serialize(next_kwargs)

    # Remove ourselves as its trigger
    task_instance.trigger_id = None

    # Set the state of the task instance to scheduled
    task_instance.state = TaskInstanceState.SCHEDULED
    task_instance.scheduled_dttm = timezone.utcnow()
    session.flush()


@handle_event_submit.register