Merge pull request #117 from ambitioninc/bug/json-cursor-wrap

Fix issue with json_cursor if django cursor is wrapped.

Author: geophphrie

README.rst

@@ -34,9 +34,8 @@ where django querysets simplify the sql generation process for simple queries.
 
 Requirements
 ------------
-* Python 2.7
-* Python 3.3, 3.4
-* Django 1.7+
+* Python 3.7 - 3.9
+* Django 2.2 - 4.1
 * Postgres 9.3+
 
 Installation

docs/release_notes.rst

@@ -1,6 +1,10 @@
 Release Notes
 =============
 
+v3.0.3
+------
+* Addressed bug in `json_cursor` if Django cursor has extra wrappers
+
 v3.0.2
 ------
 * Add `json_cursor` context to handle Django3.1.1+ no longer automatically parsing json fields

querybuilder/cursor.py

@@ -3,6 +3,45 @@
 import psycopg2
 
 
+def jsonify_cursor(django_cursor):
+    """
+    Adjust an already existing cursor to ensure it will return structured types (list or dict)
+    from jsonb columns instead of strings. Django 3.1.1+ returns strings for raw queries.
+    https://code.djangoproject.com/ticket/31956
+    https://code.djangoproject.com/ticket/31973
+    https://www.psycopg.org/docs/extras.html#psycopg2.extras.register_default_jsonb
+    """
+
+    # The thing that is returned by connection.cursor() is (normally) a Django object
+    # of type CursorWrapper that itself has the "real" cursor as a property called cursor.
+    # However, it could be a CursorDebugWrapper instead, or it could be an outer wrapper
+    # wrapping one of those. For example django-debug-toolbar wraps CursorDebugWrapper in
+    # a NormalCursorWrapper. The django-db-readonly package wraps the Django CursorWrapper
+    # in a ReadOnlyCursorWrapper. I'm not sure if they ever nest multiple levels. I tried
+    # looping with `while isinstance(inner_cursor, CursorWrapper)`, but it seems that the
+    # outer wrapper is not necessarily a subclass of the Django wrapper. My next best option
+    # is to make the assumption that we need to get to the last property called `cursor`,
+    # basically assuming that any wrapper is going to have a property called `cursor`
+    # that is the real cursor or the next-level wrapper.
+    # Another option might be to check the class of inner_cursor to see if it is the real
+    # database cursor. That would require importing more django libraries, and probably
+    # having to handle some changes in those libraries over different versions.
+
+    # We expect that there is always at least one wrapper, but we might as well handle
+    # the possibility that we get passed the inner cursor.
+    inner_cursor = django_cursor
+
+    while hasattr(inner_cursor, 'cursor'):
+        inner_cursor = inner_cursor.cursor
+
+    # Hopefully we have the right thing now, but try/catch so we can get a little better info
+    # if it is not. Another option might be an isinstance, or another function that tests the cursor?
+    try:
+        psycopg2.extras.register_default_jsonb(conn_or_curs=inner_cursor, loads=json.loads)
+    except TypeError as e:
+        raise Exception(f'jsonify_cursor: conn_or_curs was actually a {type(inner_cursor)}: {e}')
+
+
 @contextlib.contextmanager
 def json_cursor(django_database_connection):
     """
@@ -12,5 +51,5 @@ def json_cursor(django_database_connection):
     https://www.psycopg.org/docs/extras.html#psycopg2.extras.register_default_jsonb
     """
     with django_database_connection.cursor() as cursor:
-        psycopg2.extras.register_default_jsonb(conn_or_curs=cursor.cursor, loads=json.loads)
+        jsonify_cursor(cursor)
         yield cursor

querybuilder/query.py

@@ -8,12 +8,11 @@
 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
 from querybuilder.tables import TableFactory, ModelTable, QueryTable
+from querybuilder.cursor import jsonify_cursor
 
 SERIAL_DTYPES = ['serial', 'bigserial']
 
@@ -641,13 +640,8 @@ def get_cursor(self):
         :rtype: :class:`CursorDebugWrapper <django:django.db.backends.util.CursorDebugWrapper>`
         :returns: A database cursor
         """
-
-        # From Django 3.1.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)
+        jsonify_cursor(cursor)
         return cursor
 
     def from_table(self, table=None, fields='*', schema=None, **kwargs):

querybuilder/version.py

@@ -1 +1 @@
-__version__ = '3.0.2'
+__version__ = '3.0.3'