misc: use base64 to store PickledObjectField content (#41235)

src/authentic2/saml/fields.py

@@ -14,28 +14,37 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import base64
 try:
     import cPickle as pickle
 except ImportError:
     import pickle
-import six
 
 from django import forms
 from django.db import models
-from django.db.models.lookups import Exact, In
 from django.core.exceptions import ValidationError
+from django.utils import six
 from django.utils.text import capfirst
-from django.utils.encoding import force_bytes, force_str
+from django.utils.encoding import force_bytes, force_text
 from django.contrib.humanize.templatetags.humanize import apnumber
 from django.template.defaultfilters import pluralize
 
 
 def loads(value):
-    return pickle.loads(force_bytes(str(value)))
+    # value is always an unicode string
+    value = force_bytes(value)
+    try:
+        value = base64.b64decode(value)
+    except ValueError:
+        # there can be old values which are just pickle
+        pass
+    return pickle.loads(value)
 
 
 def dumps(value):
-    return force_str(pickle.dumps(value, protocol=0))
+    return PickledObject(
+        force_text(
+            base64.b64encode(pickle.dumps(value, protocol=0))))
 
 # This is a copy of http://djangosnippets.org/snippets/513/
 #
@@ -50,66 +59,36 @@ def dumps(value):
 # Initial author: Oliver Beattie
 
 
-class PickledObject(str):
+class PickledObject(six.text_type):
     """A subclass of string so it can be told whether a string is
        a pickled object or not (if the object is an instance of this class
        then it must [well, should] be a pickled one)."""
     pass
 
-try:
-    # make PickledObject compatible with Postgres
-    import psycopg2.extensions
-    psycopg2.extensions.register_adapter(PickledObject, psycopg2.extensions.QuotedString)
-except ImportError:
-    pass
-
-
-def do_pickle(value):
-    if value is not None and not isinstance(value, PickledObject):
-        value = PickledObject(dumps(value))
-    return value
-
 
 class PickledObjectField(models.Field):
-
-    def to_python(self, value):
-        if isinstance(value, PickledObject):
-            # If the value is a definite pickle; and an error is raised in
-            # de-pickling it should be allowed to propogate.
-            return loads(value)
-        else:
-            return loads(value)
-
     def from_db_value(self, value, expression, connection, context):
-        return self.to_python(value)
+        # Reading value from db
+        if value is not None:
+            value = loads(value)
+        return value
 
-    def get_db_prep_save(self, value, connection):
-        return do_pickle(value)
+    def get_prep_value(self, value):
+        # Preparing value for db
+        if value is not None and not isinstance(value, PickledObject):
+            value = dumps(value)
+        return value
 
     def get_internal_type(self):
         return 'TextField'
 
-    def value_to_string(self, obj):
-        return PickledObject(dumps(obj))
+    def get_lookup(self, lookup_name):
+        """
+        No lookup is possible.
+        """
+        raise TypeError('Lookup type %s is not supported.' % lookup_name)
 
 
-class PickledObjectFieldExact(Exact):
-
-    def get_db_prep_lookup(self, value, connection):
-        value = do_pickle(value)
-        return super(PickledObjectFieldExact, self).get_db_prep_lookup(value, connection)
-
-
-class PickledObjectFieldIn(In):
-
-    def get_db_prep_lookup(self, value, connection):
-        value = [do_pickle(v) for v in value]
-        return super(PickledObjectField, self).get_db_prep_lookup(value, connection)
-
-
-PickledObjectField.register_lookup(PickledObjectFieldExact)
-PickledObjectField.register_lookup(PickledObjectFieldIn)
-
 # This is a modified copy of http://djangosnippets.org/snippets/1200/
 #
 # We added a validate method.

tests/test_customfields.py

@@ -40,15 +40,6 @@ def test_pickled_data_integrity(value, db):
     assert KeyValue.objects.get().value == value
 
 
-def test_pickled_lookups(db, testing_data):
-    """Tests that lookups can be performed on data once stored in the database."""
-    for value in testing_data:
-        model_test = KeyValue(value=value)
-        model_test.save()
-        assert value == KeyValue.objects.get(value__exact=value).value
-        model_test.delete()
-
-
 def test_multiselectfield_data_integrity(db):
     spp = SPOptionsIdPPolicy.objects.create(name='spp')
     value = [x[0] for x in NAME_ID_FORMATS_CHOICES]