Skip to content

Commit 85ead5a

Browse files
dantownsenddantownsend
authored
1103 Improve arrow function (#1112)
* improve arrow function * fix tests * refactor, so we have an `Arrow` function * add test for nested arrow functions * skip sqlite * allow `arrow` to access keys multiple levels deep * improve the example data in the playground for JSON data * move arrow function to JSON, as it can be used by JSON or JSONB * add `arrow` function to Arrow, so it can be called recursively * change heading levels of JSON docs * move `Arrow` to operators folder * update docs * improve docstring * add `technicians` to example JSON * improve docstrings * allow `QueryString` as an arg type to `Arrow` * fix docstring error * make sure integers can be passed in * add `QueryString` as an arg type to `arrow` method * added `GetElementFromPath` * add docs for ``from_path`` * add `__getitem__` as a shortcut for the arrow method * update the docs to use the square bracket notation * explain why the method is called `arrow` * move arrow tests into separate class * add `test_multiple_levels_deep` * add tests for `for_path` * last documentation tweaks * add basic operator tests
1 parent 265b0c2 commit 85ead5a

File tree

10 files changed

+482
-106
lines changed

10 files changed

+482
-106
lines changed

‎docs/src/piccolo/schema/column_types.rst‎

Lines changed: 107 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -189,18 +189,15 @@ Storing JSON can be useful in certain situations, for example - raw API
189189
responses, data from a Javascript app, and for storing data with an unknown or
190190
changing schema.
191191

192-
====
193-
JSON
194-
====
192+
====================
193+
``JSON`` / ``JSONB``
194+
====================
195195

196196
.. autoclass:: JSON
197197

198-
=====
199-
JSONB
200-
=====
201-
202198
.. autoclass:: JSONB
203199

200+
===========
204201
Serialising
205202
===========
206203

@@ -224,6 +221,7 @@ You can also pass in a JSON string if you prefer:
224221
)
225222
await studio.save()
226223
224+
=============
227225
Deserialising
228226
=============
229227

@@ -257,29 +255,122 @@ With ``objects`` queries, we can modify the returned JSON, and then save it:
257255
studio['facilities']['restaurant'] = False
258256
await studio.save()
259257
260-
arrow
261-
=====
258+
================
259+
Getting elements
260+
================
261+
262+
``JSON`` and ``JSONB`` columns have an ``arrow`` method (representing the
263+
``->`` operator in Postgres), which is useful for retrieving a child element
264+
from the JSON data.
265+
266+
.. note:: Postgres and CockroachDB only.
262267

263-
``JSONB`` columns have an ``arrow`` function, which is useful for retrieving
264-
a subset of the JSON data:
268+
``select`` queries
269+
==================
270+
271+
If we have the following JSON stored in the ``RecordingStudio.facilities``
272+
column:
273+
274+
.. code-block:: json
275+
276+
{
277+
"instruments": {
278+
"drum_kits": 2,
279+
"electric_guitars": 10
280+
},
281+
"restaurant": true,
282+
"technicians": [
283+
{
284+
"name": "Alice Jones"
285+
},
286+
{
287+
"name": "Bob Williams"
288+
}
289+
]
290+
}
291+
292+
We can retrieve the ``restaurant`` value from the JSON object:
265293

266294
.. code-block:: python
267295
268296
>>> await RecordingStudio.select(
269-
... RecordingStudio.name,
270-
... RecordingStudio.facilities.arrow('mixing_desk').as_alias('mixing_desk')
297+
... RecordingStudio.facilities.arrow('restaurant')
298+
... .as_alias('restaurant')
271299
... ).output(load_json=True)
272-
[{'name': 'Abbey Road', 'mixing_desk': True}]
300+
[{'restaurant': True}, ...]
273301
274-
It can also be used for filtering in a where clause:
302+
As a convenience, you can use square brackets, instead of calling ``arrow``
303+
explicitly:
304+
305+
.. code-block:: python
306+
307+
>>> await RecordingStudio.select(
308+
... RecordingStudio.facilities['restaurant']
309+
... .as_alias('restaurant')
310+
... ).output(load_json=True)
311+
[{'restaurant': True}, ...]
312+
313+
You can drill multiple levels deep by calling ``arrow`` multiple times (or
314+
alternatively use the :ref:`from_path` method - see below).
315+
316+
Here we fetch the number of drum kits that the recording studio has:
317+
318+
.. code-block:: python
319+
320+
>>> await RecordingStudio.select(
321+
... RecordingStudio.facilities["instruments"]["drum_kits"]
322+
... .as_alias("drum_kits")
323+
... ).output(load_json=True)
324+
[{'drum_kits': 2}, ...]
325+
326+
If you have a JSON object which consists of arrays and objects, then you can
327+
navigate the array elements by passing in an integer to ``arrow``.
328+
329+
Here we fetch the first technician from the array:
330+
331+
.. code-block:: python
332+
333+
>>> await RecordingStudio.select(
334+
... RecordingStudio.facilities["technicians"][0]["name"]
335+
... .as_alias("technician_name")
336+
... ).output(load_json=True)
337+
338+
[{'technician_name': 'Alice Jones'}, ...]
339+
340+
``where`` clauses
341+
=================
342+
343+
The ``arrow`` operator can also be used for filtering in a where clause:
275344

