Fix alembic autogeneration and rename mismatching constraints

airflow/cli/cli_config.py

@@ -691,6 +691,12 @@ def string_lower_type(val):
     action="store_true",
     default=False,
 )
+ARG_DB_USE_MIGRATION_FILES = Arg(
+    ("-m", "--use-migration-files"),
+    help="Use migration files to perform migration",
+    action="store_true",
+    default=False,
+)
 
 # webserver
 ARG_PORT = Arg(
@@ -1525,7 +1531,7 @@ class GroupCommand(NamedTuple):
         name="reset",
         help="Burn down and rebuild the metadata database",
         func=lazy_load_command("airflow.cli.commands.db_command.resetdb"),
-        args=(ARG_YES, ARG_DB_SKIP_INIT, ARG_VERBOSE),
+        args=(ARG_YES, ARG_DB_SKIP_INIT, ARG_DB_USE_MIGRATION_FILES, ARG_VERBOSE),
     ),
     ActionCommand(
         name="upgrade",
@@ -1545,6 +1551,7 @@ class GroupCommand(NamedTuple):
             ARG_DB_FROM_REVISION,
             ARG_DB_FROM_VERSION,
             ARG_DB_RESERIALIZE_DAGS,
+            ARG_DB_USE_MIGRATION_FILES,
             ARG_VERBOSE,
         ),
         hide=True,
@@ -1568,6 +1575,7 @@ class GroupCommand(NamedTuple):
             ARG_DB_FROM_REVISION,
             ARG_DB_FROM_VERSION,
             ARG_DB_RESERIALIZE_DAGS,
+            ARG_DB_USE_MIGRATION_FILES,
             ARG_VERBOSE,
         ),
     ),

airflow/cli/commands/db_command.py

@@ -63,7 +63,7 @@ def resetdb(args):
     print(f"DB: {settings.engine.url!r}")
     if not (args.yes or input("This will drop existing tables if they exist. Proceed? (y/n)").upper() == "Y"):
         raise SystemExit("Cancelled")
-    db.resetdb(skip_init=args.skip_init)
+    db.resetdb(skip_init=args.skip_init, use_migration_files=args.use_migration_files)
 
 
 def upgradedb(args):
@@ -142,6 +142,7 @@ def migratedb(args):
         from_revision=from_revision,
         show_sql_only=args.show_sql_only,
         reserialize_dags=args.reserialize_dags,
+        use_migration_files=args.use_migration_files,
     )
     if not args.show_sql_only:
         print("Database migrating done!")

airflow/migrations/env.py

@@ -29,6 +29,9 @@
 
 def include_object(_, name, type_, *args):
     """Filter objects for autogenerating revisions."""
+    # Ignore the sqlite_sequence table, which is an internal SQLite construct
+    if name == "sqlite_sequence":
+        return False
     # Ignore _anything_ to do with Celery, or FlaskSession's tables
     if type_ == "table" and (name.startswith("celery_") or name == "session"):
         return False

airflow/migrations/versions/0142_2_9_2_fix_inconsistency_between_ORM_and_migration_files.py

