Allow a subquery to be passed into is_in and not_in (#1217)

sinisaos · dantownsend · web-flow · commit 15d18af0a626 · 2025-07-15T22:41:57.000+01:00
* adds a subquery for is_in and not_in for tables with fk relationships

* pass `QueryString` in, instead of `Select`

* allow querystring to be passed in to - cuts down on repetition

* add more tests

* tweak docs

* update docs

---------

Co-authored-by: Daniel Townsend &lt;dan@dantownsend.co.uk&gt;
diff --git a/docs/src/piccolo/query_clauses/where.rst b/docs/src/piccolo/query_clauses/where.rst
@@ -112,6 +112,45 @@ And all rows with a value not contained in the list:
         Band.name.not_in(['Terrible Band', 'Awful Band'])
     )
 
+You can also pass a subquery into the ``is_in`` clause:
+
+.. code-block:: python
+
+    await Band.select().where(
+        Band.id.is_in(
+            Concert.select(Concert.band_1).where(
+                Concert.starts >= datetime.datetime(year=2025, month=1, day=1)
+            )
+        )
+    )
+
+.. hint::
+    In SQL there are often several ways of solving the same problem. You
+    can also solve the above using :meth:`join_on <piccolo.columns.base.Column.join_on>`.
+
+    .. code-block:: python
+
+        >>> await Band.select().where(
+        ...     Band.id.join_on(Concert.band_1).starts >= datetime.datetime(
+        ...        year=2025, month=1, day=1
+        ...     )
+        ... )
+
+    Use whichever you prefer, and whichever suits the situation best.
+
+Subqueries can also be passed into the ``not_in`` clause:
+
+.. code-block:: python
+
+    await Band.select().where(
+        Band.id.not_in(
+            Concert.select(Concert.band_1).where(
+                Concert.starts >= datetime.datetime(year=2025, month=1, day=1)
+            )
+        )
+    )
+
+
 -------------------------------------------------------------------------------
 
 ``is_null`` / ``is_not_null``
diff --git a/piccolo/columns/base.py b/piccolo/columns/base.py
@@ -46,6 +46,7 @@
 
 if TYPE_CHECKING:  # pragma: no cover
     from piccolo.columns.column_types import ForeignKey
+    from piccolo.query.methods.select import Select
     from piccolo.table import Table
 
 
