Merge pull request piccolo-orm#21 from piccolo-orm/json_field

JSON field

Author: dantownsend

docs/src/piccolo/query_types/select.rst

@@ -31,6 +31,18 @@ Or making an alias to make it shorter:
 
 .. hint:: All of these examples also work with async by using .run() inside coroutines - see :ref:`SyncAndAsync`.
 
+as_alias
+--------
+
+By using ``as_alias``, the name of the row can be overriden in the response.
+
+.. code-block:: python
+
+    >>> Band.select(Band.name.as_alias('title')).run_sync()
+    [{'title': 'Rustaceans'}, {'title': 'Pythonistas'}]
+
+This is equivalent to ``SELECT name AS title FROM band`` in SQL.
+
 Joins
 -----
 

docs/src/piccolo/schema/column_types.rst

@@ -134,3 +134,56 @@ Timestamp
 =========
 
 .. autoclass:: Timestamp
+
+-------------------------------------------------------------------------------
+
+****
+JSON
+****
+
+Storing JSON can be useful in certain situations, for example - raw API
+responses, data from a Javascript app, and for storing data with an unknown or
+changing schema.
+
+====
+JSON
+====
+
+.. autoclass:: JSON
+
+=====
+JSONB
+=====
+
+.. autoclass:: JSONB
+
+arrow
+=====
+
+``JSONB`` columns have an ``arrow`` function, which is useful for retrieving
+a subset of the JSON data, and for filtering in a where clause.
+
+.. code-block:: python
+
+    # Example schema:
+    class Booking(Table):
+        data = JSONB()
+
+    Booking.create_table().run_sync()
+
+    # Example data:
+    Booking.insert(
+        Booking(data='{"name": "Alison"}'),
+        Booking(data='{"name": "Bob"}')
+    ).run_sync()
+
+    # Example queries
+    >>> Booking.select(
+    >>>     Booking.id, Booking.data.arrow('name').as_alias('name')
+    >>> ).run_sync()
+    [{'id': 1, 'name': '"Alison"'}, {'id': 2, 'name': '"Bob"'}]
+
+    >>> Booking.select(Booking.id).where(
+    >>>     Booking.data.arrow('name') == '"Alison"'
+    >>> ).run_sync()
+    [{'id': 1}]

piccolo/columns/__init__.py

