11import contextlib
22import json
3- import psycopg2
3+ from psycopg2 .extras import register_default_jsonb
4+ from psycopg2 ._json import JSONB_OID
5+
6+
7+ def jsonify_cursor (django_cursor , enabled = True ):
8+ """
9+ Adjust an already existing cursor to ensure it will return structured types (list or dict)
10+ from jsonb columns instead of strings. Django 3.1.1+ returns strings for raw queries.
11+ https://code.djangoproject.com/ticket/31956
12+ https://code.djangoproject.com/ticket/31973
13+ https://www.psycopg.org/docs/extras.html#psycopg2.extras.register_default_jsonb
14+ """
15+
16+ # The thing that is returned by connection.cursor() is (normally) a Django object
17+ # of type CursorWrapper that itself has the "real" cursor as a property called cursor.
18+ # However, it could be a CursorDebugWrapper instead, or it could be an outer wrapper
19+ # wrapping one of those. For example django-debug-toolbar wraps CursorDebugWrapper in
20+ # a NormalCursorWrapper. The django-db-readonly package wraps the Django CursorWrapper
21+ # in a ReadOnlyCursorWrapper. I'm not sure if they ever nest multiple levels. I tried
22+ # looping with `while isinstance(inner_cursor, CursorWrapper)`, but it seems that the
23+ # outer wrapper is not necessarily a subclass of the Django wrapper. My next best option
24+ # is to make the assumption that we need to get to the last property called `cursor`,
25+ # basically assuming that any wrapper is going to have a property called `cursor`
26+ # that is the real cursor or the next-level wrapper.
27+ # Another option might be to check the class of inner_cursor to see if it is the real
28+ # database cursor. That would require importing more django libraries, and probably
29+ # having to handle some changes in those libraries over different versions.
30+
31+ # This register_default_jsonb functionality in psycopg2 does not itself have a "deregister"
32+ # capability. So to deregister, we pass in a different value for the loads method; in this
33+ # case just the str() built-in, which just returns the value passed in. Note that passing
34+ # None for loads does NOT do a deregister; it uses the default value, which as it turns out
35+ # is json.loads anyway!
36+ loads_func = json .loads if enabled else str
37+
38+ # We expect that there is always at least one wrapper, but we might as well handle
39+ # the possibility that we get passed the inner cursor.
40+ inner_cursor = django_cursor
41+
42+ while hasattr (inner_cursor , 'cursor' ):
43+ inner_cursor = inner_cursor .cursor
44+
45+ # Hopefully we have the right thing now, but try/catch so we can get a little better info
46+ # if it is not. Another option might be an isinstance, or another function that tests the cursor?
47+ try :
48+ register_default_jsonb (conn_or_curs = inner_cursor , loads = loads_func )
49+ except TypeError as e :
50+ raise Exception (f'jsonify_cursor: conn_or_curs was actually a { type (inner_cursor )} { e } )
51+
52+
53+ def dejsonify_cursor (django_cursor ):
54+ """
55+ Re-adjust a cursor that was "jsonified" so it no longer performs the json.loads().
56+ """
57+ jsonify_cursor (django_cursor , enabled = False )
458
559
660@contextlib .contextmanager
@@ -12,5 +66,62 @@ def json_cursor(django_database_connection):
1266 https://www.psycopg.org/docs/extras.html#psycopg2.extras.register_default_jsonb
1367 """
1468 with django_database_connection .cursor () as cursor :
15- psycopg2 . extras . register_default_jsonb ( conn_or_curs = cursor . cursor , loads = json . loads )
69+ jsonify_cursor ( cursor )
1670 yield cursor
71+ # This should really not be necessary, because the cursor context manager will
72+ # be closing the cursor on __exit__ anyway. But just in case.
73+ dejsonify_cursor (cursor )
74+
75+
76+ def json_fetch_all_as_dict (cursor ):
77+ """
78+ Iterates over a result set and converts each row to a dictionary.
79+ The cursor passed in is assumed to have just executed a raw Postgresql query.
80+ If the cursor's columns include any with the jsonb type, the process includes
81+ examining every value from those columns. If the value is a string, a json.loads()
82+ is attempted on the value, because in Django 3.1.1 and later, this is not
83+ handled automatically for raw sql as it was before. There is no compatibility
84+ issue running with older Django versions because if the value is not a string,
85+ (e.g. it has already been converted to a list or dict), the loads() is skipped.
86+ Note that JSON decoding errors are ignored (and the original result value is provided)
87+ because it is possible that the query involved an actual json query, say on a single
88+ string property of the underlying column data. In that case, the column type is
89+ still jsonb, but the result value is a string as it should be. This ignoring of
90+ errors is the same logic used in json handling in Django's from_db_value() method.
91+
92+ :return: A list of dictionaries where each row is a dictionary
93+ :rtype: list of dict
94+ """
95+
96+ colnames = [col .name for col in cursor .description ]
97+ coltypes = [col .type_code for col in cursor .description ]
98+ # Identify any jsonb columns in the query, by column index
99+ jsonbcols = [i for i , x in enumerate (coltypes ) if x == JSONB_OID ]
100+
101+ # Optimize with a simple comprehension if we know there are no jsonb columns to handle.
102+ if not jsonbcols :
103+ return [
104+ dict (zip (colnames , row ))
105+ for row in cursor .fetchall ()
106+ ]
107+
108+ # If there are jsonb columns, intercept the result rows and run a json.loads() on any jsonb
109+ # columns that are presenting as strings.
110+ # In Django 3.1.0 they would already be a json type (e.g. dict or list) but in Django 3.1.1 it changes
111+ # and raw sql queries return strings for jsonb columns.
112+ # https://docs.djangoproject.com/en/4.0/releases/3.1.1/
113+ results = []
114+
115+ for row in cursor .fetchall ():
116+ rowvals = list (row )
117+ for colindex in jsonbcols :
118+ if type (rowvals [colindex ]) is str : # need to check type to avoid attempting to jsonify a None
119+ try :
120+ rowvals [colindex ] = json .loads (rowvals [colindex ])
121+ # It is possible that we are selecting a sub-value from the json in the column. I.e.
122+ # we got here because it IS a jsonb column, but what we selected is not json and will
123+ # fail to parse. In that case, we already have the value we want in place.
124+ except json .JSONDecodeError :
125+ pass
126+ results .append (dict (zip (colnames , rowvals )))
127+ return results
0 commit comments