Merge pull request piccolo-orm#120 from piccolo-orm/fix_protected_tablename_bug

Fix protected tablename bug

Author: dantownsend

piccolo/apps/migrations/auto/diffable_table.py

@@ -12,7 +12,7 @@
     serialise_params,
     deserialise_params,
 )
-from piccolo.table import Table
+from piccolo.table import Table, create_table_class
 
 
 def compare_dicts(dict_1, dict_2) -> t.Dict[str, t.Any]:
@@ -144,10 +144,12 @@ def to_table_class(self) -> t.Type[Table]:
         """
         Converts the DiffableTable into a Table subclass.
         """
-        _Table: t.Type[Table] = type(
-            self.class_name,
-            (Table,),
-            {column._meta.name: column for column in self.columns},
+        _Table: t.Type[Table] = create_table_class(
+            class_name=self.class_name,
+            class_kwargs={"tablename": self.tablename},
+            class_members={
+                column._meta.name: column for column in self.columns
+            },
         )
-        _Table._meta.tablename = self.tablename
+
         return _Table

piccolo/apps/migrations/auto/migration_manager.py

@@ -13,7 +13,7 @@
     RenameTable,
 )
 from piccolo.apps.migrations.auto.serialisation import deserialise_params
-from piccolo.table import Table
+from piccolo.table import Table, create_table_class
 
 
 @dataclass
@@ -329,8 +329,10 @@ async def _run_alter_columns(self, backwards=False):
             if not alter_columns:
                 continue
 
-            _Table: t.Type[Table] = type(table_class_name, (Table,), {})
-            _Table._meta.tablename = alter_columns[0].tablename
+            _Table: t.Type[Table] = create_table_class(
+                class_name=table_class_name,
+                class_kwargs={"tablename": alter_columns[0].tablename},
+            )
 
             for alter_column in alter_columns:
 
@@ -492,8 +494,10 @@ async def _run_drop_columns(self, backwards=False):
                 if not columns:
                     continue
 
-                _Table: t.Type[Table] = type(table_class_name, (Table,), {})
-                _Table._meta.tablename = columns[0].tablename
+                _Table: t.Type[Table] = create_table_class(
+                    class_name=table_class_name,
+                    class_kwargs={"tablename": columns[0].tablename},
+                )
 
                 for column in columns:
                     await _Table.alter().drop_column(
@@ -518,8 +522,9 @@ async def _run_rename_tables(self, backwards=False):
                 else rename_table.new_tablename
             )
 
-            _Table: t.Type[Table] = type(class_name, (Table,), {})
-            _Table._meta.tablename = tablename
+            _Table: t.Type[Table] = create_table_class(
+                class_name=class_name, class_kwargs={"tablename": tablename}
+            )
 
             await _Table.alter().rename_table(new_name=new_tablename).run()
 
@@ -532,8 +537,10 @@ async def _run_rename_columns(self, backwards=False):
             if not columns:
                 continue
 
-            _Table: t.Type[Table] = type(table_class_name, (Table,), {})
-            _Table._meta.tablename = columns[0].tablename
+            _Table: t.Type[Table] = create_table_class(
+                class_name=table_class_name,
+                class_kwargs={"tablename": columns[0].tablename},
+            )
 
             for rename_column in columns:
                 column = (
@@ -559,15 +566,14 @@ async def _run_add_tables(self, backwards=False):
                 ).run()
         else:
             for add_table in self.add_tables:
