Merge pull request piccolo-orm#91 from piccolo-orm/array_column_type

Array column type

Author: dantownsend

docs/src/piccolo/schema/column_types.rst

@@ -204,3 +204,20 @@ a subset of the JSON data, and for filtering in a where clause.
     >>>     Booking.data.arrow('name') == '"Alison"'
     >>> ).run_sync()
     [{'id': 1}]
+
+-------------------------------------------------------------------------------
+
+*****
+Array
+*****
+
+Arrays of data can be stored, which can be useful when you want store lots of
+values without using foreign keys.
+
+.. autoclass:: Array
+
+=============================
+Accessing individual elements
+=============================
+
+.. automethod:: Array.__getitem__

piccolo/columns/base.py

@@ -475,6 +475,11 @@ def get_sql_value(self, value: t.Any) -> t.Any:
             output = f"'{value.isoformat().replace('T', '')}'"
         elif isinstance(value, bytes):
             output = f"'{value.hex()}'"
+        elif isinstance(value, list):
+            # Convert to the array syntax.
+            output = (
+                "'{" + ", ".join([self.get_sql_value(i) for i in value]) + "}'"
+            )
         else:
             output = value
 

piccolo/columns/column_types.py

@@ -52,7 +52,7 @@ def get_querystring(
         value: t.Union[str, Varchar, Text],
         engine_type: str,
         reverse=False,
-    ):
+    ) -> QueryString:
         Concat = ConcatPostgres if engine_type == "postgres" else ConcatSQLite
 
         if isinstance(value, (Varchar, Text)):
@@ -106,7 +106,7 @@ def get_querystring(
         operator: str,
         value: t.Union[int, float, Integer],
         reverse=False,
-    ):
+    ) -> QueryString:
         if isinstance(value, Integer):
             column: Integer = value
             if len(column._meta.call_chain) > 0:
@@ -1268,8 +1268,9 @@ 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
+        instance = t.cast(JSONB, self.copy())
+        instance.json_operator = f"-> '{key}'"
+        return instance
 
     def get_select_string(self, engine_type: str, just_alias=False) -> str:
         select_string = self._meta.get_full_name(just_alias=just_alias)
