Add DagBundle information to Serialized DagVersion table (#45976)

* Add DagBundle information to Serialized DagVersion table

This will help us keep track of the different versions that we see
in DagBundle since DagModel would always have the latest bundle version.

* fixup! Add DagBundle information to Serialized DagVersion table

* Fix test

Author: ephraimbuddy

airflow/dag_processing/collection.py

@@ -177,7 +177,9 @@ def _update_dag_owner_links(dag_owner_links: dict[str, str], dm: DagModel, *, se
     )
 
 
-def _serialize_dag_capturing_errors(dag: MaybeSerializedDAG, session: Session):
+def _serialize_dag_capturing_errors(
+    dag: MaybeSerializedDAG, bundle_name, session: Session, bundle_version: str | None
+):
     """
     Try to serialize the dag to the DB, but make a note of any errors.
 
@@ -192,6 +194,8 @@ def _serialize_dag_capturing_errors(dag: MaybeSerializedDAG, session: Session):
         # We can't use bulk_write_to_db as we want to capture each error individually
         dag_was_updated = SerializedDagModel.write_dag(
             dag,
+            bundle_name=bundle_name,
+            bundle_version=bundle_version,
             min_update_interval=settings.MIN_SERIALIZED_DAG_UPDATE_INTERVAL,
             session=session,
         )
@@ -333,7 +337,11 @@ def update_dag_parsing_results_in_db(
                 DAG.bulk_write_to_db(bundle_name, bundle_version, dags, session=session)
                 # Write Serialized DAGs to DB, capturing errors
                 for dag in dags:
-                    serialize_errors.extend(_serialize_dag_capturing_errors(dag, session))
+                    serialize_errors.extend(
+                        _serialize_dag_capturing_errors(
+                            dag=dag, bundle_name=bundle_name, bundle_version=bundle_version, session=session
+                        )
+                    )
             except OperationalError:
                 session.rollback()
                 raise

airflow/migrations/versions/0050_3_0_0_add_dagbundlemodel.py

@@ -29,6 +29,7 @@
 import sqlalchemy as sa
 from alembic import op
 
+from airflow.migrations.db_types import StringID
 from airflow.utils.sqlalchemy import UtcDateTime
 
 revision = "e229247a6cb1"
@@ -59,6 +60,10 @@ def upgrade():
     with op.batch_alter_table("import_error", schema=None) as batch_op:
         batch_op.add_column(sa.Column("bundle_name", sa.String(length=250), nullable=True))
 
+    with op.batch_alter_table("dag_version", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("bundle_name", StringID(), nullable=False))
+        batch_op.add_column(sa.Column("bundle_version", StringID()))
+
 
 def downgrade():
     with op.batch_alter_table("dag", schema=None) as batch_op:
@@ -71,4 +76,8 @@ def downgrade():
     with op.batch_alter_table("import_error", schema=None) as batch_op:
         batch_op.drop_column("bundle_name")
 
+    with op.batch_alter_table("dag_version", schema=None) as batch_op:
+        batch_op.drop_column("bundle_version")
+        batch_op.drop_column("bundle_name")
+
     op.drop_table("dag_bundle")

airflow/models/dag_version.py

@@ -45,6 +45,8 @@ class DagVersion(Base):
     version_number = Column(Integer, nullable=False, default=1)
     dag_id = Column(StringID(), ForeignKey("dag.dag_id", ondelete="CASCADE"), nullable=False)
     dag_model = relationship("DagModel", back_populates="dag_versions")
+    bundle_name = Column(StringID(), nullable=False)
+    bundle_version = Column(StringID())
     dag_code = relationship(
         "DagCode",
         back_populates="dag_version",
@@ -77,6 +79,8 @@ def write_dag(
         cls,
         *,
         dag_id: str,
+        bundle_name: str,
+        bundle_version: str | None = None,
         version_number: int = 1,
         session: Session = NEW_SESSION,
     ) -> DagVersion:
@@ -99,6 +103,8 @@ def write_dag(
         dag_version = DagVersion(
             dag_id=dag_id,
             version_number=version_number,
+            bundle_name=bundle_name,
+            bundle_version=bundle_version,
         )
         log.debug("Writing DagVersion %s to the DB", dag_version)
         session.add(dag_version)

airflow/models/serialized_dag.py

@@ -165,6 +165,8 @@ def _sort_serialized_dag_dict(cls, serialized_dag: Any):
     def write_dag(
         cls,
         dag: DAG | LazyDeserializedDAG,
+        bundle_name: str,
+        bundle_version: str | None = None,
         min_update_interval: int | None = None,
         session: Session = NEW_SESSION,
     ) -> bool:
@@ -201,8 +203,11 @@ def write_dag(
         if serialized_dag_hash is not None and serialized_dag_hash == new_serialized_dag.dag_hash:
             log.debug("Serialized DAG (%s) is unchanged. Skipping writing to DB", dag.dag_id)
             return False
+
         dagv = DagVersion.write_dag(
             dag_id=dag.dag_id,
+            bundle_name=bundle_name,
+            bundle_version=bundle_version,
             session=session,
         )
         log.debug("Writing Serialized DAG: %s to the DB", dag.dag_id)
@@ -340,6 +345,8 @@ def get(cls, dag_id: str, session: Session = NEW_SESSION) -> SerializedDagModel
     @provide_session
     def bulk_sync_to_db(
         dags: list[DAG] | list[LazyDeserializedDAG],
+        bundle_name: str,
+        bundle_version: str | None = None,
         session: Session = NEW_SESSION,
     ) -> None:
         """
@@ -354,6 +361,8 @@ def bulk_sync_to_db(
         for dag in dags:
             SerializedDagModel.write_dag(
                 dag=dag,
+                bundle_name=bundle_name,
+                bundle_version=bundle_version,
                 min_update_interval=MIN_SERIALIZED_DAG_UPDATE_INTERVAL,
                 session=session,
             )

docs/apache-airflow/img/airflow_erd.sha256

@@ -1 +1 @@
-cb858681fdc7a596db20c1c5dbf93812fd011a6df1e0b5322a49a51c8476bb93
+eb25e0718c9382cdbb02368c9c3e29c90da06ddaba8e8e92d9fc53417b714039