-                _Table: t.Type[Table] = type(
-                    add_table.class_name,
-                    (Table,),
-                    {
+                _Table: t.Type[Table] = create_table_class(
+                    class_name=add_table.class_name,
+                    class_kwargs={"tablename": add_table.tablename},
+                    class_members={
                         column._meta.name: column
                         for column in add_table.columns
                     },
                 )
-                _Table._meta.tablename = add_table.tablename
 
                 await _Table.create_table().run()
 
@@ -584,10 +590,10 @@ async def _run_add_columns(self, backwards=False):
                     # be deleted.
                     continue
 
-                _Table: t.Type[Table] = type(
-                    add_column.table_class_name, (Table,), {}
+                _Table: t.Type[Table] = create_table_class(
+                    class_name=add_column.table_class_name,
+                    class_kwargs={"tablename": add_column.tablename},
                 )
-                _Table._meta.tablename = add_column.tablename
 
                 await _Table.alter().drop_column(add_column.column).run()
         else:
@@ -598,15 +604,14 @@ async def _run_add_columns(self, backwards=False):
 
                 # Define the table, with the columns, so the metaclass
                 # sets up the columns correctly.
-                _Table: t.Type[Table] = type(
-                    add_columns[0].table_class_name,
-                    (Table,),
-                    {
+                _Table: t.Type[Table] = create_table_class(
+                    class_name=add_columns[0].table_class_name,
+                    class_kwargs={"tablename": add_columns[0].tablename},
+                    class_members={
                         add_column.column._meta.name: add_column.column
                         for add_column in add_columns
                     },
                 )
-                _Table._meta.tablename = add_columns[0].tablename
 
                 for add_column in add_columns:
                     # We fetch the column from the Table, as the metaclass

piccolo/apps/migrations/auto/serialisation_legacy.py

@@ -3,7 +3,7 @@
 import typing as t
 
 from piccolo.columns.column_types import OnDelete, OnUpdate
-from piccolo.table import Table
+from piccolo.table import Table, create_table_class
 from piccolo.columns.defaults.timestamp import TimestampNow
 
 
@@ -25,11 +25,10 @@ def deserialise_legacy_params(name: str, value: str) -> t.Any:
                 "`SomeClassName` or `SomeClassName|some_table_name`."
             )
 
-        _Table: t.Type[Table] = type(
-            class_name, (Table,), {},
+        _Table: t.Type[Table] = create_table_class(
+            class_name=class_name,
+            class_kwargs={"tablename": tablename} if tablename else {},
         )
-        if tablename:
-            _Table._meta.tablename = tablename
         return _Table
 
     ###########################################################################

piccolo/table.py

@@ -3,6 +3,7 @@
 from enum import Enum
 import inspect
 import itertools
+import types
 import typing as t
 
 from piccolo.engine import Engine, engine_finder
@@ -719,3 +720,32 @@ def _table_str(cls, abbreviated=False):
         return (
             f"class {cls.__name__}({class_args}):\n" f"    {columns_string}\n"
         )
+
+
+def create_table_class(
+    class_name: str,
+    bases: t.Tuple[t.Type] = (Table,),
+    class_kwargs: t.Dict[str, t.Any] = {},
+    class_members: t.Dict[str, t.Any] = {},
+) -> t.Type[Table]:
+    """
+    Used to dynamically create ``Table``subclasses at runtime. Most users
+    will not require this. It's mostly used internally for Piccolo's
+    migrations.
+
+    :param class_name:
+        For example `'MyTable'`.
+    :param bases:
+        A tuple of parent classes - usually just `(Table,)`.
+    :param class_kwargs:
+        For example, `{'tablename': 'my_custom_tablename'}`.
+    :param class_members:
+        For example, `{'my_column': Varchar()}`.
+
+    """
+    return types.new_class(
+        name=class_name,
+        bases=bases,
+        kwds=class_kwargs,
+        exec_body=lambda namespace: namespace.update(class_members),
+    )

tests/base.py

@@ -8,7 +8,7 @@
 from piccolo.engine.finder import engine_finder
 from piccolo.engine.postgres import PostgresEngine
 from piccolo.engine.sqlite import SQLiteEngine
-from piccolo.table import Table
+from piccolo.table import Table, create_table_class
 
 
 ENGINE = engine_finder()
@@ -47,12 +47,13 @@ class DBTestCase(TestCase):
     """
 
     def run_sync(self, query):
-        _Table = type("_Table", (Table,), {})
+        _Table = create_table_class(class_name="_Table")
         return _Table.raw(query).run_sync()
 
     def table_exists(self, tablename: str) -> bool:
-        _Table: t.Type[Table] = type(tablename.upper(), (Table,), {})
-        _Table._meta.tablename = tablename
+        _Table: t.Type[Table] = create_table_class(
+            class_name=tablename.upper(), class_kwargs={"tablename": tablename}
+        )
         return _Table.table_exists().run_sync()
 
     ###########################################################################

tests/table/test_create_table_class.py

@@ -0,0 +1,42 @@
+from unittest import TestCase
+
+from piccolo.table import create_table_class
+from piccolo.columns import Varchar
+
+
+class TestCreateTableClass(TestCase):
+    def test_create_table_class(self):
+        """
+        Make sure a basic `Table` can be created successfully.
+        """
+        _Table = create_table_class(class_name="MyTable")
+        self.assertEqual(_Table._meta.tablename, "my_table")
+
+        _Table = create_table_class(
+            class_name="MyTable", class_kwargs={"tablename": "my_table_1"}
+        )
+        self.assertEqual(_Table._meta.tablename, "my_table_1")
+
+        column = Varchar()
+        _Table = create_table_class(
+            class_name="MyTable", class_members={"name": column}
+        )
+        self.assertTrue(column in _Table._meta.columns)
+
+    def test_protected_tablenames(self):
+        """
+        Make sure that the logic around protected tablenames still works as
+        expected.
+        """
+        with self.assertRaises(ValueError):
+            create_table_class(class_name="User")
+
+        with self.assertRaises(ValueError):
+            create_table_class(
+                class_name="MyUser", class_kwargs={"tablename": "user"}
+            )
+
+        # This shouldn't raise an error:
+        create_table_class(
+            class_name="User", class_kwargs={"tablename": "my_user"}
+        )