Enable disabling of dag bundle versioning (apache#47473)

This allows folks to opt in to running on the latest bundle version for each task - basically how Airflow 2 operates, even if they have a versioned bundle under the hood.

Approach:
* config options, something like use_bundle_versioning
* DAG kwarg that uses that config option as the default
* conditionally add bundle_version to the dagrun

Author: dstandish

airflow/config_templates/config.yml

@@ -2671,6 +2671,14 @@ dag_processor:
       type: integer
       example: ~
       default: "30"
+    disable_bundle_versioning:
+      description: |
+        Always run tasks with the latest code.  If set to True, the bundle version will not
+        be stored on the dag run and therefore, the latest code will always be used.
+      version_added: ~
+      type: boolean
+      example: ~
+      default: "False"
     bundle_refresh_check_interval:
       description: |
         How often the DAG processor should check if any DAG bundles are ready for a refresh, either by hitting

airflow/models/dag.py

@@ -261,6 +261,11 @@ def _create_orm_dagrun(
     triggered_by: DagRunTriggeredByType,
     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),
+        )
     run = DagRun(
         dag_id=dag.dag_id,
         run_id=run_id,
@@ -274,7 +279,7 @@ def _create_orm_dagrun(
         data_interval=data_interval,
         triggered_by=triggered_by,
         backfill_id=backfill_id,
-        bundle_version=session.scalar(select(DagModel.bundle_version).where(DagModel.dag_id == dag.dag_id)),
+        bundle_version=bundle_version,
     )
     # Load defaults into the following two fields to ensure result can be serialized detached
     run.log_template_id = int(session.scalar(select(func.max(LogTemplate.__table__.c.id))))

airflow/serialization/schema.json

@@ -206,7 +206,8 @@
           { "$ref": "#/definitions/task_group" }
         ]},
         "edge_info": { "$ref": "#/definitions/edge_info" },