276345
.. code-block:: python
277346
278347
>>> await RecordingStudio.select(RecordingStudio.name).where(
279-
... RecordingStudio.facilities.arrow('mixing_desk') == True
348+
... RecordingStudio.facilities['mixing_desk'].eq(True)
280349
... )
281350
[{'name': 'Abbey Road'}]
282351
352+
.. _from_path:
353+
354+
=============
355+
``from_path``
356+
=============
357+
358+
This works the same as ``arrow`` but is more optimised if you need to return
359+
part of a highly nested JSON structure.
360+
361+
.. code-block:: python
362+
363+
>>> await RecordingStudio.select(
364+
... RecordingStudio.facilities.from_path([
365+
... "technicians",
366+
... 0,
367+
... "name"
368+
... ]).as_alias("technician_name")
369+
... ).output(load_json=True)
370+
371+
[{'technician_name': 'Alice Jones'}, ...]
372+
373+
=============
283374
Handling null
284375
=============
285376

‎piccolo/apps/playground/commands/run.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,11 @@ def populate():
233233
RecordingStudio.facilities: {
234234
"restaurant": True,
235235
"mixing_desk": True,
236+
"instruments": {"electric_guitars": 10, "drum_kits": 2},
237+
"technicians": [
238+
{"name": "Alice Jones"},
239+
{"name": "Bob Williams"},
240+
],
236241
},
237242
}
238243
)
@@ -244,6 +249,10 @@ def populate():
244249
RecordingStudio.facilities: {
245250
"restaurant": False,
246251
"mixing_desk": True,
252+
"instruments": {"electric_guitars": 6, "drum_kits": 3},
253+
"technicians": [
254+
{"name": "Frank Smith"},
255+
],
247256
},
248257
},
249258
)

‎piccolo/columns/column_types.py‎

Lines changed: 78 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,10 @@ class Band(Table):
7070

7171
if t.TYPE_CHECKING: # pragma: no cover
7272
from piccolo.columns.base import ColumnMeta
73+
from piccolo.query.operators.json import (
74+
GetChildElement,
75+
GetElementFromPath,
76+
)
7377
from piccolo.table import Table
7478

7579

@@ -2319,6 +2323,76 @@ def column_type(self):
23192323
else:
23202324
return "JSON"
23212325

2326+
###########################################################################
2327+
2328+
def arrow(self, key: t.Union[str, int, QueryString]) -> GetChildElement:
2329+
"""
2330+
Allows a child element of the JSON structure to be returned - for
2331+
example::
2332+
2333+
>>> await RecordingStudio.select(
2334+
... RecordingStudio.facilities.arrow("restaurant")
2335+
... )
2336+
2337+
"""
2338+
from piccolo.query.operators.json import GetChildElement
2339+
2340+
alias = self._alias or self._meta.get_default_alias()
2341+
return GetChildElement(identifier=self, key=key, alias=alias)
2342+
2343+
def __getitem__(
2344+
self, value: t.Union[str, int, QueryString]
2345+
) -> GetChildElement:
2346+
"""
2347+
A shortcut for the ``arrow`` method, used for retrieving a child
2348+
element.
2349+
2350+
For example:
2351+
2352+
.. code-block:: python
2353+
2354+
>>> await RecordingStudio.select(
2355+
... RecordingStudio.facilities["restaurant"]
2356+
... )
2357+
2358+
"""
2359+
return self.arrow(key=value)
2360+
2361+
def from_path(
2362+
self,
2363+
path: t.List[t.Union[str, int]],
2364+
) -> GetElementFromPath:
2365+
"""
2366+
Allows an element of the JSON structure to be returned, which can be
2367+
arbitrarily deep. For example::
2368+
2369+
>>> await RecordingStudio.select(
2370+
... RecordingStudio.facilities.from_path([
2371+
... "technician",
2372+
... 0,
2373+
... "first_name"
2374+
... ])
2375+
... )
2376+
2377+
It's the same as calling ``arrow`` multiple times, but is more
2378+
efficient / convenient if extracting highly nested data::
2379+
2380+
>>> await RecordingStudio.select(
2381+
... RecordingStudio.facilities.arrow(
2382+
... "technician"
2383+
... ).arrow(
2384+
... 0
2385+
... ).arrow(
2386+
... "first_name"
2387+
... )
2388+
... )
2389+
2390+
"""
2391+
from piccolo.query.operators.json import GetElementFromPath
2392+
2393+
alias = self._alias or self._meta.get_default_alias()
2394+
return GetElementFromPath(identifier=self, path=path, alias=alias)
2395+
23222396
###########################################################################
23232397
# Descriptors
23242398