@@ -1,10 +1,12 @@
-from .column_types import (  # noqa
+from .column_types import (  # noqa: F401
     Boolean,
     Decimal,
     Float,
     ForeignKey,
     Integer,
     Interval,
+    JSON,
+    JSONB,
     Numeric,
     PrimaryKey,
     Real,
@@ -15,6 +17,6 @@
     UUID,
     Varchar,
 )
-from .base import Column, ForeignKeyMeta, Selectable  # noqa
-from .base import OnDelete, OnUpdate  # noqa
-from .combination import And, Or, Where  # noqa
+from .base import Column, ForeignKeyMeta, Selectable  # noqa: F401
+from .base import OnDelete, OnUpdate  # noqa: F401
+from .combination import And, Or, Where  # noqa: F401

piccolo/columns/base.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 from abc import ABCMeta, abstractmethod
+import copy
 from dataclasses import dataclass, field
 import datetime
 import decimal
@@ -226,6 +227,8 @@ def __init__(
             required=required,
         )
 
+        self.alias: t.Optional[str] = None
+
     def _validate_default(
         self,
         default: t.Any,
@@ -306,6 +309,20 @@ def __ne__(self, value) -> Where:  # type: ignore
     def __hash__(self):
         return hash(self._meta.name)
 
+    def as_alias(self, name: str) -> Column:
+        """
+        Allows column names to be changed in the result of a select.
+
+        For example:
+
+        >>> await Band.select(Band.name.as_alias('title')).run()
+        {'title': 'Pythonistas'}
+
+        """
+        column = copy.deepcopy(self)
+        column.alias = name
+        return column
+
     def get_default_value(self) -> t.Any:
         """
         If the column has a default attribute, return it. If it's callable,
@@ -323,9 +340,16 @@ def get_select_string(self, engine_type: str, just_alias=False) -> str:
         """
         How to refer to this column in a SQL query.
         """
-        return self._meta.get_full_name(just_alias=just_alias)
+        if self.alias is None:
+            return self._meta.get_full_name(just_alias=just_alias)
+        else:
+            original_name = self._meta.get_full_name(just_alias=True)
+            return f"{original_name} AS {self.alias}"
 
-    def get_sql_value(self, value: t.Any) -> str:
+    def get_where_string(self, engine_type: str) -> str:
+        return self.get_select_string(engine_type=engine_type, just_alias=True)
+
+    def get_sql_value(self, value: t.Any) -> t.Any:
         """
         When using DDL statements, we can't parameterise the values. An example
         is when setting the default for a column. So we have to convert from
@@ -390,6 +414,13 @@ def querystring(self) -> QueryString:
         if not self._meta.primary:
             default = self.get_default_value()
             sql_value = self.get_sql_value(value=default)
+            # Escape the value if it contains a pair of curly braces, otherwise
+            # an empty value will appear in the compiled querystring.
+            sql_value = (
+                sql_value.replace("{}", "{{}}")
+                if isinstance(sql_value, str)
+                else sql_value
+            )
             query += f" DEFAULT {sql_value}"
 
         return QueryString(query)

piccolo/columns/column_types.py

@@ -17,6 +17,7 @@
 )
 from piccolo.columns.defaults.uuid import UUIDArg, UUID4
 from piccolo.querystring import Unquoted, QueryString
+from piccolo.utils.encoding import dump_json
 
 if t.TYPE_CHECKING:
     from piccolo.table import Table
@@ -1039,3 +1040,67 @@ def __getattribute__(self, name: str):
             return new_column
         else:
             return value
+
+
+###############################################################################
+
+
+class JSON(Column):  # lgtm[py/missing-equals]
+    """
+    Used for storing JSON strings. The data is stored as text. This can be
+    preferable to JSONB if you just want to store and retrieve JSON without
+    querying it directly. It works with SQLite and Postgres.
+
+    :param default:
+        Either a JSON string can be provided, or a Python ``dict`` or ``list``
+        which is then converted to a JSON string.
+
+    """
+
+    value_type = str
+
+    def __init__(
+        self, default: t.Union[str, t.List, t.Dict, None] = "{}", **kwargs
+    ) -> None:
+        self._validate_default(default, (str, list, dict, None))
+
+        if isinstance(default, (list, dict)):
+            default = dump_json(default)
+
+        self.default = default
+        kwargs.update({"default": default})
+        super().__init__(**kwargs)
+
+        self.json_operator: t.Optional[str] = None
+
+
+class JSONB(JSON):
+    """
+    Used for storing JSON strings - Postgres only. The data is stored in a
+    binary format, and can be queried. Insertion can be slower (as it needs to
+    be converted to the binary format). The benefits of JSONB generally
+    outweigh the downsides.
+
+    :param default:
+        Either a JSON string can be provided, or a Python ``dict`` or ``list``
+        which is then converted to a JSON string.
+
+    """
+
+    def arrow(self, key: str) -> JSONB:
+        """
+        Allows part of the JSON structure to be returned - for example,
+        for {"a": 1}, and a key value of "a", then 1 will be returned.
+        """
+        self.json_operator = f"-> '{key}'"
+        return self
+
+    def get_select_string(self, engine_type: str, just_alias=False) -> str:
+        select_string = self._meta.get_full_name(just_alias=just_alias)
+        if self.json_operator is None:
+            return select_string
+        else:
+            if self.alias is None:
+                return f"{select_string} {self.json_operator}"
+            else:
+                return f"{select_string} {self.json_operator} AS {self.alias}"

piccolo/columns/combination.py

@@ -60,7 +60,7 @@ class Where(CombinableMixin):
 
     def __init__(
         self,
-        column: "Column",
+        column: Column,
         value: t.Any = UNDEFINED,
         values: t.Union[Iterable, Undefined] = UNDEFINED,
         operator: t.Type[ComparisonOperator] = ComparisonOperator,
@@ -91,7 +91,9 @@ def querystring(self) -> QueryString:
             args.append(self.values_querystring)
 
         template = self.operator.template.format(
-            name=self.column._meta.get_full_name(just_alias=True),
+            name=self.column.get_where_string(
+                engine_type=self.column._meta.engine_type
+            ),
             value="{}",
             values="{}",
         )

piccolo/query/base.py

@@ -5,6 +5,7 @@
 
 from piccolo.querystring import QueryString
 from piccolo.utils.sync import run_sync
+from piccolo.utils.encoding import dump_json
 
 if t.TYPE_CHECKING:
     from piccolo.table import Table  # noqa
@@ -68,14 +69,7 @@ async def _process_results(self, results):
                     else:
                         raw = list(itertools.chain(*[j.values() for j in raw]))
                 if output._output.as_json:
-                    try:
-                        import orjson
-                    except ImportError:
-                        import json
-
-                        raw = json.dumps(raw, default=str)
-                    else:
-                        raw = orjson.dumps(raw, default=str).decode("utf8")
+                    raw = dump_json(raw)
 
         return raw
 

piccolo/query/methods/select.py

@@ -248,7 +248,9 @@ def querystrings(self) -> t.Sequence[QueryString]:
 
         #######################################################################
 
-        select = "SELECT DISTINCT" if self.distinct else "SELECT"
+        select = (
+            "SELECT DISTINCT" if self.distinct_delegate._distinct else "SELECT"
+        )
         query = f"{select} {columns_str} FROM {self.table._meta.tablename}"
 
         for join in joins:

piccolo/query/mixins.py

@@ -176,6 +176,14 @@ def output(
         as_list: t.Optional[bool] = None,
         as_json: t.Optional[bool] = None,
     ):
+        """
+        :param as_list:
+            If each row only returns a single value, compile all of the results
+            into a single list.
+        :param as_json:
+            The results are serialised into JSON. It's equivalent to running
+            `json.dumps` on the result.
+        """
         if as_list is not None:
             self._output.as_list = bool(as_list)
 

piccolo/utils/encoding.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+import typing as t
+
+try:
+    import orjson
+
+    ORJSON = True
+except ImportError:
+    import json
+
+    ORJSON = False
+
+
+def dump_json(data: t.Any) -> str:
+    if ORJSON:
+        return orjson.dumps(data, default=str).decode("utf8")
+    else:
+        return json.dumps(data, default=str)