diff --git a/docs/source/deletion.rst b/docs/source/deletion.rst index c27cdcb6..9308594c 100644 --- a/docs/source/deletion.rst +++ b/docs/source/deletion.rst @@ -48,3 +48,28 @@ By default, Postgres will raise an error if any other table is referencing one o MyModel.objects.truncate(cascade=True) print(MyModel1.objects.count()) # zero records left print(MyModel2.objects.count()) # zero records left + + +Restart identity +**************** + +If specified, any sequences on the table will be restarted. + +.. code-block:: python + + from django.db import models + from psqlextra.models import PostgresModel + + class MyModel(PostgresModel): + pass + + mymodel = MyModel.objects.create() + assert mymodel.id == 1 + + MyModel.objects.truncate(restart_identity=True) # table is empty after this + print(MyModel.objects.count()) # zero records left + + # Create a new row, it should get ID 1 again because + # the sequence got restarted. + mymodel = MyModel.objects.create() + assert mymodel.id == 1 diff --git a/psqlextra/manager/manager.py b/psqlextra/manager/manager.py index 0931b38a..ee1eb58b 100644 --- a/psqlextra/manager/manager.py +++ b/psqlextra/manager/manager.py @@ -37,7 +37,10 @@ def __init__(self, *args, **kwargs): ) def truncate( - self, cascade: bool = False, using: Optional[str] = None + self, + cascade: bool = False, + restart_identity: bool = False, + using: Optional[str] = None, ) -> None: """Truncates this model/table using the TRUNCATE statement. @@ -51,14 +54,19 @@ def truncate( False, an error will be raised if there are rows in other tables referencing the rows you're trying to delete. + restart_identity: + Automatically restart sequences owned by + columns of the truncated table(s). """ connection = connections[using or "default"] table_name = connection.ops.quote_name(self.model._meta.db_table) with connection.cursor() as cursor: - sql = "TRUNCATE TABLE %s" % table_name + sql = f"TRUNCATE TABLE {table_name}" if cascade: sql += " CASCADE" + if restart_identity: + sql += " RESTART IDENTITY" cursor.execute(sql)