Fix ORM vs migration files inconsistencies

airflow/migrations/versions/0010_2_9_0_add_run_id_to_audit_log_table_and_change_event_name_length.py

@@ -59,8 +59,8 @@ def downgrade():
     if conn.dialect.name == "mssql":
         with op.batch_alter_table("log") as batch_op:
             batch_op.drop_index("idx_log_event")
-            batch_op.alter_column("event", type_=sa.String(30), nullable=False)
+            batch_op.alter_column("event", type_=sa.String(30))
             batch_op.create_index("idx_log_event", ["event"])
     else:
         with op.batch_alter_table("log") as batch_op:
-            batch_op.alter_column("event", type_=sa.String(30), nullable=False)
+            batch_op.alter_column("event", type_=sa.String(30))

airflow/migrations/versions/0028_3_0_0_drop_ab_user_id_foreign_key.py

@@ -45,6 +45,13 @@ def upgrade():
     with op.batch_alter_table("task_instance_note", schema=None) as batch_op:
         batch_op.drop_constraint("task_instance_note_user_fkey", type_="foreignkey")
 
+    if op.get_bind().dialect.name == "mysql":
+        with op.batch_alter_table("dag_run_note", schema=None) as batch_op:
+            batch_op.drop_index("dag_run_note_user_fkey")
+
+        with op.batch_alter_table("task_instance_note", schema=None) as batch_op:
+            batch_op.drop_index("task_instance_note_user_fkey")
+
 
 def downgrade():
     """Unapply Drop ab_user.id foreign key."""
@@ -53,3 +60,10 @@ def downgrade():
 
     with op.batch_alter_table("dag_run_note", schema=None) as batch_op:
         batch_op.create_foreign_key("dag_run_note_user_fkey", "ab_user", ["user_id"], ["id"])
+
+    if op.get_bind().dialect.name == "mysql":
+        with op.batch_alter_table("task_instance_note", schema=None) as batch_op:
+            batch_op.create_index("task_instance_note_user_fkey", ["user_id"], unique=False)
+
+        with op.batch_alter_table("dag_run_note", schema=None) as batch_op:
+            batch_op.create_index("dag_run_note_user_fkey", ["user_id"], unique=False)

airflow/migrations/versions/0032_3_0_0_drop_execution_date_unique.py

@@ -31,9 +31,10 @@
 
 from __future__ import annotations
 
-import sqlalchemy as sa
 from alembic import op
 
+from airflow.migrations.db_types import TIMESTAMP
+
 # revision identifiers, used by Alembic.
 revision = "1cdc775ca98f"
 down_revision = "a2c32e6c7729"
@@ -44,14 +45,24 @@
 
 def upgrade():
     with op.batch_alter_table("dag_run", schema=None) as batch_op:
-        batch_op.alter_column("execution_date", new_column_name="logical_date", existing_type=sa.TIMESTAMP)
+        batch_op.alter_column(
+            "execution_date",
+            new_column_name="logical_date",
+            existing_type=TIMESTAMP(timezone=True),
+            existing_nullable=False,
+        )
     with op.batch_alter_table("dag_run", schema=None) as batch_op:
         batch_op.drop_constraint("dag_run_dag_id_execution_date_key", type_="unique")
 
 
 def downgrade():
     with op.batch_alter_table("dag_run", schema=None) as batch_op:
