Fix database reset failing on PostgreSQL with orphaned state_id references (#20941)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: alex.s <alex.s@prefect.io>
Co-authored-by: alex.s <ajstreed1@gmail.com>

Author: devin-ai-integration[bot]

src/prefect/server/database/_migrations/versions/postgresql/2025_06_13_160644_3b86c5ea017a_drop_task_run_state_id_foreign_key_.py

@@ -6,6 +6,7 @@
 
 """
 
+import sqlalchemy as sa
 from alembic import op
 
 # revision identifiers, used by Alembic.
@@ -23,6 +24,26 @@ def upgrade():
 
 
 def downgrade():
+    # Null out any task_run.state_id values that reference non-existent
+    # task_run_state rows. After the FK was dropped (upgrade), the system no
+    # longer enforces referential integrity, so orphaned references are
+    # expected. Without this cleanup, re-adding the FK would fail with an
+    # IntegrityError on PostgreSQL.
+    connection = op.get_bind()
+    connection.execute(
+        sa.text(
+            """
+            UPDATE task_run
+            SET state_id = NULL
+            WHERE state_id IS NOT NULL
+              AND NOT EXISTS (
+                  SELECT 1 FROM task_run_state
+                  WHERE task_run_state.id = task_run.state_id
+              )
+            """
+        )
+    )
+
     # Re-add the foreign key constraint
     op.create_foreign_key(
         "fk_task_run__state_id__task_run_state",

src/prefect/server/database/_migrations/versions/sqlite/2025_06_13_164506_8bb517bae6f9_drop_task_run_state_id_foreign_key_.py

@@ -6,6 +6,7 @@
 
 """
 
+import sqlalchemy as sa
 from alembic import op
 
 # revision identifiers, used by Alembic.
@@ -22,6 +23,26 @@ def upgrade():
 
 
 def downgrade():
+    # Null out any task_run.state_id values that reference non-existent
+    # task_run_state rows. After the FK was dropped (upgrade), the system no
+    # longer enforces referential integrity, so orphaned references are
+    # expected. Without this cleanup, re-adding the FK would fail with an
+    # IntegrityError.
+    connection = op.get_bind()
+    connection.execute(
+        sa.text(
+            """
+            UPDATE task_run
+            SET state_id = NULL
+            WHERE state_id IS NOT NULL
+              AND NOT EXISTS (
+                  SELECT 1 FROM task_run_state
+                  WHERE task_run_state.id = task_run.state_id
+              )
+            """
+        )
+    )
+
     # Re-add the foreign key constraint
     with op.batch_alter_table("task_run") as batch_op:
         batch_op.create_foreign_key(

src/prefect/server/database/interface.py

@@ -77,8 +77,38 @@ async def create_db(self) -> None:
         await self.run_migrations_upgrade()
 
     async def drop_db(self) -> None:
-        """Drop the database"""
-        await self.run_migrations_downgrade(revision="base")
+        """Drop the database by removing all tables directly.
+
+        This reflects the actual database schema and drops every table rather
+        than running all Alembic downgrade migrations in reverse.  Running
+        downgrades is fragile because individual migration downgrade steps may
+        fail on real-world data (e.g. re-adding a foreign key constraint when
+        orphaned references exist).  Dropping tables directly is both faster
+        and more robust.
+
+        Reflection is used instead of ``Base.metadata.drop_all()`` so that
+        tables created by migrations but not tracked in the ORM (e.g.
+        ``deployment_version``, ``alembic_version``) are also removed.
+        """
+        engine = await self.engine()
+        async with engine.begin() as conn:
+            # Disable FK checks for SQLite so that tables can be dropped in
+            # any order without triggering constraint errors.
+            dialect = get_dialect(self.database_config.connection_url)
+            is_sqlite = dialect.name == "sqlite"
+            if is_sqlite:
+                await conn.execute(sa.text("PRAGMA foreign_keys = OFF"))
+
+            try:
+                # Reflect the actual database schema so we capture every
+                # table, including migration-only tables not present in the
+                # ORM metadata.
+                metadata = sa.MetaData()
+                await conn.run_sync(metadata.reflect)
+                await conn.run_sync(metadata.drop_all)
+            finally:
+                if is_sqlite:
+                    await conn.execute(sa.text("PRAGMA foreign_keys = ON"))
 
     async def run_migrations_upgrade(self) -> None:
         """Run all upgrade migrations"""

tests/server/database/test_migrations.py

@@ -936,3 +936,105 @@ async def test_migrate_flow_run_notifications_to_automations(db: PrefectDBInterf
             ]
     finally:
         await run_sync_in_worker_thread(alembic_upgrade)
+
+
+async def test_downgrade_with_orphaned_task_run_state_ids(db, flow):
+    """
+    Tests that migration 3b86c5ea017a can be downgraded even when
+    task_run.state_id references a non-existent task_run_state row.
+
+    This is the scenario reported in GitHub issue #20939: after the FK
+    constraint was dropped (upgrade), the system no longer enforces
+    referential integrity. Orphaned state_id values are expected in
+    production. The downgrade must clean them up before re-adding the FK.
+    """
+    connection_url = PREFECT_API_DATABASE_CONNECTION_URL.value()
+    dialect = get_dialect(connection_url)
+
+    if dialect.name == "postgresql":
+        # revision just after the FK-drop migration
+        fk_drop_revision = "3b86c5ea017a"
+        pre_fk_drop_revision = "aa1234567890"
+    else:
+        fk_drop_revision = "8bb517bae6f9"
+        pre_fk_drop_revision = "bb2345678901"
+
+    flow_run_id = uuid4()
+    task_run_id = uuid4()
+    orphaned_state_id = uuid4()  # Does NOT exist in task_run_state
+
+    try:
+        # Downgrade to the FK-drop migration (FK is removed at this point)
+        await run_sync_in_worker_thread(alembic_downgrade, revision=fk_drop_revision)
+
+        # Insert a task_run with an orphaned state_id (no corresponding
+        # task_run_state row). This simulates real-world data after the FK
+        # was dropped.
+        session = await db.session()
+        async with session:
+            await session.execute(
+                sa.text(
+                    "INSERT INTO flow_run (id, name, flow_id)"
+                    f" VALUES ('{flow_run_id}', 'test-flow-run', '{flow.id}');"
+                )
+            )
+            await session.execute(
+                sa.text(
+                    "INSERT INTO task_run (id, name, task_key, dynamic_key,"
+                    f" flow_run_id, state_id) VALUES ('{task_run_id}',"
+                    f" 'test-task', 'test-task', '0', '{flow_run_id}',"
+                    f" '{orphaned_state_id}');"
+                )
+            )
+            await session.commit()
+
+        # Downgrade past the FK-drop migration — this re-adds the FK.
+        # Before the fix, this would fail with IntegrityError.
+        await run_sync_in_worker_thread(
+            alembic_downgrade, revision=pre_fk_drop_revision
+        )
+
+        # Verify the orphaned state_id was cleaned up (nulled)
+        session = await db.session()
+        async with session:
+            result = (
+                await session.execute(
+                    sa.text(f"SELECT state_id FROM task_run WHERE id = '{task_run_id}'")
+                )
+            ).scalar()
+            assert result is None, "Orphaned state_id should be nulled during downgrade"
+    finally:
+        await run_sync_in_worker_thread(alembic_upgrade)
+
+
+async def test_drop_db_removes_all_tables(db):
+    """
+    Tests that drop_db() successfully removes all tables even when the
+    database contains data that would cause downgrade migrations to fail.
+
+    drop_db() uses metadata.drop_all() instead of running the full Alembic
+    downgrade chain, making it immune to data-dependent migration failures.
+    """
+    engine = await db.engine()
+
+    # Verify tables exist before drop
+    async with engine.connect() as conn:
+        table_names = await conn.run_sync(
+            lambda sync_conn: sa.inspect(sync_conn).get_table_names()
+        )
+    assert len(table_names) > 0, "Database should have tables before drop"
+
+    # Drop the database
+    await db.drop_db()
+
+    # Verify all ORM tables and alembic_version are gone
+    async with engine.connect() as conn:
+        table_names = await conn.run_sync(
+            lambda sync_conn: sa.inspect(sync_conn).get_table_names()
+        )
+    assert len(table_names) == 0, (
+        f"All tables should be dropped, but found: {table_names}"
+    )
+
+    # Recreate the database for other tests
+    await db.create_db()