Support Truncate with Cascade (#73)

* Support Truncate with Cascade --------- Co-authored-by: chinskiy <[email protected]>
preply · Jan 9, 2024 · 93d63c2 · 93d63c2
1 parent a6ea29a
commit 93d63c2
Show file tree

Hide file tree

Showing 9 changed files with 91 additions and 11 deletions.
diff --git a/dj_anonymizer/anonymizer.py b/dj_anonymizer/anonymizer.py
@@ -86,6 +86,9 @@ def clean(self, only=None):
         for queryset in clean_list:
             print(f'Cleaning {self.key(queryset.model)}')
             if getattr(queryset, 'truncate') is True:
-                truncate_table(queryset.model)
+                if getattr(queryset, 'cascade') is True:
+                    truncate_table(queryset.model, True)
+                else:
+                    truncate_table(queryset.model, False)
             else:
                 queryset.delete()
diff --git a/dj_anonymizer/register_models.py b/dj_anonymizer/register_models.py
@@ -13,9 +13,11 @@
 
 class AnonymBase:
     truncate = False
+    cascade = False
 
-    def __init__(self, truncate=False):
+    def __init__(self, truncate=False, cascade=False):
         self.truncate = truncate
+        self.cascade = cascade
 
     @classmethod
     def get_fields_names(cls):
@@ -122,6 +124,7 @@ def register_clean(models):
 
         queryset = model.objects.all()
         queryset.truncate = cls_anonym.truncate
+        queryset.cascade = cls_anonym.cascade
         if Anonymizer.key(model) in Anonymizer.clean_models.keys():
             raise ValueError(
                 f'Model {Anonymizer.key(model)} '

diff --git a/dj_anonymizer/utils.py b/dj_anonymizer/utils.py
@@ -13,6 +13,11 @@
     'oracle': 'TRUNCATE TABLE',
 }
 
+VENDOR_TO_CASCADE = {
+    'postgresql': 'CASCADE',
+    'oracle': 'CASCADE',
+}
+
 
 def import_if_exist(filename):
     """
@@ -27,7 +32,7 @@ def import_if_exist(filename):
         spec.loader.exec_module(mod)
 
 
-def truncate_table(model):
+def truncate_table(model, cascade=False):
     """
     Generate and execute via Django ORM proper SQL to truncate table
     """
@@ -42,11 +47,21 @@ def truncate_table(model):
             "Database vendor %s is not supported" % vendor
         )
 
+    cascade_op = ''
+    try:
+        if cascade:
+            cascade_op = VENDOR_TO_CASCADE[vendor]
+    except KeyError:
+        raise NotImplementedError(
+            "DB vendor %s does not support TRUNCATE with CASCADE" % vendor
+        )
+
     dbtable = '"{}"'.format(model._meta.db_table)
 
-    sql = '{operation} {dbtable}'.format(
+    sql = '{operation} {dbtable} {cascade}'.format(
         operation=operation,
         dbtable=dbtable,
-    )
+        cascade=cascade_op,
+    ).strip()
     with connection.cursor() as c:
         c.execute(sql)
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -1,6 +1,10 @@
 Changelog
 =============
 
+0.6.0
+----------
+* Feature: adds possibility to truncate with cascade option (`#73 <https://github.com/preply/dj_anonymizer/pull/73>`__)
+
 0.5.1
 ----------
 * Bugfix: fix issue when model manager is overridden (`#71 <https://github.com/preply/dj_anonymizer/pull/71>`__)

diff --git a/docs/conf.py b/docs/conf.py
@@ -26,7 +26,7 @@
 # The short X.Y version
 version = ''
 # The full version, including alpha/beta/rc tags
-release = '0.5.1'
+release = '0.6.0'
 
 
 # -- General configuration ---------------------------------------------------

diff --git a/docs/usage.rst b/docs/usage.rst
@@ -65,7 +65,7 @@ Model registration
 
    * `models` - list of tuples `(model, cls_anonym)`, where `model` is a model class and `cls_anonym` - anonymization class, inherited form `AnonymBase` with specified queryset for deletion or just `AnonymBase`.
 
-   If `AnonymBase` class have `truncate=True`, parameter table will be truncated instead of performing an SQL delete query.
+   If `AnonymBase` class have `truncate=True`, parameter table will be truncated instead of performing an SQL delete query. Additionally, `cascade=True` may be set to truncate foreign-key related tables (supported for postgresql and oracle).
 
 .. function:: register_skip(models)
 
@@ -186,7 +186,19 @@ Example 2 - truncate all data from model `User`::
         (User, AnonymBase(truncate=True)),
     ])
 
-Example 3 - delete all data from model `User`, except user with id=1::
+Example 3 - truncate all data from model `User` with cascade option::
+
+    from django.contrib.auth.models import User
+
+    from dj_anonymizer.register_models import AnonymBase
+    from dj_anonymizer.register_models import register_clean
+
+
+    register_clean([
+        (User, AnonymBase(truncate=True, cascade=True)),
+    ])
+
+Example 4 - delete all data from model `User`, except user with id=1::
 
     from django.contrib.auth.models import User
 

diff --git a/setup.py b/setup.py
@@ -13,7 +13,7 @@ def read(*rnames):
     name='dj_anonymizer',
     packages=['dj_anonymizer'],
     include_package_data=True,
-    version='0.5.1',
+    version='0.6.0-dev',
     description='This project helps anonymize production database '
                 + 'with fake data of any kind.',
     long_description=(read('README.md')),

diff --git a/tests/test_register_models.py b/tests/test_register_models.py
@@ -3,6 +3,7 @@
 
 import pytest
 from django.contrib.auth.models import Group, Permission, User
+from django.contrib.contenttypes.models import ContentType
 from django.db.models.query import QuerySet
 
 from dj_anonymizer import fields, register_models
@@ -216,10 +217,11 @@ def test_register_clean():
     register_models.register_clean([
         (User, register_models.AnonymBase),
         (Permission, register_models.AnonymBase(truncate=True)),
-        (Group, register_models.AnonymBase())
+        (Group, register_models.AnonymBase()),
+        (ContentType, register_models.AnonymBase(truncate=True, cascade=True)),
     ])
 
-    assert len(Anonymizer.clean_models) == 3
+    assert len(Anonymizer.clean_models) == 4
     assert len(Anonymizer.skip_models) == 0
     assert len(Anonymizer.anonym_models) == 0
 
@@ -229,6 +231,8 @@ def test_register_clean():
         Anonymizer.clean_models.keys()
     assert 'django.contrib.auth.models.Group' in \
         Anonymizer.clean_models.keys()
+    assert 'django.contrib.contenttypes.models.ContentType' in \
+        Anonymizer.clean_models.keys()
 
     assert isinstance(
         Anonymizer.clean_models['django.contrib.auth.models.User'],
@@ -242,16 +246,37 @@ def test_register_clean():
         Anonymizer.clean_models['django.contrib.auth.models.Group'],
         QuerySet
     )
+    assert isinstance(
+        Anonymizer.clean_models[
+            'django.contrib.contenttypes.models.ContentType'
+        ],
+        QuerySet
+    )
 
     assert Anonymizer.clean_models[
         "django.contrib.auth.models.User"
     ].model is User
     assert Anonymizer.clean_models[
         "django.contrib.auth.models.Permission"
     ].model is Permission
+    assert Anonymizer.clean_models[
+        "django.contrib.auth.models.Permission"
+    ].truncate is True
+    assert Anonymizer.clean_models[
+        "django.contrib.auth.models.Permission"
+    ].cascade is False
     assert Anonymizer.clean_models[
         "django.contrib.auth.models.Group"
     ].model is Group
+    assert Anonymizer.clean_models[
+        "django.contrib.contenttypes.models.ContentType"
+    ].model is ContentType
+    assert Anonymizer.clean_models[
+        "django.contrib.contenttypes.models.ContentType"
+    ].truncate is True
+    assert Anonymizer.clean_models[
+        "django.contrib.contenttypes.models.ContentType"
+    ].cascade is True
 
 
 @pytest.mark.django_db

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -36,3 +36,21 @@ def test_truncate_table(mock_connections):
 
     with pytest.raises(NotImplementedError):
         truncate_table(User)
+
+
+@mock.patch('dj_anonymizer.utils.connections')
+def test_truncate_table_with_cascade(mock_connections):
+    mock_cursor = mock_connections.\
+        __getitem__(DEFAULT_DB_ALIAS).\
+        cursor.return_value.__enter__.return_value
+    mock_connections.__getitem__(DEFAULT_DB_ALIAS).vendor = 'postgresql'
+
+    truncate_table(User, True)
+    mock_cursor.execute.assert_called_once_with(
+        'TRUNCATE TABLE "auth_user" CASCADE'
+    )
+
+    mock_connections.__getitem__(DEFAULT_DB_ALIAS).vendor = 'sqlite'
+
+    with pytest.raises(NotImplementedError):
+        truncate_table(User, True)