-        batch_op.alter_column("logical_date", new_column_name="execution_date", existing_type=sa.TIMESTAMP)
+        batch_op.alter_column(
+            "logical_date",
+            new_column_name="execution_date",
+            existing_type=TIMESTAMP(timezone=True),
+            existing_nullable=False,
+        )
     with op.batch_alter_table("dag_run", schema=None) as batch_op:
         batch_op.create_unique_constraint(
             "dag_run_dag_id_execution_date_key",

airflow/migrations/versions/0033_3_0_0_add_tables_for_backfill.py

@@ -46,10 +46,10 @@ def upgrade():
     op.create_table(
         "backfill",
         sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
-        sa.Column("dag_id", sa.String(length=250), nullable=True),
+        sa.Column("dag_id", sa.String(length=250), nullable=False),
         sa.Column("from_date", airflow.utils.sqlalchemy.UtcDateTime(timezone=True), nullable=False),
         sa.Column("to_date", airflow.utils.sqlalchemy.UtcDateTime(timezone=True), nullable=False),
-        sa.Column("dag_run_conf", sqlalchemy_jsonfield.jsonfield.JSONField(), nullable=True),
+        sa.Column("dag_run_conf", sqlalchemy_jsonfield.jsonfield.JSONField(), nullable=False),
         sa.Column("is_paused", sa.Boolean(), nullable=True),
         sa.Column("max_active_runs", sa.Integer(), nullable=False),
         sa.Column("created_at", airflow.utils.sqlalchemy.UtcDateTime(timezone=True), nullable=False),
@@ -65,6 +65,10 @@ def upgrade():
         sa.Column("sort_ordinal", sa.Integer(), nullable=False),
         sa.PrimaryKeyConstraint("id", name=op.f("backfill_dag_run_pkey")),
         sa.UniqueConstraint("backfill_id", "dag_run_id", name="ix_bdr_backfill_id_dag_run_id"),
+        sa.ForeignKeyConstraint(
+            ["backfill_id"], ["backfill.id"], name="bdr_backfill_fkey", ondelete="cascade"
+        ),
+        sa.ForeignKeyConstraint(["dag_run_id"], ["dag_run.id"], name="bdr_dag_run_fkey", ondelete="set null"),
     )
 
 

airflow/migrations/versions/0036_3_0_0_add_name_field_to_dataset_model.py

@@ -61,9 +61,7 @@ def upgrade():
     # Add 'name' column. Set it to nullable for now.
     with op.batch_alter_table("dataset", schema=None) as batch_op:
         batch_op.add_column(sa.Column("name", _STRING_COLUMN_TYPE))
-        batch_op.add_column(
-            sa.Column("group", _STRING_COLUMN_TYPE, default=str, server_default="", nullable=False)
-        )
+        batch_op.add_column(sa.Column("group", _STRING_COLUMN_TYPE, default="", nullable=False))
     # Fill name from uri column.
     with Session(bind=op.get_bind()) as session:
         session.execute(sa.text("update dataset set name=uri"))

airflow/migrations/versions/0040_3_0_0_add_exception_reason_and_logical_date_.py

@@ -42,7 +42,7 @@
 def upgrade():
     """Apply Add exception_reason and logical_date to BackfillDagRun."""
     with op.batch_alter_table("backfill", schema=None) as batch_op:
-        batch_op.add_column(sa.Column("reprocess_behavior", sa.String(length=250), nullable=True))
+        batch_op.add_column(sa.Column("reprocess_behavior", sa.String(length=250), nullable=False))
     with op.batch_alter_table("backfill_dag_run", schema=None) as batch_op:
         batch_op.add_column(sa.Column("exception_reason", sa.String(length=250), nullable=True))
         batch_op.add_column(sa.Column("logical_date", UtcDateTime(timezone=True), nullable=False))

airflow/migrations/versions/0041_3_0_0_rename_dataset_as_asset.py

@@ -285,13 +285,6 @@ def upgrade():
             unique=False,
         )
 
-        batch_op.create_foreign_key(
-            constraint_name="toar_asset_fkey",
-            referent_table="asset",
-            local_cols=["asset_id"],
-            remote_cols=["id"],
-            ondelete="CASCADE",
-        )
         batch_op.create_foreign_key(
             constraint_name="toar_dag_id_fkey",
             referent_table="dag",
@@ -321,13 +314,6 @@ def upgrade():
             unique=False,
         )
 
