Merge pull request #3 from M1hacka/save-bytes-to-hll

Added ability to save HllField as it has been got from database

setup.py

@@ -13,7 +13,7 @@
 
 setup(
     name='django-pg-hll',
-    version='1.0.1',
+    version='1.1.0',
     packages=['django_pg_hll'],
     package_dir={'': 'src'},
     url='https://github.com/M1hacka/django-pg-hll',

src/django_pg_hll/fields.py

@@ -4,7 +4,9 @@
 from django.db.models import BinaryField, IntegerField
 from django.db.models.lookups import Transform
 
-from .values import HllEmpty
+import six
+
+from .values import HllEmpty, HllFromHex
 
 __all__ = ['HllField']
 
@@ -51,6 +53,15 @@ def rel_db_type(self, connection):
     def get_internal_type(self):
         return self.__class__.__name__
 
+    def get_db_prep_value(self, value, connection, prepared=False):
+        # Psycopg2 returns Binary results as hex string, prefixed by \x
+        # BinaryField requires bytes to be saved
+        # But none of these can be converted to HLL by postgres directly
+        if isinstance(value, bytes) or isinstance(value, six.string_types) and value.startswith(r'\x'):
+            return HllFromHex(value, db_type=self.db_type(connection))
+        else:
+            return super(HllField, self).get_db_prep_value(value, connection, prepared=prepared)
+
     def get_default(self):
         if self.has_default() and not callable(self.default):
             return self.default

src/django_pg_hll/values.py

@@ -27,6 +27,26 @@ class HllCombinedExpression(HllJoinMixin, CombinedExpression):
     pass
 
 
+class HllFromHex(six.with_metaclass(ABCMeta, Func)):
+    """
+    Constructs hll that can be saved from binary data (or it's psycopg representation)
+    """
+    def __init__(self, data, *args, **extra):
+        db_type = extra.pop('db_type', 'hll')
+
+        # Psycopg2 returns Binary results as hex string, prefixed by \x but requires bytes for saving.
+        if isinstance(data, six.string_types) and data.startswith(r'\x'):
+            data = bytearray.fromhex(data[2:])
+        elif isinstance(data, bytes):
+            pass
+        else:
+            raise ValueError('data should be bytes instance or string starting with \\x')
+
+        self.template = extra.get('template', '%(expressions)s::{}'.format(db_type))
+
+        super(HllFromHex, self).__init__(Value(data), *args, **extra)
+
+
 class HllValue(six.with_metaclass(ABCMeta, HllJoinMixin, Func)):
     pass
 
@@ -77,9 +97,6 @@ def parse_data(cls, data):  # type: (Any) -> HllDataValue
 
 
 class HllPrimitiveValue(six.with_metaclass(ABCMeta, HllDataValue)):
-    # Abstract class property
-    db_type = None
-
     def __init__(self, data, **extra):  # type: (Any, **dict) -> None
         """
         :param data: Data to build value from

tests/test_hll_field.py

@@ -1,5 +1,6 @@
 from unittest import skipIf
 
+import six
 from django.db import connection
 from django.db.models import F
 from django.test import TestCase
@@ -70,6 +71,21 @@ def test_update(self):
         TestModel.objects.filter(id=100501).update(hll_field=HllInteger(1) | F('hll_field'))
         self.assertEqual(1, TestModel.objects.annotate(card=Cardinality('hll_field')).filter(id=100501).
                          values_list('card', flat=True)[0])
+
+    def test_hex_convertion(self):
+        instance = TestModel.objects.get(id=100501)
+        instance.hll_field = HllInteger(1) | F('hll_field')
+        instance.save()
+
+        instance.refresh_from_db()
+
+        self.assertIsInstance(instance.hll_field, six.string_types)
+        self.assertEqual(instance.hll_field[:2], r'\x')
+
+        instance.save()
+
+        self.assertEqual(1, TestModel.objects.annotate(card=Cardinality('hll_field')).filter(id=100501).
+                         values_list('card', flat=True)[0])
 
 
 class TestAggregation(TestCase):