@@ -2337,10 +2411,10 @@ def __set__(self, obj, value: t.Union[str, t.Dict]):
23372411

23382412
class JSONB(JSON):
23392413
"""
2340-
Used for storing JSON strings - Postgres only. The data is stored in a
2341-
binary format, and can be queried. Insertion can be slower (as it needs to
2342-
be converted to the binary format). The benefits of JSONB generally
2343-
outweigh the downsides.
2414+
Used for storing JSON strings - Postgres / CochroachDB only. The data is
2415+
stored in a binary format, and can be queried more efficiently. Insertion
2416+
can be slower (as it needs to be converted to the binary format). The
2417+
benefits of JSONB generally outweigh the downsides.
23442418
23452419
:param default:
23462420
Either a JSON string can be provided, or a Python ``dict`` or ``list``
@@ -2352,41 +2426,6 @@ class JSONB(JSON):
23522426
def column_type(self):
23532427
return "JSONB" # Must be defined, we override column_type() in JSON()
23542428

2355-
def arrow(self, key: str) -> JSONB:
2356-
"""
2357-
Allows part of the JSON structure to be returned - for example,
2358-
for {"a": 1}, and a key value of "a", then 1 will be returned.
2359-
"""
2360-
instance = t.cast(JSONB, self.copy())
2361-
instance.json_operator = f"-> '{key}'"
2362-
return instance
2363-
2364-
def get_select_string(
2365-
self, engine_type: str, with_alias: bool = True
2366-
) -> QueryString:
2367-
select_string = self._meta.get_full_name(with_alias=False)
2368-
2369-
if self.json_operator is not None:
2370-
select_string += f" {self.json_operator}"
2371-
2372-
if with_alias:
2373-
alias = self._alias or self._meta.get_default_alias()
2374-
select_string += f' AS "{alias}"'
2375-
2376-
return QueryString(select_string)
2377-
2378-
def eq(self, value) -> Where:
2379-
"""
2380-
See ``Boolean.eq`` for more details.
2381-
"""
2382-
return self.__eq__(value)
2383-
2384-
def ne(self, value) -> Where:
2385-
"""
2386-
See ``Boolean.ne`` for more details.
2387-
"""
2388-
return self.__ne__(value)
2389-
23902429
###########################################################################
23912430
# Descriptors
23922431

‎piccolo/query/base.py‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from piccolo.columns.column_types import JSON, JSONB
77
from piccolo.custom_types import QueryResponseType, TableInstance
88
from piccolo.query.mixins import ColumnsDelegate
9+
from piccolo.query.operators.json import JSONQueryString
910
from piccolo.querystring import QueryString
1011
from piccolo.utils.encoding import load_json
1112
from piccolo.utils.objects import make_nested_object
@@ -65,16 +66,20 @@ async def _process_results(self, results) -> QueryResponseType:
6566
self, "columns_delegate", None
6667
)
6768

69+
json_column_names: t.List[str] = []
70+
6871
if columns_delegate is not None:
69-
json_columns = [
70-
i
71-
for i in columns_delegate.selected_columns
72-
if isinstance(i, (JSON, JSONB))
73-
]
72+
json_columns: t.List[t.Union[JSON, JSONB]] = []
73+
74+
for column in columns_delegate.selected_columns:
75+
if isinstance(column, (JSON, JSONB)):
76+
json_columns.append(column)
77+
elif isinstance(column, JSONQueryString):
78+
if alias := column._alias:
79+
json_column_names.append(alias)
7480
else:
7581
json_columns = self.table._meta.json_columns
7682

77-
json_column_names = []
7883
for column in json_columns:
7984
if column._alias is not None:
8085
json_column_names.append(column._alias)

‎piccolo/query/operators/__init__.py‎

Whitespace-only changes.

0 commit comments

Comments
 (0)