JSON capable cursor (#112)

* Adjusted query.py to use the same technique as json_cursor to correctly handle jsonb columns in Django 3.1.1+
* Eliminated the changes that were working around test result diffs in Dj3.1+

Author: geophphrie

querybuilder/query.py

@@ -8,7 +8,8 @@
 from django.apps import apps
 get_model = apps.get_model
 import six
-
+import json
+import psycopg2
 
 from querybuilder.fields import FieldFactory, CountField, MaxField, MinField, SumField, AvgField
 from querybuilder.helpers import set_value_for_keypath, copy_instance
@@ -640,7 +641,14 @@ def get_cursor(self):
         :rtype: :class:`CursorDebugWrapper <django:django.db.backends.util.CursorDebugWrapper>`
         :returns: A database cursor
         """
-        return self.connection.cursor()
+
+        # From Django 3.1 forward, json columns in raw select statements return a string of json instead of a
+        # json type such as a dict or list. But we can tell psycopg2 to put the
+        # json.loads() call back in place. Technically we would only need this addition for cursors being used
+        # for a SELECT, but it should not cause any issues for other operations.
+        cursor = self.connection.cursor()
+        psycopg2.extras.register_default_jsonb(conn_or_curs=cursor.cursor, loads=json.loads)
+        return cursor
 
     def from_table(self, table=None, fields='*', schema=None, **kwargs):
         """

querybuilder/tests/json_tests.py

@@ -1,5 +1,4 @@
 import unittest
-from django import VERSION
 from django.test.utils import override_settings
 from querybuilder.fields import JsonField
 from querybuilder.query import Query, JsonQueryset
@@ -47,11 +46,7 @@ def test_one(self):
             )
         )
 
-        # Django 3.1 changes the raw queryset behavior so querybuilder isn't going to change that behavior
-        if self.is_31_or_above():
-            self.assertEqual(query.select(), [{'my_two_alias': '"two"'}])
-        else:
-            self.assertEqual(query.select(), [{'my_two_alias': 'two'}])
+        self.assertEqual(query.select(), [{'my_two_alias': 'two'}])
 
         query = Query().from_table(MetricRecord, fields=[one_field]).where(**{
             one_field.get_where_key(): '1'
@@ -64,11 +59,7 @@ def test_one(self):
             )
         )
 
-        # Django 3.1 changes the raw queryset behavior so querybuilder isn't going to change that behavior
-        if self.is_31_or_above():
-            self.assertEqual(query.select(), [{'my_one_alias': '1'}])
-        else:
-            self.assertEqual(query.select(), [{'my_one_alias': 1}])
+        self.assertEqual(query.select(), [{'my_one_alias': 1}])
 
         query = Query().from_table(MetricRecord, fields=[one_field]).where(**{
             one_field.get_where_key(): '2'
@@ -82,12 +73,14 @@ def test_one(self):
         )
         self.assertEqual(query.select(), [])
 
-    def is_31_or_above(self):
-        if VERSION[0] == 3 and VERSION[1] >= 1:
-            return True
-        elif VERSION[0] > 3:
-            return True
-        return False
+    # Currently unused (but maybe again sometime) function to check django version for test results
+    # from django import VERSION
+    # def is_31_or_above(self):
+    #     if VERSION[0] == 3 and VERSION[1] >= 1:
+    #         return True
+    #     elif VERSION[0] > 3:
+    #         return True
+    #     return False
 
 
 @override_settings(DEBUG=True)