[v3-3-test] Add the option to select bundle version parameter on dag run trigger endpoint (#61550)

airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py

@@ -228,6 +228,7 @@ class TriggerDAGRunPostBody(StrictBaseModel):
     conf: dict | None = Field(default_factory=dict)
     note: str | None = None
     partition_key: str | None = None
+    bundle_version: str | None = None
 
     @model_validator(mode="after")
     def check_data_intervals(self):

airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml

@@ -15110,6 +15110,11 @@ components:
           - type: string
           - type: 'null'
           title: Partition Key
+        bundle_version:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Bundle Version
       additionalProperties: false
       type: object
       title: MaterializeAssetBody
@@ -16643,6 +16648,11 @@ components:
           - type: string
           - type: 'null'
           title: Partition Key
+        bundle_version:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Bundle Version
       additionalProperties: false
       type: object
       required:

airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py

@@ -73,7 +73,7 @@
 from airflow.api_fastapi.logging.decorators import action_logging
 from airflow.assets.manager import asset_manager
 from airflow.configuration import conf
-from airflow.exceptions import ParamValidationError
+from airflow.exceptions import DagVersionNotFound, ParamValidationError
 from airflow.models.asset import (
     AssetAliasModel,
     AssetDagRunQueue,
@@ -82,6 +82,7 @@
     AssetWatcherModel,
     TaskOutletAssetReference,
 )
+from airflow.models.dag_version import DagVersion
 from airflow.typing_compat import Unpack
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunTriggeredByType, DagRunType
@@ -445,14 +446,31 @@ def materialize_asset(
 
     dag = get_latest_version_of_dag(dag_bag, dag_id, session)
 
-    if dag.allowed_run_types is not None and DagRunType.ASSET_MATERIALIZATION not in dag.allowed_run_types:
-        raise HTTPException(
-            status.HTTP_400_BAD_REQUEST,
-            f"Dag with dag_id: '{dag_id}' does not allow asset materialization runs",
-        )
+    resolved_body = body or MaterializeAssetBody()
 
     try:
-        params = (body or MaterializeAssetBody()).validate_context(dag)
+        preloaded_dag_version = None
+        context_dag = dag
+        if resolved_body.bundle_version is not None and not dag.disable_bundle_versioning:
+            preloaded_dag_version = DagVersion.get_latest_version(
+                dag_id, bundle_version=resolved_body.bundle_version, load_serialized_dag=True, session=session
+            )
+            if not preloaded_dag_version:
+                raise DagVersionNotFound(
+                    f"DAG with dag_id: '{dag_id}' does not have a version for bundle_version '{resolved_body.bundle_version}'"
+                )
+            context_dag = preloaded_dag_version.serialized_dag.dag
+
+        if (
+            context_dag.allowed_run_types is not None
+            and DagRunType.ASSET_MATERIALIZATION not in context_dag.allowed_run_types
+        ):
+            raise HTTPException(
+                status.HTTP_400_BAD_REQUEST,
+                f"Dag with dag_id: '{dag_id}' does not allow asset materialization runs",
+            )
+
+        params = resolved_body.validate_context(context_dag)
         return dag.create_dagrun(
             run_id=params["run_id"],
             logical_date=params["logical_date"],
@@ -467,9 +485,13 @@ def materialize_asset(
             partition_date=params["partition_date"],
             note=params["note"],
             session=session,
+            bundle_version=resolved_body.bundle_version,
+            dag_version=preloaded_dag_version,
         )
     except (ParamValidationError, ValueError) as e:
         raise HTTPException(status.HTTP_400_BAD_REQUEST, str(e)) from e
+    except DagVersionNotFound as e:
+        raise HTTPException(status.HTTP_404_NOT_FOUND, str(e)) from e
 
 
 @assets_router.get(

airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py

@@ -108,7 +108,7 @@
     perform_clear_dag_run,
 )
 from airflow.api_fastapi.logging.decorators import action_logging
-from airflow.exceptions import ParamValidationError
+from airflow.exceptions import DagVersionNotFound, ParamValidationError
 from airflow.models import DagModel, DagRun
 from airflow.models.asset import AssetEvent
 from airflow.models.dag_version import DagVersion
@@ -690,21 +690,44 @@ def trigger_dag_run(
             f"Dag with dag_id: '{dag_id}' has import errors and cannot be triggered",
         )
 
-    if dm.allowed_run_types is not None and DagRunType.MANUAL not in dm.allowed_run_types:
-        raise HTTPException(
-            status.HTTP_400_BAD_REQUEST,
-            f"Dag with dag_id: '{dag_id}' does not allow manual runs",
-        )
-
     referer = request.headers.get("referer")
     if referer:
         triggered_by = DagRunTriggeredByType.UI
     else:
         triggered_by = DagRunTriggeredByType.REST_API
 
-    dag = get_latest_version_of_dag(dag_bag, dag_id, session)
     try:
-        params = body.validate_context(dag)
+        dag = get_latest_version_of_dag(dag_bag, dag_id, session)
+        preloaded_dag_version = None
+        context_dag = dag
+        if body.bundle_version is not None and not dag.disable_bundle_versioning:
+            preloaded_dag_version = DagVersion.get_latest_version(
+                dag_id, bundle_version=body.bundle_version, load_serialized_dag=True, session=session
+            )
+            if not preloaded_dag_version:
+                raise DagVersionNotFound(
+                    f"DAG with dag_id: '{dag_id}' does not have a version for bundle_version '{body.bundle_version}'"
+                )
+            context_dag = preloaded_dag_version.serialized_dag.dag
+
+        if (
+            context_dag.allowed_run_types is not None
+            and DagRunType.MANUAL not in context_dag.allowed_run_types
+        ):
+            raise HTTPException(
+                status.HTTP_400_BAD_REQUEST,
+                f"Dag with dag_id: '{dag_id}' does not allow manual runs",
+            )
+
+        params = body.validate_context(context_dag)
+
+        if body.bundle_version is not None:
+            if dag.disable_bundle_versioning:
+                raise HTTPException(
+                    status.HTTP_400_BAD_REQUEST,
+                    f"DAG with dag_id: '{dag_id}' does not support bundle versioning",
+                )
+
         dag_run = dag.create_dagrun(
             run_id=params["run_id"],
             logical_date=params["logical_date"],
@@ -716,17 +739,22 @@ def trigger_dag_run(
             triggering_user_name=user.get_name(),
             state=DagRunState.QUEUED,
             partition_key=params["partition_key"],
+            bundle_version=body.bundle_version,
+            dag_version=preloaded_dag_version,
             partition_date=params["partition_date"],
             session=session,
         )
+
+        dag_run_note = body.note
+        if dag_run_note:
+            current_user_id = user.get_id()
+            dag_run.note = (dag_run_note, current_user_id)
+        return dag_run
+
     except (ParamValidationError, ValueError) as e:
         raise HTTPException(status.HTTP_400_BAD_REQUEST, str(e)) from e
-
-    dag_run_note = body.note
-    if dag_run_note:
-        current_user_id = user.get_id()
-        dag_run.note = (dag_run_note, current_user_id)
-    return dag_run
+    except DagVersionNotFound as e:
+        raise HTTPException(status.HTTP_404_NOT_FOUND, str(e)) from e
 
 
 @dag_run_router.get(

airflow-core/src/airflow/exceptions.py

@@ -141,6 +141,10 @@ class DagRunNotFound(AirflowNotFoundException):
     """Raise when a DAG Run is not available in the system."""
 
 
+class DagVersionNotFound(AirflowNotFoundException):
+    """Raised when a DagVersion for the given dag_id / bundle_version is not found."""
+
+
 class DagNotPartitionedError(ValueError):
     """Raise when a partition_key is supplied for a Dag that is not partitioned."""
 

airflow-core/src/airflow/models/dag_version.py

@@ -154,6 +154,7 @@ def _latest_version_select(
         bundle_version: str | None = None,
         load_dag_model: bool = False,
         load_bundle_model: bool = False,
+        load_serialized_dag: bool = False,
     ) -> Select:
         """
         Get the select object to get the latest version of the DAG.
@@ -162,7 +163,7 @@ def _latest_version_select(
         :return: The select object.
         """
         query = select(cls).where(cls.dag_id == dag_id)
-        if bundle_version:
+        if bundle_version is not None:
             query = query.where(cls.bundle_version == bundle_version)
 
         if load_dag_model:
@@ -171,6 +172,9 @@ def _latest_version_select(
         if load_bundle_model:
             query = query.options(joinedload(cls.bundle))
 
+        if load_serialized_dag:
+            query = query.options(joinedload(cls.serialized_dag))
+
         # Order by version_number, not created_at: version_number is monotonic and unique per
         # dag_id, so it is deterministic even when two versions share a created_at timestamp.
         # write_dag relies on this select to compute the next version_number; ordering by
@@ -188,6 +192,7 @@ def get_latest_version(
         bundle_version: str | None = None,
         load_dag_model: bool = False,
         load_bundle_model: bool = False,
+        load_serialized_dag: bool = False,
         session: Session = NEW_SESSION,
     ) -> DagVersion | None:
         """
@@ -197,6 +202,7 @@ def get_latest_version(
         :param session: The database session.
         :param load_dag_model: Whether to load the DAG model.
         :param load_bundle_model: Whether to load the DagBundle model.
+        :param load_serialized_dag: Whether to eagerly load the serialized DAG.
         :return: The latest version of the DAG or None if not found.
         """
         return session.scalar(
@@ -205,6 +211,7 @@ def get_latest_version(
                 bundle_version=bundle_version,
                 load_dag_model=load_dag_model,
                 load_bundle_model=load_bundle_model,
+                load_serialized_dag=load_serialized_dag,
             )
         )
 

airflow-core/src/airflow/serialization/definitions/dag.py

@@ -36,6 +36,7 @@
 from airflow.exceptions import (
     AirflowException,
     DagNotPartitionedError,
+    DagVersionNotFound,
     InvalidPartitionKeyError,
     NodeNotFound,
     TaskNotFound,
@@ -581,8 +582,10 @@ def create_dagrun(
         creating_job_id: int | None = None,
         backfill_id: NonNegativeInt | None = None,
         partition_key: str | None = None,
+        bundle_version: str | None = None,
         partition_date: datetime.datetime | None = None,
         note: str | None = None,
+        dag_version: DagVersion | None = None,
         session: Session = NEW_SESSION,
     ) -> DagRun:
         """
@@ -650,10 +653,28 @@ def create_dagrun(
                     f"is reserved for {inferred_run_type.value} runs"
                 )
 
-        self.validate_partition_key(partition_key)
-
         # todo: AIP-78 add verification that if run type is backfill then we have a backfill id
-        copied_params = self.params.deep_merge(conf)
+
+        # When triggering against a specific bundle version, resolve that version first so
+        # partition_key and conf are validated against it (not the live/latest dag).
+        if bundle_version is not None:
+            if self.disable_bundle_versioning:
+                raise ValueError(f"DAG with dag_id: '{self.dag_id}' does not support bundle versioning")
+            if dag_version is None:
+                dag_version = DagVersion.get_latest_version(
+                    self.dag_id, bundle_version=bundle_version, load_serialized_dag=True, session=session
+                )
+                if not dag_version:
+                    raise DagVersionNotFound(
+                        f"DAG with dag_id: '{self.dag_id}' does not have a version for bundle_version '{bundle_version}'"
+                    )
+            params_dag = dag_version.serialized_dag.dag
+        else:
+            params_dag = self
+
+        params_dag.validate_partition_key(partition_key)
+
+        copied_params = params_dag.params.deep_merge(conf)
         copied_params.validate()
         orm_dagrun = _create_orm_dagrun(
             dag=self,
@@ -670,12 +691,15 @@ def create_dagrun(
             triggered_by=triggered_by,
             triggering_user_name=triggering_user_name,
             partition_key=partition_key,
+            bundle_version=bundle_version,
             partition_date=partition_date,
             note=note,
+            dag_version=dag_version,
+            resolved_dag=params_dag if bundle_version is not None else None,
             session=session,
         )
 
-        if self.deadline:
+        if params_dag.deadline:
             self._process_dagrun_deadline_alerts(orm_dagrun, session)
 
         return orm_dagrun
@@ -1391,16 +1415,25 @@ def _create_orm_dagrun(
     triggered_by: DagRunTriggeredByType,
     triggering_user_name: str | None = None,
     partition_key: str | None = None,
+    bundle_version: str | None = None,
     partition_date: datetime.datetime | None = None,
     note: str | None = None,
+    dag_version: DagVersion | None = None,
+    resolved_dag: SerializedDAG | None = None,
     session: Session = NEW_SESSION,
 ) -> DagRun:
-    bundle_version = None
-    if not dag.disable_bundle_versioning:
-        bundle_version = session.scalar(
-            select(DagModel.bundle_version).where(DagModel.dag_id == dag.dag_id),
-        )
-    dag_version = DagVersion.get_latest_version(dag.dag_id, session=session)
+    resolved_bundle_version: str | None = None
+    use_resolved_dag = False
+    if dag_version is not None:
+        resolved_bundle_version = bundle_version
+        use_resolved_dag = True
+    else:
+        if not dag.disable_bundle_versioning:
+            resolved_bundle_version = session.scalar(
+                select(DagModel.bundle_version).where(DagModel.dag_id == dag.dag_id)
+            )
+        dag_version = DagVersion.get_latest_version(dag.dag_id, session=session)
+
     if not dag_version:
         raise AirflowException(f"Cannot create DagRun for DAG {dag.dag_id} because the dag is not serialized")
 
@@ -1418,7 +1451,7 @@ def _create_orm_dagrun(
         triggered_by=triggered_by,
         triggering_user_name=triggering_user_name,
         backfill_id=backfill_id,
-        bundle_version=bundle_version,
+        bundle_version=resolved_bundle_version,
         partition_key=partition_key,
         partition_date=partition_date,
         note=note,
@@ -1431,6 +1464,8 @@ def _create_orm_dagrun(
     session.add(run)
     session.flush()
     run.dag = dag
+    if use_resolved_dag:
+        run.dag = resolved_dag if resolved_dag is not None else dag_version.serialized_dag.dag
     # create the associated task instances
     # state is None at the moment of creation
     run.verify_integrity(session=session, dag_version_id=dag_version.id)