-        batch_op.create_foreign_key(
-            constraint_name="adrq_asset_fkey",
-            referent_table="asset",
-            local_cols=["asset_id"],
-            remote_cols=["id"],
-            ondelete="CASCADE",
-        )
         batch_op.create_foreign_key(
             constraint_name="adrq_dag_fkey",
             referent_table="dag",
@@ -565,7 +551,6 @@ def downgrade():
     with op.batch_alter_table("task_outlet_dataset_reference", schema=None) as batch_op:
         batch_op.alter_column("asset_id", new_column_name="dataset_id", type_=sa.Integer(), nullable=False)
 
-        batch_op.drop_constraint("toar_asset_fkey", type_="foreignkey")
         batch_op.drop_constraint("toar_dag_id_fkey", type_="foreignkey")
 
         _rename_index(
@@ -600,7 +585,6 @@ def downgrade():
     with op.batch_alter_table("dataset_dag_run_queue", schema=None) as batch_op:
         batch_op.alter_column("asset_id", new_column_name="dataset_id", type_=sa.Integer(), nullable=False)
 
-        batch_op.drop_constraint("adrq_asset_fkey", type_="foreignkey")
         batch_op.drop_constraint("adrq_dag_fkey", type_="foreignkey")
 
         _rename_pk_constraint(

airflow/migrations/versions/0042_3_0_0_create_fks_as_datasets_are_renamed.py

@@ -0,0 +1,55 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+create foreign key constraints for assets.
+
+Revision ID: c4a1639f0f67
+Revises: 05234396c6fc
+Create Date: 2024-11-22 09:49:41.813016
+
+"""
+
+from __future__ import annotations
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "c4a1639f0f67"
+down_revision = "05234396c6fc"
+branch_labels = None
+depends_on = None
+airflow_version = "3.0.0"
+
+
+def upgrade():
+    """Apply create foreign key constraints for assets."""
+    with op.batch_alter_table("asset_dag_run_queue", schema=None) as batch_op:
+        batch_op.create_foreign_key("adrq_asset_fkey", "asset", ["asset_id"], ["id"], ondelete="CASCADE")
+
+    with op.batch_alter_table("task_outlet_asset_reference", schema=None) as batch_op:
+        batch_op.create_foreign_key("toar_asset_fkey", "asset", ["asset_id"], ["id"], ondelete="CASCADE")
+
+
+def downgrade():
+    """Unapply create foreign key constraints for assets."""
+    with op.batch_alter_table("task_outlet_asset_reference", schema=None) as batch_op:
+        batch_op.drop_constraint("toar_asset_fkey", type_="foreignkey")
+
+    with op.batch_alter_table("asset_dag_run_queue", schema=None) as batch_op:
+        batch_op.drop_constraint("adrq_asset_fkey", type_="foreignkey")

airflow/migrations/versions/0042_3_0_0_add_uuid_primary_key_to_task_instance_.py → airflow/migrations/versions/0043_3_0_0_add_uuid_primary_key_to_task_instance_.py

@@ -20,7 +20,7 @@
 Add UUID primary key to ``task_instance`` table.
 
 Revision ID: d59cbbef95eb
-Revises: 05234396c6fc
+Revises: c4a1639f0f67
 Create Date: 2024-10-21 22:39:12.394079
 """
 
@@ -33,7 +33,7 @@
 
 # revision identifiers, used by Alembic.
 revision = "d59cbbef95eb"
-down_revision = "05234396c6fc"
+down_revision = "c4a1639f0f67"
 branch_labels = "None"
 depends_on = None
 airflow_version = "3.0.0"
@@ -167,6 +167,32 @@ def _get_type_id_column(dialect_name: str) -> sa.types.TypeEngine:
         return sa.String(36)
 
 
+def create_foreign_keys():
+    for fk in ti_fk_constraints:
+        if fk["table"] in ["task_instance_history", "task_map"]:
+            continue
+        with op.batch_alter_table(fk["table"]) as batch_op:
+            batch_op.create_foreign_key(
+                constraint_name=fk["fk"],
+                referent_table=ti_table,
+                local_cols=ti_fk_cols,
+                remote_cols=ti_fk_cols,
+                ondelete="CASCADE",
+            )
+    for fk in ti_fk_constraints:
+        if fk["table"] not in ["task_instance_history", "task_map"]:
+            continue
+        with op.batch_alter_table(fk["table"]) as batch_op:
+            batch_op.create_foreign_key(
+                constraint_name=fk["fk"],
+                referent_table=ti_table,
+                local_cols=ti_fk_cols,
+                remote_cols=ti_fk_cols,
+                ondelete="CASCADE",
+                onupdate="CASCADE",
+            )
+
+
 def upgrade():
     """Add UUID primary key to task instance table."""
     conn = op.get_bind()
@@ -232,15 +258,7 @@ def upgrade():
         batch_op.create_primary_key("task_instance_pkey", ["id"])
 
     # Create foreign key constraints
-    for fk in ti_fk_constraints:
-        with op.batch_alter_table(fk["table"]) as batch_op:
-            batch_op.create_foreign_key(
-                constraint_name=fk["fk"],
-                referent_table=ti_table,
-                local_cols=ti_fk_cols,
-                remote_cols=ti_fk_cols,
-                ondelete="CASCADE",
-            )
+    create_foreign_keys()
 
 
 def downgrade():
@@ -270,12 +288,4 @@ def downgrade():
         batch_op.create_primary_key("task_instance_pkey", ti_fk_cols)
 
     # Re-add foreign key constraints
-    for fk in ti_fk_constraints:
-        with op.batch_alter_table(fk["table"]) as batch_op:
-            batch_op.create_foreign_key(
-                constraint_name=fk["fk"],
-                referent_table=ti_table,
-                local_cols=ti_fk_cols,
-                remote_cols=ti_fk_cols,
-                ondelete="CASCADE",
-            )
+    create_foreign_keys()

airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py → airflow/migrations/versions/0050_3_0_0_remove_pickled_data_from_xcom_table.py

@@ -180,3 +180,5 @@ def downgrade():
 
         with op.batch_alter_table("xcom", schema=None) as batch_op:
             batch_op.drop_column("value_old")
+
+    op.drop_table("_xcom_archive", if_exists=True)

airflow/models/__init__.py

@@ -63,7 +63,6 @@ def import_all_models():
     import airflow.models.serialized_dag
     import airflow.models.taskinstancehistory
     import airflow.models.tasklog
-    import airflow.providers.fab.auth_manager.models
 
 
 def __getattr__(name):

airflow/models/asset.py

@@ -119,7 +119,7 @@ class AssetAliasModel(Base):
             ),
             "mysql",
         ),
-        default=str,
+        default="",
         nullable=False,
     )
 