@@ -599,18 +600,40 @@ def _validate_choices(
 
         return True
 
-    def is_in(self, values: list[Any]) -> Where:
-        if len(values) == 0:
-            raise ValueError(
-                "The `values` list argument must contain at least one value."
-            )
+    def is_in(self, values: Union[Select, QueryString, list[Any]]) -> Where:
+        from piccolo.query.methods.select import Select
+
+        if isinstance(values, list):
+            if len(values) == 0:
+                raise ValueError(
+                    "The `values` list argument must contain at least one "
+                    "value."
+                )
+        elif isinstance(values, Select):
+            if len(values.columns_delegate.selected_columns) != 1:
+                raise ValueError(
+                    "A sub select must only return a single column."
+                )
+            values = values.querystrings[0]
+
         return Where(column=self, values=values, operator=In)
 
-    def not_in(self, values: list[Any]) -> Where:
-        if len(values) == 0:
-            raise ValueError(
-                "The `values` list argument must contain at least one value."
-            )
+    def not_in(self, values: Union[Select, QueryString, list[Any]]) -> Where:
+        from piccolo.query.methods.select import Select
+
+        if isinstance(values, list):
+            if len(values) == 0:
+                raise ValueError(
+                    "The `values` list argument must contain at least one "
+                    "value."
+                )
+        elif isinstance(values, Select):
+            if len(values.columns_delegate.selected_columns) != 1:
+                raise ValueError(
+                    "A sub select must only return a single column."
+                )
+            values = values.querystrings[0]
+
         return Where(column=self, values=values, operator=NotIn)
 
     def like(self, value: str) -> Where:
diff --git a/piccolo/columns/combination.py b/piccolo/columns/combination.py
@@ -146,7 +146,7 @@ def __init__(
         self,
         column: Column,
         value: Any = UNDEFINED,
-        values: Union[CustomIterable, Undefined] = UNDEFINED,
+        values: Union[CustomIterable, Undefined, QueryString] = UNDEFINED,
         operator: type[ComparisonOperator] = ComparisonOperator,
     ) -> None:
         """
@@ -156,7 +156,7 @@ def __init__(
         self.column = column
 
         self.value = value if value == UNDEFINED else self.clean_value(value)
-        if values == UNDEFINED:
+        if (values == UNDEFINED) or isinstance(values, QueryString):
             self.values = values
         else:
             self.values = [self.clean_value(i) for i in values]  # type: ignore
@@ -192,6 +192,9 @@ def clean_value(self, value: Any) -> Any:
     def values_querystring(self) -> QueryString:
         values = self.values
 
+        if isinstance(values, QueryString):
+            return values
+
         if isinstance(values, Undefined):
             raise ValueError("values is undefined")
 
diff --git a/tests/columns/test_combination.py b/tests/columns/test_combination.py
@@ -1,6 +1,6 @@
 import unittest
 
-from tests.example_apps.music.tables import Band
+from tests.example_apps.music.tables import Band, Concert
 
 
 class TestWhere(unittest.TestCase):
@@ -29,6 +29,20 @@ def test_is_in(self):
         with self.assertRaises(ValueError):
             Band.name.is_in([])
 
+    def test_is_in_subquery(self):
+        _where = Band.id.is_in(
+            Concert.select(Concert.band_1).where(Concert.band_1 == 1)
+        )
+        sql = _where.__str__()
+        self.assertEqual(
+            sql,
+            '"band"."id" IN (SELECT ALL "concert"."band_1" AS "band_1" FROM "concert" WHERE "concert"."band_1" = 1)',  # noqa: E501
+        )
+
+        # a sub select must only return a single column
+        with self.assertRaises(ValueError):
+            Band.id.is_in(Concert.select().where(Concert.band_1 == 1))
+
     def test_not_in(self):
         _where = Band.name.not_in(["CSharps"])
         sql = _where.__str__()
@@ -37,6 +51,20 @@ def test_not_in(self):
         with self.assertRaises(ValueError):
             Band.name.not_in([])
 
+    def test_not_in_subquery(self):
+        _where = Band.id.not_in(
+            Concert.select(Concert.band_1).where(Concert.band_1 == 1)
+        )
+        sql = _where.__str__()
+        self.assertEqual(
+            sql,
+            '"band"."id" NOT IN (SELECT ALL "concert"."band_1" AS "band_1" FROM "concert" WHERE "concert"."band_1" = 1)',  # noqa: E501
+        )
+
+        # a sub select must only return a single column
+        with self.assertRaises(ValueError):
+            Band.id.not_in(Concert.select().where(Concert.band_1 == 1))
+
 
 class TestAnd(unittest.TestCase):
     def test_get_column_values(self):
diff --git a/tests/table/test_select.py b/tests/table/test_select.py
@@ -258,6 +258,64 @@ def test_where_greater_than(self):
 
         self.assertEqual(response, [{"name": "Rustaceans"}])
 
+    def test_is_in(self):
+        self.insert_rows()
+
+        response = (
+            Band.select(Band.name)
+            .where(Band.manager._.name.is_in(["Guido"]))
+            .run_sync()
+        )
+
+        self.assertListEqual(response, [{"name": "Pythonistas"}])
+
+    def test_is_in_subquery(self):
+        self.insert_rows()
+
+        # This is a contrived example, just for testing.
+        response = (
+            Band.select(Band.name)
+            .where(
+                Band.manager.is_in(
+                    Manager.select(Manager.id).where(Manager.name == "Guido")
+                )
+            )
+            .run_sync()
+        )
+
+        self.assertListEqual(response, [{"name": "Pythonistas"}])
+
+    def test_not_in(self):
+        self.insert_rows()
+
+        response = (
+            Band.select(Band.name)
+            .where(Band.manager._.name.not_in(["Guido"]))
+            .run_sync()
+        )
+
+        self.assertListEqual(
+            response, [{"name": "Rustaceans"}, {"name": "CSharps"}]
+        )
+
+    def test_not_in_subquery(self):
+        self.insert_rows()
+
+        # This is a contrived example, just for testing.
+        response = (
+            Band.select(Band.name)
+            .where(
+                Band.manager.not_in(
+                    Manager.select(Manager.id).where(Manager.name == "Guido")
+                )
+            )
+            .run_sync()
+        )
+
+        self.assertListEqual(
+            response, [{"name": "Rustaceans"}, {"name": "CSharps"}]
+        )
+
     def test_where_is_null(self):
         self.insert_rows()
 