-        "dag_dependencies": { "$ref": "#/definitions/dag_dependencies" }
+        "dag_dependencies": { "$ref": "#/definitions/dag_dependencies" },
+        "disable_bundle_versioning": {"type":  "boolean"}
       },
       "required": [
         "dag_id",

task-sdk/src/airflow/sdk/definitions/dag.py

@@ -439,6 +439,7 @@ def __rich_repr__(self):
 
     has_on_success_callback: bool = attrs.field(init=False)
     has_on_failure_callback: bool = attrs.field(init=False)
+    disable_bundle_versioning: bool = attrs.field(init=True)
 
     def __attrs_post_init__(self):
         from airflow.utils import timezone
@@ -510,6 +511,12 @@ def _default_timetable(instance: DAG):
         else:
             return _create_timetable(schedule, instance.timezone)
 
+    @disable_bundle_versioning.default
+    def _disable_bundle_versioning_default(self):
+        from airflow.configuration import conf as airflow_conf
+
+        return airflow_conf.getboolean("dag_processor", "disable_bundle_versioning")
+
     @timezone.default
     def _extract_tz(instance):
         import pendulum

task-sdk/tests/task_sdk/definitions/test_dag.py

@@ -282,7 +282,7 @@ def test_continuous_schedule_linmits_max_active_runs(self):
             dag = DAG("continuous", start_date=DEFAULT_DATE, schedule="@continuous", max_active_runs=25)
 
 
-# Test some of the arg valiadtion. This is not all the validations we perform, just some of them.
+# Test some of the arg validation. This is not all the validations we perform, just some of them.
 @pytest.mark.parametrize(
     ["attr", "value"],
     [

tests/models/test_dag.py

@@ -98,6 +98,7 @@
 from tests_common.test_utils.asserts import assert_queries_count
 from tests_common.test_utils.db import (
     clear_db_assets,
+    clear_db_dag_bundles,
     clear_db_dags,
     clear_db_runs,
     clear_db_serialized_dags,
@@ -120,9 +121,11 @@
 def clear_dags():
     clear_db_dags()
     clear_db_serialized_dags()
+    clear_db_dag_bundles()
     yield
     clear_db_dags()
     clear_db_serialized_dags()
+    clear_db_dag_bundles()
 
 
 @pytest.fixture
@@ -3493,3 +3496,38 @@ def test_validate_setup_teardown_trigger_rule(self):
                 Exception, match="Setup tasks must be followed with trigger rule ALL_SUCCESS."
             ):
                 dag.validate_setup_teardown()
+
+
+@pytest.mark.parametrize(
+    "disable, bundle_version, expected",
+    [
+        (True, "some-version", None),
+        (False, "some-version", "some-version"),
+    ],
+)
+def test_disable_bundle_versioning(disable, bundle_version, expected, dag_maker, session, clear_dags):
+    """When bundle versioning is disabled for a dag, the dag run should not have a bundle version."""
+
+    def hello():
+        print("hello")
+
+    with dag_maker(disable_bundle_versioning=disable, session=session, serialized=True) as dag:
+        PythonOperator(task_id="hi", python_callable=hello)
+
+    assert dag.disable_bundle_versioning is disable
+
+    # the dag *always* has bundle version
+    dag_model = session.scalar(select(DagModel).where(DagModel.dag_id == dag.dag_id))
+    dag_model.bundle_version = bundle_version
+    session.commit()
+
+    dr = dag.create_dagrun(
+        run_id="abcoercuhcrh",
+        run_after=pendulum.now(),
+        run_type="manual",
+        triggered_by=DagRunTriggeredByType.TEST,
+        state=None,
+    )
+
+    # but it only gets stamped on the dag run when bundle versioning not disabled
+    assert dr.bundle_version == expected

tests/serialization/test_dag_serialization.py

@@ -86,6 +86,7 @@
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.xcom import XCOM_RETURN_KEY
 
+from tests_common.test_utils.config import conf_vars
 from tests_common.test_utils.mock_operators import (
     AirflowLink2,
     CustomOperator,
@@ -154,6 +155,7 @@
         },
         "is_paused_upon_creation": False,
         "dag_id": "simple_dag",
+        "disable_bundle_versioning": False,
         "doc_md": "### DAG Tutorial Documentation",
         "fileloc": None,
         "_processor_dags_folder": f"{repo_root}/tests/dags",
@@ -2068,6 +2070,38 @@ def test_dag_on_failure_callback_roundtrip(self, passed_failure_callback, expect
 
         assert deserialized_dag.has_on_failure_callback is expected_value
 
+    @pytest.mark.parametrize(
+        "dag_arg, conf_arg, expected",
+        [
+            (True, "True", True),
+            (True, "False", True),
+            (False, "True", False),
+            (False, "False", False),
+            (None, "True", True),
+            (None, "False", False),
+        ],
+    )
+    def test_dag_disable_bundle_versioning_roundtrip(self, dag_arg, conf_arg, expected):
+        """
+        Test that when disable_bundle_versioning is passed to the DAG, has_disable_bundle_versioning is stored
+        in Serialized JSON blob. And when it is de-serialized dag.has_disable_bundle_versioning is set to True.
+
+        When the callback is not set, has_disable_bundle_versioning should not be stored in Serialized blob
+        and so default to False on de-serialization
+        """
+        with conf_vars({("dag_processor", "disable_bundle_versioning"): conf_arg}):
+            kwargs = {}
+            kwargs["disable_bundle_versioning"] = dag_arg
+            dag = DAG(
+                dag_id="test_dag_disable_bundle_versioning_roundtrip",
+                schedule=None,
+                **kwargs,
+            )
+            BaseOperator(task_id="simple_task", dag=dag, start_date=datetime(2019, 8, 1))
+            serialized_dag = SerializedDAG.to_dict(dag)
+            deserialized_dag = SerializedDAG.from_dict(serialized_dag)
+            assert deserialized_dag.disable_bundle_versioning is expected
+
     @pytest.mark.parametrize(
         "object_to_serialized, expected_output",
         [