Improve truncate_models for cascade delete (#2100)

* Improve truncate_models for cascade delete

* Fix

Author: abondar

docs/contrib/unittest.rst

@@ -246,7 +246,12 @@ Utility Functions
 truncate_all_models
 -------------------
 
-Truncate all model tables in the current context:
+Truncate all model tables in the current context. The function handles foreign key
+constraints automatically:
+
+- **PostgreSQL**: Uses a single ``TRUNCATE ... CASCADE`` statement (fast, single round-trip).
+- **Other databases**: Deletes in topological order — child tables are emptied before the
+  parent tables they reference, avoiding FK constraint violations.
 
 .. code-block:: python
 
@@ -257,7 +262,7 @@ Truncate all model tables in the current context:
         # Create some data
         await User.create(name="Test")
 
-        # Truncate all tables
+        # Truncate all tables (FK-safe)
         await truncate_all_models()
 
         # Tables are now empty

tests/cli/test_cli.py

@@ -203,6 +203,7 @@ async def fake_migrate(**kwargs) -> None:
 
 @pytest.mark.asyncio
 async def test_upgrade_alias(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    _write_package(tmp_path, "cli_app")
     module_name = _write_settings(
         tmp_path,
         """
@@ -234,6 +235,7 @@ async def fake_migrate(**kwargs) -> None:
 
 @pytest.mark.asyncio
 async def test_downgrade_defaults_to_first(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    _write_package(tmp_path, "cli_app")
     module_name = _write_settings(
         tmp_path,
         """

tests/fields/test_time.py

@@ -80,8 +80,8 @@ async def test_datetime_create(db):
     obj = await model.get(id=obj0.id)
     assert obj.datetime == now
     assert obj.datetime_null is None
-    assert obj.datetime_auto - now < timedelta(microseconds=20000)
-    assert obj.datetime_add - now < timedelta(microseconds=20000)
+    assert obj.datetime_auto - now < timedelta(seconds=1)
+    assert obj.datetime_add - now < timedelta(seconds=1)
     datetime_auto = obj.datetime_auto
     sleep(0.012)
     await obj.save()

tests/test_truncate.py

@@ -0,0 +1,131 @@
+"""
+Tests for truncate_all_models() and _topological_sort_models().
+
+Verifies FK-aware truncation ordering and the PostgreSQL TRUNCATE CASCADE path.
+"""
+
+import pytest
+
+from tests.testmodels import Employee, Event, MinRelation, Reporter, Team, Tournament
+from tortoise import Tortoise
+from tortoise.contrib.test import _topological_sort_models, truncate_all_models
+
+# ---------------------------------------------------------------------------
+# _topological_sort_models — unit tests on real model metadata
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_topological_sort_children_before_parents(db):
+    """Event (FK→Tournament) must come before Tournament in delete order."""
+    sorted_models = _topological_sort_models([Tournament, Event])
+    assert sorted_models.index(Event) < sorted_models.index(Tournament)
+
+
+@pytest.mark.asyncio
+async def test_topological_sort_input_order_independent(db):
+    """Result must be the same regardless of input order."""
+    order_a = _topological_sort_models([Tournament, Event])
+    order_b = _topological_sort_models([Event, Tournament])
+    assert order_a == order_b
+
+
+@pytest.mark.asyncio
+async def test_topological_sort_self_referential_fk(db):
+    """Self-referential FK (Employee→Employee) must not cause infinite loop."""
+    result = _topological_sort_models([Employee])
+    assert result == [Employee]
+
+
+@pytest.mark.asyncio
+async def test_topological_sort_no_fk_models(db):
+    """Models without FK relationships are still included."""
+    result = _topological_sort_models([Team])
+    assert result == [Team]
+
+
+@pytest.mark.asyncio
+async def test_topological_sort_all_models(db):
+    """Sorting all registered models succeeds and includes every model."""
+    all_models = list(Tortoise.apps.get_models_iterable())
+    sorted_models = _topological_sort_models(all_models)
+    assert set(sorted_models) == set(all_models)
+
+
+@pytest.mark.asyncio
+async def test_topological_sort_multi_level_chain(db):
+    """MinRelation→Tournament and MinRelation→Team: MinRelation before both parents."""
+    sorted_models = _topological_sort_models([Tournament, Team, MinRelation])
+    assert sorted_models.index(MinRelation) < sorted_models.index(Tournament)
+    assert sorted_models.index(MinRelation) < sorted_models.index(Team)
+
+
+@pytest.mark.asyncio
+async def test_topological_sort_multiple_fks_on_one_model(db):
+    """Event has FKs to both Tournament and Reporter — must come before both."""
+    sorted_models = _topological_sort_models([Tournament, Reporter, Event])
+    assert sorted_models.index(Event) < sorted_models.index(Tournament)
+    assert sorted_models.index(Event) < sorted_models.index(Reporter)
+
+
+# ---------------------------------------------------------------------------
+# truncate_all_models — integration tests against real DB
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_truncate_empty_db(db):
+    """Truncating when tables are empty should succeed without error."""
+    await truncate_all_models()
+
+
+@pytest.mark.asyncio
+async def test_truncate_clears_data(db):
+    """Data created before truncation is gone after truncation."""
+    tournament = await Tournament.create(name="Test Tournament")
+    await Event.create(name="Test Event", tournament=tournament)
+
+    await truncate_all_models()
+
+    assert await Tournament.all().count() == 0
+    assert await Event.all().count() == 0
+
+
+@pytest.mark.asyncio
+async def test_truncate_with_fk_constraints(db):
+    """Truncation succeeds even with FK constraints (child→parent)."""
+    t = await Tournament.create(name="T1")
+    await Event.create(name="E1", tournament=t)
+    await Event.create(name="E2", tournament=t)
+
+    # This would fail with arbitrary order on strict FK enforcement
+    await truncate_all_models()
+
+    assert await Event.all().count() == 0
+    assert await Tournament.all().count() == 0
+
+
+@pytest.mark.asyncio
+async def test_truncate_with_self_referential_fk(db):
+    """Self-referential FK (Employee→Employee) doesn't break truncation."""
+    boss = await Employee.create(name="Boss")
+    await Employee.create(name="Worker", manager=boss)
+
+    await truncate_all_models()
+
+    assert await Employee.all().count() == 0
+
+
+@pytest.mark.asyncio
+async def test_truncate_raises_when_apps_not_loaded(db_simple):
+    """truncate_all_models raises ValueError when apps aren't loaded."""
+    from tortoise.context import get_current_context
+
+    ctx = get_current_context()
+    saved_apps = ctx._apps
+    ctx._apps = {}
+    try:
+        with pytest.raises(ValueError, match="apps are not loaded"):
+            await truncate_all_models()
+    finally:
+        ctx._apps = saved_apps

tortoise/contrib/test/__init__.py

@@ -28,13 +28,16 @@ async def test_sqlite_only(db):
 import typing
 from collections.abc import Callable, Coroutine
 from functools import partial, wraps
-from typing import ParamSpec, TypeVar, cast
+from typing import TYPE_CHECKING, ParamSpec, TypeVar, cast
 from unittest import SkipTest, expectedFailure, skip, skipIf, skipUnless
 
 from tortoise import Tortoise
 from tortoise.connection import get_connection
 from tortoise.context import TortoiseContext, tortoise_test_context
 
+if TYPE_CHECKING:
+    from tortoise.models import Model
+
 T = TypeVar("T")
 P = ParamSpec("P")
 AsyncFunc = Callable[P, Coroutine[None, None, T]]
@@ -70,19 +73,88 @@ async def truncate_all_models() -> None:
     This is a utility function for test cleanup that deletes all rows from
     all registered model tables.
 
-    Note: This is a naive implementation that may fail with M2M relations
-    and non-cascade foreign keys.
+    On PostgreSQL, uses ``TRUNCATE ... CASCADE`` for a single fast statement.
+    On other databases, deletes in topological (FK dependency) order so that
+    child rows are removed before parent rows they reference.
 
     Raises:
         ValueError: If Tortoise.apps is not loaded.
     """
     if not Tortoise.apps:
         raise ValueError("apps are not loaded")
-    for model in Tortoise.apps.get_models_iterable():
-        quote_char = model._meta.db.query_class.SQL_CONTEXT.quote_char
-        await model._meta.db.execute_script(
-            f"DELETE FROM {quote_char}{model._meta.db_table}{quote_char}"  # nosec
-        )
+
+    models = list(Tortoise.apps.get_models_iterable())
+
+    if not models:
+        return
+
+    db = models[0]._meta.db
+    dialect = db.capabilities.dialect
+
+    if dialect == "postgres":
+        # PostgreSQL supports TRUNCATE with CASCADE — single statement, fast
+        tables = ", ".join(f'"{m._meta.db_table}"' for m in models)
+        await db.execute_script(f"TRUNCATE {tables} CASCADE")
+    else:
+        # For other dialects, topologically sort by FK dependencies (children first)
+        sorted_models = _topological_sort_models(models)
+
+        # Disable FK checks to handle self-referential and circular FK constraints
+        if dialect == "mysql":
+            await db.execute_script("SET FOREIGN_KEY_CHECKS = 0")
+        elif dialect == "sqlite":
+            await db.execute_script("PRAGMA foreign_keys = OFF")
+
+        try:
+            for model in sorted_models:
+                quote_char = model._meta.db.query_class.SQL_CONTEXT.quote_char
+                await model._meta.db.execute_script(
+                    f"DELETE FROM {quote_char}{model._meta.db_table}{quote_char}"  # nosec
+                )
+        finally:
+            if dialect == "mysql":
+                await db.execute_script("SET FOREIGN_KEY_CHECKS = 1")
+            elif dialect == "sqlite":
+                await db.execute_script("PRAGMA foreign_keys = ON")
+
+
+def _topological_sort_models(models: list[type[Model]]) -> list[type[Model]]:
+    """Sort models so children come before parents (safe delete order).
+
+    Uses Kahn's algorithm on FK dependencies. Models that depend on others
+    via ForeignKey are placed *before* the models they reference, ensuring
+    child rows are deleted before parent rows.
+    """
+    from tortoise.fields.relational import ForeignKeyFieldInstance
+
+    model_set = set(models)
+    # Build adjacency for delete order: parent -> children that must be deleted first
+    # If Event has FK to Tournament, then Tournament depends on Event being deleted first
+    deps: dict[type[Model], set[type[Model]]] = {m: set() for m in models}
+    for model in models:
+        for field in model._meta.fields_map.values():
+            if isinstance(field, ForeignKeyFieldInstance):
+                related = field.related_model
+                if related in model_set and related is not model:
+                    deps[related].add(model)
+
+    # Kahn's algorithm — emit models whose deps are already emitted
+    sorted_models: list[type[Model]] = []
+    no_deps = [m for m in models if not deps[m]]
+    while no_deps:
+        m = no_deps.pop()
+        sorted_models.append(m)
+        for other in models:
+            deps[other].discard(m)
+            if not deps[other] and other not in sorted_models and other not in no_deps:
+                no_deps.append(other)
+
+    # Append any remaining (circular deps — fallback)
+    for m in models:
+        if m not in sorted_models:
+            sorted_models.append(m)
+
+    return sorted_models
 
 
 _FT = TypeVar("_FT", bound=Callable[..., typing.Any])