docs/apache-airflow/img/airflow_erd.sha256

@@ -1 +1 @@
-7748eec981f977cc97b852d1fe982aebe24ec2d090ae8493a65cea101f9d42a5
+9f404e900b4d166d170d878fc00cad8e4397a4288167eebef1d77281a881e92d

docs/apache-airflow/migrations-ref.rst

@@ -53,7 +53,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``486ac7936b78``        | ``d59cbbef95eb`` | ``3.0.0``         | remove scheduler_lock column.                                |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
-| ``d59cbbef95eb``        | ``05234396c6fc`` | ``3.0.0``         | Add UUID primary key to ``task_instance`` table.             |
+| ``d59cbbef95eb``        | ``c4a1639f0f67`` | ``3.0.0``         | Add UUID primary key to ``task_instance`` table.             |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``c4a1639f0f67``        | ``05234396c6fc`` | ``3.0.0``         | create foreign key constraints for assets.                   |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``05234396c6fc``        | ``3a8972ecb8f9`` | ``3.0.0``         | Rename dataset as asset.                                     |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

scripts/in_container/run_generate_migration.sh

@@ -20,5 +20,7 @@
 
 cd "${AIRFLOW_SOURCES}" || exit 1
 cd "airflow" || exit 1
-airflow db reset
+airflow db reset -y
+airflow db downgrade -n 2.10.3 -y
+airflow db migrate -r heads
 alembic revision --autogenerate -m "${@}"

tests/utils/test_db.py

@@ -96,6 +96,8 @@ def test_database_schema_and_sqlalchemy_model_are_in_sync(self):
             lambda t: (t[0] == "remove_table" and t[1].name == "sqlite_sequence"),
             # fab version table
             lambda t: (t[0] == "remove_table" and t[1].name == "alembic_version_fab"),
+            # Ignore _xcom_archive table
+            lambda t: (t[0] == "remove_table" and t[1].name == "_xcom_archive"),
         ]
 
         for ignore in ignores:

tests_common/test_utils/db.py

@@ -62,7 +62,12 @@ def initial_db_init():
     from airflow.www.extensions.init_appbuilder import init_appbuilder
     from airflow.www.extensions.init_auth_manager import get_auth_manager
 
+    from tests_common.test_utils.compat import AIRFLOW_V_3_0_PLUS
+
     db.resetdb()
+    if AIRFLOW_V_3_0_PLUS:
+        db.downgrade(to_revision="5f2621c13b39")
+        db.upgradedb(to_revision="head")
     db.bootstrap_dagbag()
     # minimal app to add roles
     flask_app = Flask(__name__)