minor tweak to M2M so it flattens all_columns (#730)

piccolo/columns/m2m.py

@@ -12,6 +12,7 @@
     ForeignKey,
     LazyTableReference,
 )
+from piccolo.utils.list import flatten
 from piccolo.utils.sync import run_sync
 
 if t.TYPE_CHECKING:
@@ -45,9 +46,9 @@ def __init__(
         self.m2m = m2m
         self.load_json = load_json
 
-        safe_types = [int, str]
+        safe_types = (int, str)
 
-        # If the columns can be serialised / deserialise as JSON, then we
+        # If the columns can be serialised / deserialised as JSON, then we
         # can fetch the data all in one go.
         self.serialisation_safe = all(
             (column.__class__.value_type in safe_types)
@@ -408,7 +409,10 @@ def __init__(
         )
 
     def __call__(
-        self, *columns: Column, as_list: bool = False, load_json: bool = False
+        self,
+        *columns: t.Union[Column, t.List[Column]],
+        as_list: bool = False,
+        load_json: bool = False,
     ) -> M2MSelect:
         """
         :param columns:
@@ -420,14 +424,16 @@ def __call__(
         :param load_json:
             If ``True``, any JSON strings are loaded as Python objects.
         """
-        if not columns:
-            columns = tuple(self._meta.secondary_table._meta.columns)
+        columns_ = flatten(columns)
+
+        if not columns_:
+            columns_ = self._meta.secondary_table._meta.columns
 
-        if as_list and len(columns) != 1:
+        if as_list and len(columns_) != 1:
             raise ValueError(
                 "`as_list` is only valid with a single column argument"
             )
 
         return M2MSelect(
-            *columns, m2m=self, as_list=as_list, load_json=load_json
+            *columns_, m2m=self, as_list=as_list, load_json=load_json
         )

piccolo/query/mixins.py

@@ -10,6 +10,7 @@
 from piccolo.columns.column_types import ForeignKey
 from piccolo.custom_types import Combinable
 from piccolo.querystring import QueryString
+from piccolo.utils.list import flatten
 from piccolo.utils.sql_values import convert_to_sql_value
 
 if t.TYPE_CHECKING:  # pragma: no cover
@@ -440,13 +441,7 @@ def columns(self, *columns: t.Union[Selectable, t.List[Selectable]]):
             in which case we flatten the list.
 
         """
-        _columns = []
-        for column in columns:
-            if isinstance(column, list):
-                _columns.extend(column)
-            else:
-                _columns.append(column)
-
+        _columns = flatten(columns)
         combined = list(self.selected_columns) + _columns
         self.selected_columns = combined
 

piccolo/utils/list.py

@@ -0,0 +1,27 @@
+import typing as t
+
+ElementType = t.TypeVar("ElementType")
+
+
+def flatten(
+    items: t.Sequence[t.Union[ElementType, t.List[ElementType]]]
+) -> t.List[ElementType]:
+    """
+    Takes a sequence of elements, and flattens it out. For example::
+
+        >>> flatten(['a', ['b', 'c']])
+        ['a', 'b', 'c']
+
+    We need this for situations like this::
+
+        await Band.select(Band.name, Band.manager.all_columns())
+
+    """
+    _items: t.List[ElementType] = []
+    for item in items:
+        if isinstance(item, list):
+            _items.extend(item)
+        else:
+            _items.append(item)
+
+    return _items

tests/columns/test_m2m.py

@@ -273,6 +273,41 @@ def test_select_id(self):
             ],
         )
 
+    @engines_skip("cockroach")
+    def test_select_all_columns(self):
+        """
+        Make sure ``all_columns`` can be passed in as an argument. ``M2M``
+        should flatten the arguments. Reported here:
+
+        https://github.com/piccolo-orm/piccolo/issues/728
+
+        🐛 Cockroach bug: https://github.com/cockroachdb/cockroach/issues/71908 "could not decorrelate subquery" error under asyncpg
+
+        """  # noqa: E501
+        response = Band.select(
+            Band.name, Band.genres(Genre.all_columns(exclude=(Genre.id,)))
+        ).run_sync()
+        self.assertEqual(
+            response,
+            [
+                {
+                    "name": "Pythonistas",
+                    "genres": [
+                        {"name": "Rock"},
+                        {"name": "Folk"},
+                    ],
+                },
+                {"name": "Rustaceans", "genres": [{"name": "Folk"}]},
+                {
+                    "name": "C-Sharps",
+                    "genres": [
+                        {"name": "Rock"},
+                        {"name": "Classical"},
+                    ],
+                },
+            ],
+        )
+
     def test_add_m2m(self):
         """
         Make sure we can add items to the joining table.

tests/utils/test_list.py

@@ -0,0 +1,8 @@
+from unittest import TestCase
+
+from piccolo.utils.list import flatten
+
+
+class TestFlatten(TestCase):
+    def test_flatten(self):
+        self.assertListEqual(flatten(["a", ["b", "c"]]), ["a", "b", "c"])