@@ -0,0 +1,297 @@
+#
+# 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.
+
+"""Fix inconsistency between ORM and migration files.
+
+Revision ID: 686269002441
+Revises: bff083ad727d
+Create Date: 2024-04-15 14:19:49.913797
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import literal
+
+# revision identifiers, used by Alembic.
+revision = "686269002441"
+down_revision = "bff083ad727d"
+branch_labels = None
+depends_on = None
+airflow_version = "2.9.2"
+
+
+def upgrade():
+    """Apply Update missing constraints."""
+    conn = op.get_bind()
+    if conn.dialect.name == "mysql":
+        # TODO: Rewrite these queries to use alembic when lowest MYSQL version supports IF EXISTS
+        conn.execute(
+            sa.text("""
+        set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE
+            CONSTRAINT_SCHEMA = DATABASE() AND
+            TABLE_NAME        = 'connection' AND
+            CONSTRAINT_NAME   = 'unique_conn_id' AND
+            CONSTRAINT_TYPE   = 'UNIQUE') = true,'ALTER TABLE connection
+            drop constraint unique_conn_id','select 1');
+
+        prepare stmt from @var;
+        execute stmt;
+        deallocate prepare stmt;
+        """)
+        )
+        # Dropping the below and recreating cause there's no IF NOT EXISTS in mysql
+        conn.execute(
+            sa.text("""
+                set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE
+                    CONSTRAINT_SCHEMA = DATABASE() AND
+                    TABLE_NAME        = 'connection' AND
+                    CONSTRAINT_NAME   = 'connection_conn_id_uq' AND
+                    CONSTRAINT_TYPE   = 'UNIQUE') = true,'ALTER TABLE connection
+                    drop constraint connection_conn_id_uq','select 1');
+
+                prepare stmt from @var;
+                execute stmt;
+                deallocate prepare stmt;
+                """)
+        )
+    elif conn.dialect.name == "sqlite":
+        # SQLite does not support DROP CONSTRAINT
+        # We have to recreate the table without the constraint
+        conn.execute(sa.text("PRAGMA foreign_keys=off"))
+        conn.execute(
+            sa.text("""
+        CREATE TABLE connection_new (
+                id INTEGER NOT NULL,
+                conn_id VARCHAR(250) NOT NULL,
+                conn_type VARCHAR(500) NOT NULL,
+                host VARCHAR(500),
+                schema VARCHAR(500),
+                login TEXT,
+                password TEXT,
+                port INTEGER,
+                extra TEXT,
+                is_encrypted BOOLEAN,
+                is_extra_encrypted BOOLEAN,
+                description VARCHAR(5000),
+                CONSTRAINT connection_pkey PRIMARY KEY (id),
+                CONSTRAINT connection_conn_id_uq UNIQUE (conn_id)
+        )
+        """)
+        )
+        conn.execute(sa.text("INSERT INTO connection_new SELECT * FROM connection"))
+        conn.execute(sa.text("DROP TABLE connection"))
+        conn.execute(sa.text("ALTER TABLE connection_new RENAME TO connection"))
+        conn.execute(sa.text("PRAGMA foreign_keys=on"))
+    else:
+        op.execute("ALTER TABLE connection DROP CONSTRAINT IF EXISTS unique_conn_id")
+        # Dropping and recreating cause there's no IF NOT EXISTS
+        op.execute("ALTER TABLE connection DROP CONSTRAINT IF EXISTS connection_conn_id_uq")
+
+    with op.batch_alter_table("connection") as batch_op:
+        batch_op.create_unique_constraint(batch_op.f("connection_conn_id_uq"), ["conn_id"])
+
+    max_cons = sa.table("dag", sa.column("max_consecutive_failed_dag_runs"))
+    op.execute(max_cons.update().values(max_consecutive_failed_dag_runs=literal("0")))
+    with op.batch_alter_table("dag") as batch_op:
+        batch_op.alter_column("max_consecutive_failed_dag_runs", existing_type=sa.Integer(), nullable=False)
+
+    with op.batch_alter_table("task_instance") as batch_op:
+        batch_op.drop_constraint("task_instance_dag_run_fkey", type_="foreignkey")
+
+    with op.batch_alter_table("task_reschedule") as batch_op:
+        batch_op.drop_constraint("task_reschedule_dr_fkey", type_="foreignkey")
+
+    if conn.dialect.name == "mysql":
+        conn.execute(
+            sa.text("""
+                        set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE
+                            CONSTRAINT_SCHEMA = DATABASE() AND
+                            TABLE_NAME        = 'dag_run' AND
+                            CONSTRAINT_NAME   = 'dag_run_dag_id_execution_date_uq' AND
+                            CONSTRAINT_TYPE   = 'UNIQUE') = true,'ALTER TABLE dag_run
+                            drop constraint dag_run_dag_id_execution_date_uq','select 1');
+
+                        prepare stmt from @var;
+                        execute stmt;
+                        deallocate prepare stmt;
+                        """)
+        )
+        conn.execute(
+            sa.text("""
+                        set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE
+                            CONSTRAINT_SCHEMA = DATABASE() AND
+                            TABLE_NAME        = 'dag_run' AND
+                            CONSTRAINT_NAME   = 'dag_run_dag_id_run_id_uq' AND
+                            CONSTRAINT_TYPE   = 'UNIQUE') = true,'ALTER TABLE dag_run
+                            drop constraint dag_run_dag_id_run_id_uq','select 1');
+
+                        prepare stmt from @var;
+                        execute stmt;
+                        deallocate prepare stmt;
+                        """)
+        )
+        # below we drop and recreate the constraints because there's no IF NOT EXISTS
+        conn.execute(
+            sa.text("""
+                                set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE
+                                    CONSTRAINT_SCHEMA = DATABASE() AND
+                                    TABLE_NAME        = 'dag_run' AND
+                                    CONSTRAINT_NAME   = 'dag_run_dag_id_execution_date_key' AND
+                                    CONSTRAINT_TYPE   = 'UNIQUE') = true,'ALTER TABLE dag_run
+                                    drop constraint dag_run_dag_id_execution_date_key','select 1');
+
+                                prepare stmt from @var;
+                                execute stmt;
+                                deallocate prepare stmt;
+                                """)
+        )
+        conn.execute(
+            sa.text("""
+                            set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE
+                                CONSTRAINT_SCHEMA = DATABASE() AND
+                                TABLE_NAME        = 'dag_run' AND
+                                CONSTRAINT_NAME   = 'dag_run_dag_id_run_id_key' AND
+                                CONSTRAINT_TYPE   = 'UNIQUE') = true,'ALTER TABLE dag_run
+                                drop constraint dag_run_dag_id_run_id_key','select 1');
+
+                            prepare stmt from @var;
+                            execute stmt;
+                            deallocate prepare stmt;
+                            """)
+        )
+        with op.batch_alter_table("callback_request", schema=None) as batch_op:
+            batch_op.alter_column(
+                "processor_subdir",
+                existing_type=sa.Text(length=2000),
+                type_=sa.String(length=2000),
+                existing_nullable=True,
+            )
+
+        with op.batch_alter_table("dag", schema=None) as batch_op:
+            batch_op.alter_column(
+                "processor_subdir",
+                existing_type=sa.Text(length=2000),
+                type_=sa.String(length=2000),
+                existing_nullable=True,
+            )
+
+        with op.batch_alter_table("import_error", schema=None) as batch_op:
+            batch_op.alter_column(
+                "processor_subdir",
+                existing_type=sa.Text(length=2000),
+                type_=sa.String(length=2000),
+                existing_nullable=True,
+            )
+
+        with op.batch_alter_table("serialized_dag", schema=None) as batch_op:
+            batch_op.alter_column(
+                "processor_subdir",
+                existing_type=sa.Text(length=2000),
+                type_=sa.String(length=2000),
+                existing_nullable=True,
+            )
+
+    elif conn.dialect.name == "sqlite":
+        # SQLite does not support DROP CONSTRAINT
+        # We have to recreate the table without the constraint
+        conn.execute(sa.text("PRAGMA foreign_keys=off"))
+        conn.execute(
+            sa.text("""
+            CREATE TABLE dag_run_new (
+                id INTEGER NOT NULL,
+                dag_id VARCHAR(250) NOT NULL,
+                queued_at TIMESTAMP,
+                execution_date TIMESTAMP NOT NULL,
+                start_date TIMESTAMP,
+                end_date TIMESTAMP,
+                state VARCHAR(50),
+                run_id VARCHAR(250) NOT NULL,
+                creating_job_id INTEGER,
+                external_trigger BOOLEAN,
+                run_type VARCHAR(50) NOT NULL,
+                conf BLOB,
+                data_interval_start TIMESTAMP,
+                data_interval_end TIMESTAMP,
+                last_scheduling_decision TIMESTAMP,
+                dag_hash VARCHAR(32),
+                log_template_id INTEGER,
+                updated_at TIMESTAMP,
+                clear_number INTEGER DEFAULT '0' NOT NULL,
+                CONSTRAINT dag_run_pkey PRIMARY KEY (id),
+                CONSTRAINT dag_run_dag_id_execution_date_key UNIQUE (dag_id, execution_date),
+                CONSTRAINT dag_run_dag_id_run_id_key UNIQUE (dag_id, run_id),
+                CONSTRAINT task_instance_log_template_id_fkey FOREIGN KEY(log_template_id) REFERENCES log_template (id) ON DELETE NO ACTION
+            )
+        """)
+        )
+
+        conn.execute(sa.text("INSERT INTO dag_run_new SELECT * FROM dag_run"))
+        conn.execute(sa.text("DROP TABLE dag_run"))
+        conn.execute(sa.text("ALTER TABLE dag_run_new RENAME TO dag_run"))
+        conn.execute(sa.text("PRAGMA foreign_keys=on"))
+        with op.batch_alter_table("dag_run") as batch_op:
+            batch_op.create_index("dag_id_state", ["dag_id", "state"], if_not_exists=True)
+            batch_op.create_index("idx_dag_run_dag_id", ["dag_id"], if_not_exists=True)
+            batch_op.create_index(
+                "idx_dag_run_running_dags",
+                ["state", "dag_id"],
+                sqlite_where=sa.text("state='running'"),
+                if_not_exists=True,
+            )
+            batch_op.create_index(
+                "idx_dag_run_queued_dags",
+                ["state", "dag_id"],
+                sqlite_where=sa.text("state='queued'"),
+                if_not_exists=True,
+            )
+
+    else:
+        op.execute("ALTER TABLE dag_run DROP CONSTRAINT IF EXISTS dag_run_dag_id_execution_date_uq")
+        op.execute("ALTER TABLE dag_run DROP CONSTRAINT IF EXISTS dag_run_dag_id_run_id_uq")
+        # below we drop and recreate the constraints because there's no IF NOT EXISTS
+        op.execute("ALTER TABLE dag_run DROP CONSTRAINT IF EXISTS dag_run_dag_id_execution_date_key")
+        op.execute("ALTER TABLE dag_run DROP CONSTRAINT IF EXISTS dag_run_dag_id_run_id_key")
+
+    with op.batch_alter_table("dag_run") as batch_op:
+        batch_op.create_unique_constraint("dag_run_dag_id_execution_date_key", ["dag_id", "execution_date"])
+        batch_op.create_unique_constraint("dag_run_dag_id_run_id_key", ["dag_id", "run_id"])
+
+    with op.batch_alter_table("task_instance") as batch_op:
+        batch_op.create_foreign_key(
+            "task_instance_dag_run_fkey",
+            "dag_run",
+            ["dag_id", "run_id"],
+            ["dag_id", "run_id"],
+            ondelete="CASCADE",
+        )
+
+    with op.batch_alter_table("task_reschedule") as batch_op:
+        batch_op.create_foreign_key(
+            "task_reschedule_dr_fkey",
+            "dag_run",
+            ["dag_id", "run_id"],
+            ["dag_id", "run_id"],
+            ondelete="CASCADE",
+        )
+
+
+def downgrade():
+    """NO downgrade because this is to make ORM consistent with the database."""

airflow/migrations/versions/0142_2_10_0_add_new_executor_field_to_db.py → airflow/migrations/versions/0143_2_10_0_add_new_executor_field_to_db.py

@@ -19,7 +19,7 @@
 """add new executor field to db.
 
 Revision ID: 677fdbb7fc54
-Revises: 1949afb29106
+Revises: 686269002441
 Create Date: 2024-04-01 15:26:59.186579
 
 """
@@ -31,7 +31,7 @@
 
 # revision identifiers, used by Alembic.
 revision = "677fdbb7fc54"
-down_revision = "bff083ad727d"
+down_revision = "686269002441"
 branch_labels = None
 depends_on = None
 airflow_version = "2.10.0"

airflow/models/dagrun.py

@@ -147,7 +147,7 @@ class DagRun(Base, LoggingMixin):
     # Keeps track of the number of times the dagrun had been cleared.
     # This number is incremented only when the DagRun is re-Queued,
     # when the DagRun is cleared.
-    clear_number = Column(Integer, default=0, nullable=False)
+    clear_number = Column(Integer, default=0, nullable=False, server_default="0")
 
     # Remove this `if` after upgrading Sphinx-AutoAPI
     if not TYPE_CHECKING and "BUILDING_AIRFLOW_DOCS" in os.environ: