11import contextlib
22import json
3- import psycopg2
3+ from psycopg2 .extras import register_default_jsonb
4+ from psycopg2 ._json import JSONB_OID
45
56
6- def jsonify_cursor (django_cursor ):
7+ def jsonify_cursor (django_cursor , enabled = True ):
78 """
89 Adjust an already existing cursor to ensure it will return structured types (list or dict)
910 from jsonb columns instead of strings. Django 3.1.1+ returns strings for raw queries.
@@ -27,6 +28,13 @@ def jsonify_cursor(django_cursor):
2728 # database cursor. That would require importing more django libraries, and probably
2829 # having to handle some changes in those libraries over different versions.
2930
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+
3038 # We expect that there is always at least one wrapper, but we might as well handle
3139 # the possibility that we get passed the inner cursor.
3240 inner_cursor = django_cursor
@@ -37,11 +45,18 @@ def jsonify_cursor(django_cursor):
3745 # Hopefully we have the right thing now, but try/catch so we can get a little better info
3846 # if it is not. Another option might be an isinstance, or another function that tests the cursor?
3947 try :
40- psycopg2 . extras . register_default_jsonb (conn_or_curs = inner_cursor , loads = json . loads )
48+ register_default_jsonb (conn_or_curs = inner_cursor , loads = loads_func )
4149 except TypeError as e :
4250 raise Exception (f'jsonify_cursor: conn_or_curs was actually a { type (inner_cursor )} { e } )
4351
4452
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 )
58+
59+
4560@contextlib .contextmanager
4661def json_cursor (django_database_connection ):
4762 """
@@ -53,3 +68,60 @@ def json_cursor(django_database_connection):
5368 with django_database_connection .cursor () as cursor :
5469 jsonify_cursor (cursor )
5570 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