be able to select all related columns (piccolo-orm#161)

Author: dantownsend

piccolo/columns/column_types.py

@@ -1006,13 +1006,20 @@ class Band(Table):
 
     **Joins**
 
-    Can also use it to perform joins:
+    You also use it to perform joins:
 
     .. code-block:: python
 
         >>> await Band.select(Band.name, Band.manager.name).first().run()
         {'name': 'Pythonistas', 'manager.name': 'Guido'}
 
+    To retrieve all of the columns in the related table:
+
+    .. code-block:: python
+
+        >>> await Band.select(Band.name, *Band.manager.all_columns()).first().run()
+        {'name': 'Pythonistas', 'manager.id': 1, 'manager.name': 'Guido'}
+
     To get a referenced row as an object:
 
     .. code-block:: python
@@ -1217,11 +1224,27 @@ def set_proxy_columns(self):
         the ``ForeignKey`` column for each column in the table being pointed
         to.
         """
-        _foreign_key_meta = object.__getattribute__(self, "_foreign_key_meta")
-        for column in _foreign_key_meta.resolved_references._meta.columns:
+        _fk_meta = object.__getattribute__(self, "_foreign_key_meta")
+        for column in _fk_meta.resolved_references._meta.columns:
             _column: Column = column.copy()
             setattr(self, _column._meta.name, _column)
-            _foreign_key_meta.proxy_columns.append(_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):
         """

tests/columns/test_foreignkey.py

@@ -3,32 +3,34 @@
 
 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
 
 
-class Manager(Table):
+class Manager1(Table, tablename="manager"):
     name = Varchar()
     manager = ForeignKey("self", null=True)
 
 
-class Band1(Table):
-    manager = ForeignKey(references=Manager)
+class Band1(Table, tablename="band"):
+    manager = ForeignKey(references=Manager1)
 
 
-class Band2(Table):
-    manager = ForeignKey(references="Manager")
+class Band2(Table, tablename="band"):
+    manager = ForeignKey(references="Manager1")
 
 
-class Band3(Table):
+class Band3(Table, tablename="band"):
     manager = ForeignKey(
         references=LazyTableReference(
-            table_class_name="Manager",
+            table_class_name="Manager1",
             module_path="tests.columns.test_foreignkey",
         )
     )
 
 
-class Band4(Table):
-    manager = ForeignKey(references="tests.columns.test_foreignkey.Manager")
+class Band4(Table, tablename="band"):
+    manager = ForeignKey(references="tests.columns.test_foreignkey.Manager1")
 
 
 class TestForeignKeySelf(TestCase):
@@ -38,18 +40,18 @@ class TestForeignKeySelf(TestCase):
     """
 
     def setUp(self):
-        Manager.create_table().run_sync()
+        Manager1.create_table().run_sync()
 
     def test_foreign_key_self(self):
-        manager = Manager(name="Mr Manager")
+        manager = Manager1(name="Mr Manager")
         manager.save().run_sync()
 
-        worker = Manager(name="Mr Worker", manager=manager.id)
+        worker = Manager1(name="Mr Worker", manager=manager.id)
         worker.save().run_sync()
 
         response = (
-            Manager.select(Manager.name, Manager.manager.name)
-            .order_by(Manager.name)
+            Manager1.select(Manager1.name, Manager1.manager.name)
+            .order_by(Manager1.name)
             .run_sync()
         )
         self.assertEqual(
@@ -61,7 +63,7 @@ def test_foreign_key_self(self):
         )
 
     def tearDown(self):
-        Manager.alter().drop_table().run_sync()
+        Manager1.alter().drop_table().run_sync()
 
 
 class TestForeignKeyString(TestCase):
@@ -71,23 +73,25 @@ class TestForeignKeyString(TestCase):
     """
 
     def setUp(self):
-        Manager.create_table().run_sync()
+        Manager1.create_table().run_sync()
 
     def test_foreign_key_string(self):
         Band2.create_table().run_sync()
         self.assertEqual(
-            Band2.manager._foreign_key_meta.resolved_references, Manager
+            Band2.manager._foreign_key_meta.resolved_references,
+            Manager1,
         )
         Band2.alter().drop_table().run_sync()
 
         Band4.create_table().run_sync()
         self.assertEqual(
-            Band4.manager._foreign_key_meta.resolved_references, Manager
+            Band4.manager._foreign_key_meta.resolved_references,
+            Manager1,
         )
         Band4.alter().drop_table().run_sync()
 
     def tearDown(self):
-        Manager.alter().drop_table().run_sync()
+        Manager1.alter().drop_table().run_sync()
 
 
 class TestForeignKeyRelativeError(TestCase):
@@ -98,7 +102,7 @@ def test_foreign_key_relative_error(self):
         """
         with self.assertRaises(ValueError) as manager:
 
-            class Band(Table):
+            class BandRelative(Table, tablename="band"):
                 manager = ForeignKey("..example_app.tables.Manager", null=True)
 
         self.assertEqual(
@@ -112,14 +116,14 @@ def test_foreign_key_references(self):
         Make sure foreign key references are stored correctly on the table
         which is the target of the ForeignKey.
         """
-        self.assertEqual(len(Manager._meta.foreign_key_references), 5)
+        self.assertEqual(len(Manager1._meta.foreign_key_references), 5)
 
         self.assertTrue(Band1.manager in Manager._meta.foreign_key_references)
         self.assertTrue(Band2.manager in Manager._meta.foreign_key_references)
         self.assertTrue(Band3.manager in Manager._meta.foreign_key_references)
         self.assertTrue(Band4.manager in Manager._meta.foreign_key_references)
         self.assertTrue(
-            Manager.manager in Manager._meta.foreign_key_references
+            Manager1.manager in Manager1._meta.foreign_key_references
         )
 
 
@@ -128,8 +132,6 @@ def test_lazy_reference_to_app(self):
         """
         Make sure a LazyTableReference to a Table within a Piccolo app works.
         """
-        from tests.example_app.tables import Manager
-
         reference = LazyTableReference(
             table_class_name="Manager", app_name="example_app"
         )
@@ -153,18 +155,49 @@ def test_recursion_limit(self):
         if the call chain is too large.
         """
         # Should be fine:
-        column: Column = Manager.manager.name
+        column: Column = Manager1.manager.name
         self.assertTrue(len(column._meta.call_chain), 1)
         self.assertTrue(isinstance(column, Varchar))
 
         with self.assertRaises(Exception):
-            Manager.manager.manager.manager.manager.manager.manager.manager.manager.manager.manager.manager.name  # noqa
+            Manager1.manager.manager.manager.manager.manager.manager.manager.manager.manager.manager.manager.name  # noqa
 
     def test_recursion_time(self):
         """
         Make sure that a really large call chain doesn't take too long.
         """
         start = time.time()
-        Manager.manager.manager.manager.manager.manager.manager.name
+        Manager1.manager.manager.manager.manager.manager.manager.name
         end = time.time()
         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()
+
+    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()
+        self.assertEqual(
+            result,
+            [
+                {
+                    "name": "Pythonistas",
+                    "manager.id": 1,
+                    "manager.name": "Guido",
+                }
+            ],
+        )