From d624ee8730d55ade5c883bd59ac471c0311f1d41 Mon Sep 17 00:00:00 2001
From: Alexandr Artemyev <mogost@gmail.com>
Date: Tue, 17 Sep 2024 14:59:11 +0500
Subject: [PATCH] Better already migrated detection (#588)

---
 .gitignore                               |  1 +
 constance/migrations/0003_drop_pickle.py | 27 ++++++++++++++----------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/.gitignore b/.gitignore
index 99840f04..a8df5eab 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,3 +10,4 @@ test.db
 coverage.xml
 docs/_build
 .idea
+constance/_version.py
diff --git a/constance/migrations/0003_drop_pickle.py b/constance/migrations/0003_drop_pickle.py
index 5e37a229..4abd268d 100644
--- a/constance/migrations/0003_drop_pickle.py
+++ b/constance/migrations/0003_drop_pickle.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import pickle
 from base64 import b64decode
@@ -11,6 +12,16 @@
 logger = logging.getLogger(__name__)
 
 
+def is_already_migrated(value):
+    try:
+        data = json.loads(value)
+        if isinstance(data, dict) and set(data.keys()) == {'__type__', '__value__'}:
+            return True
+    except (json.JSONDecodeError, TypeError):
+        return False
+    return False
+
+
 def import_module_attr(path):
     package, module = path.rsplit('.', 1)
     return getattr(import_module(package), module)
@@ -20,8 +31,9 @@ def migrate_pickled_data(apps, schema_editor) -> None:  # pragma: no cover
     Constance = apps.get_model('constance', 'Constance')
 
     for constance in Constance.objects.exclude(value=None):
-        constance.value = dumps(pickle.loads(b64decode(constance.value.encode())))  # noqa: S301
-        constance.save(update_fields=['value'])
+        if not is_already_migrated(constance.value):
+            constance.value = dumps(pickle.loads(b64decode(constance.value.encode())))  # noqa: S301
+            constance.save(update_fields=['value'])
 
     if settings.BACKEND in ('constance.backends.redisd.RedisBackend', 'constance.backends.redisd.CachingRedisBackend'):
         import redis
@@ -39,15 +51,8 @@ def migrate_pickled_data(apps, schema_editor) -> None:  # pragma: no cover
         for key in settings.CONFIG:
             prefixed_key = f'{_prefix}{key}'
             value = _rd.get(prefixed_key)
-            if value is not None:
-                try:
-                    redis_migrated_data[prefixed_key] = dumps(pickle.loads(value))  # noqa: S301
-                except pickle.UnpicklingError as e:
-                    if value.startswith(b'{"__'):
-                        # Seems like we're facing already migrated data
-                        # Might be related to defaults and when config was accessed while django inits for migration
-                        continue
-                    raise e
+            if value is not None and not is_already_migrated(value):
+                redis_migrated_data[prefixed_key] = dumps(pickle.loads(value))  # noqa: S301
         for prefixed_key, value in redis_migrated_data.items():
             _rd.set(prefixed_key, value)