@@ -1343,3 +1344,95 @@ class Blob(Bytea):
     """
 
     pass
+
+
+###############################################################################
+
+
+class Array(Column):
+    """
+    Used for storing lists of data.
+
+    **Example**
+
+    .. code-block:: python
+
+        class Ticket(Table):
+            seat_numbers = Array(base_column=Integer())
+
+        # Create
+        >>> Ticket(seat_numbers=[34, 35, 36]).save().run_sync()
+
+        # Query
+        >>> Ticket.select(Ticket.seat_numbers).run_sync()
+        {'seat_numbers': [34, 35, 36]}
+
+    """
+
+    value_type = list
+
+    def __init__(
+        self,
+        base_column: Column,
+        default: t.Union[t.List, t.Callable[[], t.List], None] = list,
+        **kwargs,
+    ) -> None:
+        if isinstance(base_column, ForeignKey):
+            raise ValueError("Arrays of ForeignKeys aren't allowed.")
+
+        self._validate_default(default, (list, None))
+
+        self.base_column = base_column
+        self.default = default
+        self.index: t.Optional[int] = None
+        kwargs.update({"base_column": base_column, "default": default})
+        super().__init__(**kwargs)
+
+    @property
+    def column_type(self):
+        engine_type = self._meta.table._meta.db.engine_type
+        if engine_type == "postgres":
+            return f"{self.base_column.column_type}[]"
+        elif engine_type == "sqlite":
+            return "ARRAY"
+        raise Exception("Unrecognized engine type")
+
+    def __getitem__(self, value: int) -> Array:
+        """
+        Allows queries which retrieve an item from the array. The index starts
+        with 0 for the first value. If you were to write the SQL by hand, the
+        first index would be 1 instead:
+
+        https://www.postgresql.org/docs/current/arrays.html
+
+        However, we keep the first index as 0 to fit better with Python.
+
+        For example:
+
+        .. code-block:: python
+
+            >>> Ticket.select(Ticket.seat_numbers[0]).first().run_sync
+            {'seat_numbers': 325}
+
+
+        """
+        if isinstance(value, int):
+            if value < 0:
+                raise ValueError("Only positive integers are allowed.")
+
+            instance = t.cast(Array, self.copy())
+
+            # We deliberately add 1, as Postgres treats the first array element
+            # as index 1.
+            instance.index = value + 1
+            return instance
+        else:
+            raise ValueError("Only integers can be used for indexing.")
+
+    def get_select_string(self, engine_type: str, just_alias=False) -> str:
+        select_string = self._meta.get_full_name(just_alias=just_alias)
+
+        if isinstance(self.index, int):
+            return f"{select_string}[{self.index}]"
+        else:
+            return select_string

piccolo/engine/sqlite.py

@@ -14,6 +14,8 @@
 from piccolo.engine.exceptions import TransactionError
 from piccolo.query.base import Query
 from piccolo.querystring import QueryString
+
+from piccolo.utils.encoding import dump_json, load_json
 from piccolo.utils.sync import run_sync
 
 ###############################################################################
@@ -39,14 +41,14 @@ def convert_uuid_in(value) -> str:
     return str(value)
 
 
-def convert_time_in(value) -> str:
+def convert_time_in(value: datetime.time) -> str:
     """
     Converts the time value being passed into sqlite.
     """
     return value.isoformat()
 
 
-def convert_date_in(value):
+def convert_date_in(value: datetime.date):
     """
     Converts the date value being passed into sqlite.
     """
@@ -64,13 +66,24 @@ def convert_datetime_in(value: datetime.datetime) -> str:
     return str(value)
 
 
-def convert_timedelta_in(value):
+def convert_timedelta_in(value: datetime.timedelta):
     """
     Converts the timedelta value being passed into sqlite.
     """
     return value.total_seconds()
 
 
+def convert_array_in(value: list):
+    """
+    Converts a list value into a string.
+    """
+    if len(value) > 0:
+        if type(value[0]) not in [str, int, float]:
+            raise ValueError("Can only serialise str, int and float.")
+
+    return dump_json(value)
+
+
 # Out
 
 
@@ -129,6 +142,14 @@ def convert_timestamptz_out(value: bytes) -> datetime.datetime:
     return datetime.datetime.fromisoformat(value.decode("utf8"))
 
 
+def convert_array_out(value: bytes) -> t.List:
+    """
+    If the value if from an array column, deserialise the string back into a
+    list.
+    """
+    return load_json(value.decode("utf8"))
+
+
 sqlite3.register_converter("Numeric", convert_numeric_out)
 sqlite3.register_converter("Integer", convert_int_out)
 sqlite3.register_converter("UUID", convert_uuid_out)
@@ -137,13 +158,15 @@ def convert_timestamptz_out(value: bytes) -> datetime.datetime:
 sqlite3.register_converter("Seconds", convert_seconds_out)
 sqlite3.register_converter("Boolean", convert_boolean_out)
 sqlite3.register_converter("Timestamptz", convert_timestamptz_out)
+sqlite3.register_converter("Array", convert_array_out)
 
 sqlite3.register_adapter(Decimal, convert_numeric_in)
 sqlite3.register_adapter(uuid.UUID, convert_uuid_in)
 sqlite3.register_adapter(datetime.time, convert_time_in)
 sqlite3.register_adapter(datetime.date, convert_date_in)
 sqlite3.register_adapter(datetime.datetime, convert_datetime_in)
 sqlite3.register_adapter(datetime.timedelta, convert_timedelta_in)
+sqlite3.register_adapter(list, convert_array_in)
 
 ###############################################################################
 

piccolo/utils/encoding.py

@@ -16,3 +16,10 @@ def dump_json(data: t.Any) -> str:
         return orjson.dumps(data, default=str).decode("utf8")
     else:
         return json.dumps(data, default=str)
+
+
+def load_json(data: str) -> t.Any:
+    if ORJSON:
+        return orjson.loads(data)
+    else:
+        return json.loads(data)

tests/columns/test_array.py

@@ -0,0 +1,41 @@
+from unittest import TestCase
+
+from piccolo.table import Table
+from piccolo.columns.column_types import Array, Integer
+from tests.base import postgres_only
+
+
+class MyTable(Table):
+    value = Array(base_column=Integer())
+
+
+class TestArrayPostgres(TestCase):
+    """
+    Make sure an Array column can be created.
+    """
+
+    def setUp(self):
+        MyTable.create_table().run_sync()
+
+    def tearDown(self):
+        MyTable.alter().drop_table().run_sync()
+
+    def test_storage(self):
+        """
+        Make sure data can be stored and retrieved.
+        """
+        MyTable(value=[1, 2, 3]).save().run_sync()
+
+        row = MyTable.objects().first().run_sync()
+        self.assertEqual(row.value, [1, 2, 3])
+
+    @postgres_only
+    def test_index(self):
+        """
+        Indexes should allow individual array elements to be queried.
+        """
+        MyTable(value=[1, 2, 3]).save().run_sync()
+
+        self.assertEqual(
+            MyTable.select(MyTable.value[0]).first().run_sync(), {"value": 1}
+        )

tests/utils/test_encoding.py

@@ -0,0 +1,12 @@
+from unittest import TestCase
+
+from piccolo.utils.encoding import dump_json, load_json
+
+
+class TestEncodingDecoding(TestCase):
+    def test_dump_load(self):
+        """
+        Test dumping then loading an object.
+        """
+        payload = {"a": [1, 2, 3]}
+        self.assertEqual(load_json(dump_json(payload)), payload)