Merge pull request piccolo-orm#92 from piccolo-orm/array_filtering

added support for `any` and `all` for array filtering

Author: dantownsend

docs/src/piccolo/schema/column_types.rst

@@ -221,3 +221,15 @@ Accessing individual elements
 =============================
 
 .. automethod:: Array.__getitem__
+
+===
+any
+===
+
+.. automethod:: Array.any
+
+===
+all
+===
+
+.. automethod:: Array.all

piccolo/columns/column_types.py

@@ -12,6 +12,7 @@
     OnDelete,
     OnUpdate,
 )
+from piccolo.columns.combination import Where
 from piccolo.columns.defaults.date import DateArg, DateCustom, DateNow
 from piccolo.columns.defaults.interval import IntervalArg, IntervalCustom
 from piccolo.columns.defaults.time import TimeArg, TimeCustom, TimeNow
@@ -26,6 +27,7 @@
     TimestamptzNow,
 )
 from piccolo.columns.defaults.uuid import UUID4, UUIDArg
+from piccolo.columns.operators.comparison import ArrayAll, ArrayAny
 from piccolo.columns.operators.string import ConcatPostgres, ConcatSQLite
 from piccolo.columns.reference import LazyTableReference
 from piccolo.querystring import QueryString, Unquoted
@@ -1416,6 +1418,12 @@ def __getitem__(self, value: int) -> Array:
 
 
         """
+        engine_type = self._meta.table._meta.db.engine_type
+        if engine_type != "postgres":
+            raise ValueError(
+                "Only Postgres supports array indexing currently."
+            )
+
         if isinstance(value, int):
             if value < 0:
                 raise ValueError("Only positive integers are allowed.")
@@ -1436,3 +1444,39 @@ def get_select_string(self, engine_type: str, just_alias=False) -> str:
             return f"{select_string}[{self.index}]"
         else:
             return select_string
+
+    def any(self, value: t.Any) -> Where:
+        """
+        Check if any of the items in the array match the given value.
+
+        .. code-block:: python
+
+            >>> Ticket.select().where(Ticket.seat_numbers.any(510)).run_sync()
+
+        """
+        engine_type = self._meta.table._meta.db.engine_type
+
+        if engine_type == "postgres":
+            return Where(column=self, value=value, operator=ArrayAny)
+        elif engine_type == "sqlite":
+            return self.like(f"%{value}%")
+        else:
+            raise ValueError("Unrecognised engine type")
+
+    def all(self, value: t.Any) -> Where:
+        """
+        Check if all of the items in the array match the given value.
+
+        .. code-block:: python
+
+            >>> Ticket.select().where(Ticket.seat_numbers.all(510)).run_sync()
+
+        """
+        engine_type = self._meta.table._meta.db.engine_type
+
+        if engine_type == "postgres":
+            return Where(column=self, value=value, operator=ArrayAll)
+        elif engine_type == "sqlite":
+            raise ValueError("Unsupported by SQLite")
+        else:
+            raise ValueError("Unrecognised engine type")

piccolo/columns/operators/comparison.py

@@ -56,3 +56,11 @@ class LessThan(ComparisonOperator):
 
 class LessEqualThan(ComparisonOperator):
     template = "{name} <= {value}"
+
+
+class ArrayAny(ComparisonOperator):
+    template = "{value} = ANY ({name})"
+
+
+class ArrayAll(ComparisonOperator):
+    template = "{value} = ALL ({name})"

tests/columns/test_array.py

@@ -39,3 +39,50 @@ def test_index(self):
         self.assertEqual(
             MyTable.select(MyTable.value[0]).first().run_sync(), {"value": 1}
         )
+
+    @postgres_only
+    def test_all(self):
+        """
+        Make sure rows can be retrieved where all items in an array match a
+        given value.
+        """
+        MyTable(value=[1, 1, 1]).save().run_sync()
+
+        self.assertEqual(
+            MyTable.select(MyTable.value)
+            .where(MyTable.value.all(1))
+            .first()
+            .run_sync(),
+            {"value": [1, 1, 1]},
+        )
+
+        self.assertEqual(
+            MyTable.select(MyTable.value)
+            .where(MyTable.value.all(0))
+            .first()
+            .run_sync(),
+            None,
+        )
+
+    def test_any(self):
+        """
+        Make sure rows can be retrieved where any items in an array match a
+        given value.
+        """
+        MyTable(value=[1, 2, 3]).save().run_sync()
+
+        self.assertEqual(
+            MyTable.select(MyTable.value)
+            .where(MyTable.value.any(1))
+            .first()
+            .run_sync(),
+            {"value": [1, 2, 3]},
+        )
+
+        self.assertEqual(
+            MyTable.select(MyTable.value)
+            .where(MyTable.value.any(0))
+            .first()
+            .run_sync(),
+            None,
+        )