First attempt at supporting expressions in hstore

Author: Photonios

psqlextra/compiler.py

@@ -1,6 +1,8 @@
 from django.core.exceptions import SuspiciousOperation
 from django.db.models.sql.compiler import SQLInsertCompiler, SQLUpdateCompiler
 
+from psqlextra.expressions import HStoreValue
+
 
 class PostgresReturningUpdateCompiler(SQLUpdateCompiler):
     """Compiler for SQL UPDATE statements that return
@@ -16,6 +18,31 @@ def execute_sql(self, _result_type):
 
         return primary_keys
 
+    def as_sql(self):
+        self._prepare_query_values()
+        return super().as_sql()
+
+    def _prepare_query_values(self):
+        """Extra prep on query values by converting
+        dictionaries into :see:HStoreValue expressions.
+
+        This allows putting expressions in a dictionary.
+        The :see:HStoreValue will take care of resolving
+        the expressions inside the dictionary."""
+
+        new_query_values = []
+        for field, model, val in self.query.values:
+            if isinstance(val, dict):
+                val = HStoreValue(val)
+
+            new_query_values.append((
+                field,
+                model,
+                val
+            ))
+
+        self.query.values = new_query_values
+
     def _form_returning(self):
         """Builds the RETURNING part of the query."""
 

psqlextra/expressions.py

@@ -1,6 +1,62 @@
 from django.db.models import expressions, CharField
 
 
+class HStoreValue(expressions.Expression):
+    """Represents a HStore value.
+
+    The base PostgreSQL implementation Django provides,
+    always represents HStore values as dictionaries,
+    but this doesn't work if you want to use expressions
+    inside hstore values."""
+
+    def __init__(self, value):
+        """Initializes a new instance."""
+
+        self.value = value
+
+    def resolve_expression(self, *args, **kwargs):
+        """Resolves expressions inside the dictionary."""
+
+        result = dict()
+        for key, value in self.value.items():
+            if hasattr(value, 'resolve_expression'):
+                result[key] = value.resolve_expression(
+                    *args, **kwargs)
+            else:
+                result[key] = value
+
+        return HStoreValue(result)
+
+    def as_sql(self, compiler, connection):
+        """Compiles the HStore value into SQL.
+
+        Compiles expressions contained in the values
+        of HStore entries as well.
+
+        Given a dictionary like:
+
+            dict(key1='val1', key2='val2')
+
+        The resulting SQL will be:
+
+            hstore(ARRAY['key1', 'val1'], ARRAY['key2', 'val2'])
+        """
+
+        result = []
+        for key, value in self.value.items():
+            if hasattr(value, 'as_sql'):
+                sql, params = value.as_sql(compiler, connection)
+                result.append('ARRAY[\'%s\', %s]' % (
+                    key, sql % params))
+            elif value is not None:
+                result.append('ARRAY[\'%s\', \'%s\']' % ((
+                    key, value)))
+            else:
+                result.append('ARRAY[\'%s\', NULL]' % key)
+
+        return 'hstore(%s)' % ','.join(result), []
+
+
 class HStoreColumn(expressions.Col):
     """HStoreColumn expression.
 

psqlextra/fields/hstore_field.py

@@ -2,6 +2,8 @@
 
 from django.contrib.postgres.fields import HStoreField as DjangoHStoreField
 
+from psqlextra.expressions import HStoreValue
+
 
 class HStoreField(DjangoHStoreField):
     """Improved version of Django's :see:HStoreField that
@@ -22,6 +24,18 @@ def __init__(self, *args,
         self.uniqueness = uniqueness
         self.required = required
 
+    def get_db_prep_value(self, value, connection, prepared=False):
+        """Override the base class so it doesn't cast all values
+        to strings.
+
+        psqlextra supports expressions in hstore fields, so casting
+        all values to strings is a bad idea."""
+
+        if not value:
+            return None
+
+        return value
+
     def deconstruct(self):
         """Gets the values to pass to :see:__init__ when
         re-creating this object."""

tests/test_query.py

@@ -31,3 +31,25 @@ def test_annotate_hstore_key_ref():
     )
 
     assert queryset['english_title'] == 'english'
+
+
+def test_hstore_f_ref():
+    """Tests whether F(..) expressions can be used in
+    hstore values when performing update queries."""
+
+    model = get_fake_model({
+        'name': models.CharField(max_length=255),
+        'name_new': HStoreField()
+    })
+
+    model.objects.create(
+        name='waqas',
+        name_new=dict(en='swen')
+    )
+
+    model.objects.update(
+        name_new=dict(en=models.F('name'))
+    )
+
+    inst = model.objects.all().first()
+    assert inst.name_new.get('en') == 'waqas'