From 9e31d79b21b9343c7477c4b62ebb309eb5467363 Mon Sep 17 00:00:00 2001 From: Marcos Estevez Date: Mon, 8 Jan 2024 17:43:17 +0100 Subject: [PATCH] Support Truncate with Cascade --- dj_anonymizer/anonymizer.py | 5 ++++- dj_anonymizer/register_models.py | 5 ++++- dj_anonymizer/utils.py | 2 +- docs/changelog.rst | 4 ++++ docs/conf.py | 2 +- docs/usage.rst | 16 ++++++++++++++-- setup.py | 2 +- tests/test_register_models.py | 27 +++++++++++++++++++++++++-- tests/test_utils.py | 5 +++-- 9 files changed, 57 insertions(+), 11 deletions(-) diff --git a/dj_anonymizer/anonymizer.py b/dj_anonymizer/anonymizer.py index c3eb678..487106e 100644 --- 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 index 22abf67..1f6bcd1 100644 --- 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 index 323f58f..149bac6 100644 --- a/dj_anonymizer/utils.py +++ b/dj_anonymizer/utils.py @@ -53,7 +53,7 @@ def truncate_table(model, cascade=False): cascade_op = VENDOR_TO_CASCADE[vendor] except KeyError: raise NotImplementedError( - "Database vendor %s does not support TRUNCATE with CASCADE" % vendor + "DB vendor %s does not support TRUNCATE with CASCADE" % vendor ) dbtable = '"{}"'.format(model._meta.db_table) diff --git a/docs/changelog.rst b/docs/changelog.rst index d20a6fd..3d10e7b 100644 --- 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 `__) + 0.5.1 ---------- * Bugfix: fix issue when model manager is overridden (`#71 `__) diff --git a/docs/conf.py b/docs/conf.py index fc5c112..75434ee 100644 --- 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 index c0fb032..558c805 100644 --- 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 index 7d2cb35..0e98cb8 100644 --- 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', 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 index a91b5f9..6594d6d 100644 --- 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,6 +246,10 @@ 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" @@ -249,9 +257,24 @@ def test_register_clean(): 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 index fd2d35b..0a8aa7c 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -46,10 +46,11 @@ def test_truncate_table_with_cascade(mock_connections): 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_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) -