all_columns fix (piccolo-orm#210)

* fixing bug with all_columns if called from too deep

* ``ColumnsDelegate`` auto unpacks lists of columns

This is useful if someone uses ``all_columns`` and forgets to unpack them.

await Band.select(Band.name, Band.manager.all_columns()).run()

The above now works - before this was required:

await Band.select(Band.name, *Band.manager.all_columns()).run()

* update docs, to mention that unpacking `all_columns` is now optional

Author: dantownsend

docs/src/piccolo/query_types/select.rst

@@ -61,8 +61,9 @@ The joins can go several layers deep.
     >>> Concert.select(Concert.id, Concert.band_1.manager.name).run_sync()
     [{'id': 1, 'band_1.manager.name': 'Guido'}]
 
-If you want all of the columns from a related table, there's a useful shortcut
-which saves you from typing them all out:
+If you want all of the columns from a related table you can use
+``all_columns``, which is a useful shortcut which saves you from typing them
+all out:
 
 .. code-block:: python
 
@@ -72,6 +73,10 @@ which saves you from typing them all out:
         {'name': 'Rustaceans', 'manager.id': 2, 'manager.name': 'Graydon'}
     ]
 
+    # In Piccolo > 0.41.0 you no longer need to explicitly unpack ``all_columns``.
+    # This is equivalent:
+    >>> Band.select(Band.name, Band.manager.all_columns()).run_sync()
+
 You can also get the response as nested dictionaries, which can be very useful:
 
 .. code-block:: python

piccolo/columns/base.py

@@ -137,7 +137,7 @@ class ColumnMeta:
     _table: t.Optional[t.Type[Table]] = None
 
     # Used by Foreign Keys:
-    call_chain: t.List["ForeignKey"] = field(default_factory=lambda: [])
+    call_chain: t.List["ForeignKey"] = field(default_factory=list)
     table_alias: t.Optional[str] = None
 
     @property

piccolo/columns/column_types.py

@@ -1220,6 +1220,22 @@ def copy(self) -> ForeignKey:
         column._foreign_key_meta = self._foreign_key_meta.copy()
         return column
 
+    def all_columns(self):
+        """
+        Allow a user to access all of the columns on the related table.
+
+        For example:
+
+        Band.select(Band.name, *Band.manager.all_columns()).run_sync()
+
+        """
+        _fk_meta = object.__getattribute__(self, "_foreign_key_meta")
+
+        return [
+            getattr(self, column._meta.name)
+            for column in _fk_meta.resolved_references._meta.columns
+        ]
+
     def set_proxy_columns(self):
         """
         In order to allow a fluent interface, where tables can be traversed
@@ -1233,22 +1249,6 @@ def set_proxy_columns(self):
             setattr(self, _column._meta.name, _column)
             _fk_meta.proxy_columns.append(_column)
 
-        def all_columns():
-            """
-            Allow a user to access all of the columns on the related table.
-
-            For example:
-
-            Band.select(Band.name, *Band.manager.all_columns()).run_sync()
-
-            """
-            return [
-                getattr(self, column._meta.name)
-                for column in _fk_meta.resolved_references._meta.columns
-            ]
-
-        setattr(self, "all_columns", all_columns)
-
     def __getattribute__(self, name: str):
         """
         Returns attributes unmodified unless they're Column instances, in which

piccolo/query/mixins.py

@@ -237,8 +237,22 @@ class ColumnsDelegate:
 
     selected_columns: t.Sequence[Selectable] = field(default_factory=list)
 
-    def columns(self, *columns: Selectable):
-        combined = list(self.selected_columns) + list(columns)
+    def columns(self, *columns: t.Union[Selectable, t.List[Selectable]]):
+        """
+        :param columns:
+            We accept ``Selectable`` and ``List[Selectable]`` here, in case
+            someone passes in a list by accident when using ``all_columns()``,
+            in which case we flatten the list.
+
+        """
+        _columns = []
+        for column in columns:
+            if isinstance(column, list):
+                _columns.extend(column)
+            else:
+                _columns.append(column)
+
+        combined = list(self.selected_columns) + _columns
         self.selected_columns = combined
 
     def remove_secret_columns(self):

tests/columns/test_foreignkey.py

@@ -3,8 +3,7 @@
 
 from piccolo.columns import Column, ForeignKey, LazyTableReference, Varchar
 from piccolo.table import Table
-from tests.base import DBTestCase
-from tests.example_app.tables import Band, Manager
+from tests.example_app.tables import Band, Concert, Manager
 
 
 class Manager1(Table, tablename="manager"):
@@ -172,32 +171,36 @@ def test_recursion_time(self):
         self.assertTrue(end - start < 1.0)
 
 
-class TestAllColumns(DBTestCase):
-    def setUp(self):
-        Manager.create_table().run_sync()
-        manager = Manager(name="Guido")
-        manager.save().run_sync()
-
-        Band.create_table().run_sync()
-        Band(manager=manager, name="Pythonistas").save().run_sync()
-
-    def tearDown(self):
-        Band.alter().drop_table().run_sync()
-        Manager.alter().drop_table().run_sync()
-
+class TestAllColumns(TestCase):
     def test_all_columns(self):
         """
         Make sure you can retrieve all columns from a related table, without
         explicitly specifying them.
         """
