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

@@ -210,7 +210,7 @@ def upgrade():
             columns=["dag_id"],
             unique=False,
         )
-
+    with op.batch_alter_table("dag_schedule_asset_alias_reference", schema=None) as batch_op:
         batch_op.create_foreign_key(
             constraint_name="dsaar_asset_alias_fkey",
             referent_table="asset_alias",
@@ -284,14 +284,8 @@ def upgrade():
             columns=["dag_id"],
             unique=False,
         )
-
-        batch_op.create_foreign_key(
-            constraint_name="toar_asset_fkey",
-            referent_table="asset",
-            local_cols=["asset_id"],
-            remote_cols=["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")
         batch_op.create_foreign_key(
             constraint_name="toar_dag_id_fkey",
             referent_table="dag",
@@ -320,14 +314,8 @@ def upgrade():
             columns=["target_dag_id"],
             unique=False,
         )
-
-        batch_op.create_foreign_key(
-            constraint_name="adrq_asset_fkey",
-            referent_table="asset",
-            local_cols=["asset_id"],
-            remote_cols=["id"],
-            ondelete="CASCADE",
-        )
+    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")
         batch_op.create_foreign_key(
             constraint_name="adrq_dag_fkey",
             referent_table="dag",
@@ -564,7 +552,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")
 

airflow/migrations/versions/0042_3_0_0_add_uuid_primary_key_to_task_instance_.py

@@ -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

@@ -180,3 +180,5 @@ def downgrade():
 
         with op.batch_alter_table("xcom", schema=None) as batch_op:
             batch_op.drop_column("value_old")
+
+    op.execute(sa.text("DROP TABLE IF EXISTS _xcom_archive"))

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 @@
-aa9e2e5b2a52af1e92bc876727ad5e8958e291315096fc5249a9afa2c21a5d06
+b42b04b6cc47650cb9e7a37258a6e8e99bdca2677253715505b8ad287192bf72

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/utils/test_db_cleanup.py

@@ -277,11 +277,13 @@ def test__cleanup_table(self, table_name, date_add_kwargs, expected_to_delete, e
 
     @pytest.mark.parametrize(
         "skip_archive, expected_archives",
-        [pytest.param(True, 0, id="skip_archive"), pytest.param(False, 1, id="do_archive")],
+        [pytest.param(True, 1, id="skip_archive"), pytest.param(False, 2, id="do_archive")],
     )
     def test__skip_archive(self, skip_archive, expected_archives):
         """
         Verify that running cleanup_table with drops the archives when requested.
+
+        Archived tables from DB migration should be kept when skip_archive is True.
         """
         base_date = pendulum.DateTime(2022, 1, 1, tzinfo=pendulum.timezone("UTC"))
         num_tis = 10

tests_common/test_utils/db.py

@@ -73,7 +73,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")
     _bootstrap_dagbag()
     # minimal app to add roles
     flask_app = Flask(__name__)