-        result = Band.select(Band.name, *Band.manager.all_columns()).run_sync()
+        all_columns = Band.manager.all_columns()
+        self.assertEqual(all_columns, [Band.manager.id, Band.manager.name])
+
+        # Make sure the call chains are also correct.
         self.assertEqual(
-            result,
-            [
-                {
-                    "name": "Pythonistas",
-                    "manager.id": 1,
-                    "manager.name": "Guido",
-                }
-            ],
+            all_columns[0]._meta.call_chain, Band.manager.id._meta.call_chain
+        )
+        self.assertEqual(
+            all_columns[1]._meta.call_chain, Band.manager.name._meta.call_chain
+        )
+
+    def test_all_columns_deep(self):
+        """
+        Make sure ``all_columns`` works when the joins are several layers deep.
+        """
+        all_columns = Concert.band_1.manager.all_columns()
+        self.assertEqual(all_columns, [Band.manager.id, Band.manager.name])
+
+        # Make sure the call chains are also correct.
+        self.assertEqual(
+            all_columns[0]._meta.call_chain,
+            Concert.band_1.manager.id._meta.call_chain,
+        )
+        self.assertEqual(
+            all_columns[1]._meta.call_chain,
+            Concert.band_1.manager.name._meta.call_chain,
         )

tests/query/test_mixins.py

@@ -0,0 +1,28 @@
+from unittest import TestCase
+
+from piccolo.query.mixins import ColumnsDelegate
+from tests.example_app.tables import Band
+
+
+class TestColumnsDelegate(TestCase):
+    def test_list_unpacking(self):
+        """
+        The ``ColumnsDelegate`` should unpack a list of columns if passed in by
+        mistake, without the user unpacking them explicitly.
+
+        .. code-block:: python
+
+            # These two should both work the same:
+            await Band.select([Band.id, Band.name]).run()
+            await Band.select(Band.id, Band.name).run()
+
+        """
+        columns_delegate = ColumnsDelegate()
+
+        columns_delegate.columns([Band.name])
+        self.assertEqual(columns_delegate.selected_columns, [Band.name])
+
+        columns_delegate.columns([Band.id])
+        self.assertEqual(
+            columns_delegate.selected_columns, [Band.name, Band.id]
+        )

tests/table/test_delete.py

@@ -11,7 +11,6 @@ def test_delete(self):
         Band.delete().where(Band.name == "CSharps").run_sync()
 
         response = Band.count().where(Band.name == "CSharps").run_sync()
-        print(f"response = {response}")
 
         self.assertEqual(response, 0)
 

tests/table/test_join.py

@@ -25,11 +25,6 @@ def setUp(self):
         for table in self.tables:
             table.create_table().run_sync()
 
-    def tearDown(self):
-        for table in reversed(self.tables):
-            table.alter().drop_table().run_sync()
-
-    def test_join(self):
         manager_1 = Manager(name="Guido")
         manager_1.save().run_sync()
 
@@ -42,14 +37,19 @@ def test_join(self):
         band_2 = Band(name="Rustaceans", manager=manager_2.id)
         band_2.save().run_sync()
 
-        venue = Venue(name="Grand Central")
+        venue = Venue(name="Grand Central", capacity=1000)
         venue.save().run_sync()
 
         save_query = Concert(
             band_1=band_1.id, band_2=band_2.id, venue=venue.id
         ).save()
         save_query.run_sync()
 
+    def tearDown(self):
+        for table in reversed(self.tables):
+            table.alter().drop_table().run_sync()
+
+    def test_join(self):
         select_query = Concert.select(
             Concert.band_1.name,
             Concert.band_2.name,
@@ -73,3 +73,49 @@ def test_join(self):
         select_query = Concert.select(Concert.band_1.manager.name)
         response = select_query.run_sync()
         self.assertEqual(response, [{"band_1.manager.name": "Guido"}])
+
+    def test_all_columns(self):
+        """
+        Make sure you can retrieve all columns from a related table, without
+        explicitly specifying them.
+        """
+        result = (
+            Band.select(Band.name, *Band.manager.all_columns())
+            .first()
+            .run_sync()
+        )
+        self.assertEqual(
+            result,
+            {
+                "name": "Pythonistas",
+                "manager.id": 1,
+                "manager.name": "Guido",
+            },
+        )
+
+    def test_all_columns_deep(self):
+        """
+        Make sure that ``all_columns`` can be used several layers deep.
+        """
+        result = (
+            Concert.select(
+                *Concert.venue.all_columns(),
+                *Concert.band_1.manager.all_columns(),
+                *Concert.band_2.manager.all_columns(),
+            )
+            .first()
+            .run_sync()
+        )
+
+        self.assertEqual(
+            result,
+            {
+                "venue.id": 1,
+                "venue.name": "Grand Central",
+                "venue.capacity": 1000,
+                "band_1.manager.id": 1,
+                "band_1.manager.name": "Guido",
+                "band_2.manager.id": 2,
+                "band_2.manager.name": "Graydon",
+            },
+        )

tests/table/test_objects.py

@@ -32,9 +32,6 @@ def test_offset_postgres(self):
         """
         self.insert_rows()
         response = Band.objects().order_by(Band.name).offset(1).run_sync()
-
-        print(f"response = {response}")
-
         self.assertEqual(
             [i.name for i in response], ["Pythonistas", "Rustaceans"]
         )
@@ -54,8 +51,6 @@ def test_offset_sqlite(self):
 
         response = query.run_sync()
 
-        print(f"response = {response}")
-
         self.assertEqual(
             [i.name for i in response], ["Pythonistas", "Rustaceans"]
         )

tests/table/test_select.py

@@ -13,7 +13,6 @@ def test_query_all_columns(self):
         self.insert_row()
 
         response = Band.select().run_sync()
-        print(f"response = {response}")
 
         self.assertDictEqual(
             response[0],
@@ -24,7 +23,6 @@ def test_query_some_columns(self):
         self.insert_row()
 
         response = Band.select(Band.name).run_sync()
-        print(f"response = {response}")
 
         self.assertDictEqual(response[0], {"name": "Pythonistas"})