piccolo-orm#242 bug-fix (piccolo-orm#246)

* This should fix piccolo-orm#242

* add test for unknown column types

Co-authored-by: Daniel Townsend <[email protected]>

Author: AliSayyah

piccolo/apps/schema/commands/generate.py

@@ -309,52 +309,52 @@ class Schema(Table, db=engine):
             data_type = pg_row_meta.data_type
             column_type = COLUMN_TYPE_MAP.get(data_type, None)
             column_name = pg_row_meta.column_name
-
-            if column_type:
-                kwargs: t.Dict[str, t.Any] = {
-                    "null": pg_row_meta.is_nullable == "YES",
-                    "unique": constraints.is_unique(column_name=column_name),
-                }
-
-                if constraints.is_primary_key(column_name=column_name):
-                    kwargs["primary_key"] = True
-                    if column_type == Integer:
-                        column_type = Serial
-
-                if constraints.is_foreign_key(column_name=column_name):
-                    fk_constraint_name = (
-                        constraints.get_foreign_key_constraint_name(
-                            column_name=column_name
-                        )
-                    )
-                    column_type = ForeignKey
-                    referenced_tablename = await get_foreign_key_reference(
-                        table_class=Schema, constraint_name=fk_constraint_name
+            if not column_type:
+                warnings.append(f"{tablename}.{column_name} ['{data_type}']")
+                column_type = Column
+
+            kwargs: t.Dict[str, t.Any] = {
+                "null": pg_row_meta.is_nullable == "YES",
+                "unique": constraints.is_unique(column_name=column_name),
+            }
+
+            if constraints.is_primary_key(column_name=column_name):
+                kwargs["primary_key"] = True
+                if column_type == Integer:
+                    column_type = Serial
+
+            if constraints.is_foreign_key(column_name=column_name):
+                fk_constraint_name = (
+                    constraints.get_foreign_key_constraint_name(
+                        column_name=column_name
                     )
-                    if referenced_tablename:
-                        kwargs["references"] = create_table_class(
-                            _snake_to_camel(referenced_tablename)
-                        )
-                    else:
-                        kwargs["references"] = ForeignKeyPlaceholder
-
-                imports.add(
-                    "from piccolo.columns.column_types import "
-                    + column_type.__name__
                 )
+                column_type = ForeignKey
+                referenced_tablename = await get_foreign_key_reference(
+                    table_class=Schema, constraint_name=fk_constraint_name
+                )
+                if referenced_tablename:
+                    kwargs["references"] = create_table_class(
+                        _snake_to_camel(referenced_tablename)
+                    )
+                else:
+                    kwargs["references"] = ForeignKeyPlaceholder
+
+            imports.add(
+                "from piccolo.columns.column_types import "
+                + column_type.__name__
+            )
 
-                if column_type is Varchar:
-                    kwargs["length"] = pg_row_meta.character_maximum_length
+            if column_type is Varchar:
+                kwargs["length"] = pg_row_meta.character_maximum_length
 
-                column = column_type(**kwargs)
+            column = column_type(**kwargs)
 
-                serialised_params = serialise_params(column._meta.params)
-                for extra_import in serialised_params.extra_imports:
-                    imports.add(extra_import.__repr__())
+            serialised_params = serialise_params(column._meta.params)
+            for extra_import in serialised_params.extra_imports:
+                imports.add(extra_import.__repr__())
 
-                columns[column_name] = column
-            else:
-                warnings.append(f"{tablename}.{column_name} ['{data_type}']")
+            columns[column_name] = column
 
         table = create_table_class(
             class_name=_snake_to_camel(tablename),

tests/apps/schema/commands/test_generate.py

@@ -10,6 +10,7 @@
     generate,
     get_output_schema,
 )
+from piccolo.columns.base import Column
 from piccolo.columns.column_types import Varchar
 from piccolo.table import Table
 from piccolo.utils.sync import run_sync
@@ -71,7 +72,7 @@ def test_get_output_schema(self):
         self._compare_table_columns(SmallTable, SmallTable_)
 
     @patch("piccolo.apps.schema.commands.generate.print")
-    def test_generate(self, print_: MagicMock):
+    def test_generate_command(self, print_: MagicMock):
         """
         Test the main generate command runs without errors.
         """
@@ -81,3 +82,29 @@ def test_generate(self, print_: MagicMock):
         # Make sure the output is valid Python code (will raise a SyntaxError
         # exception otherwise).
         ast.parse(file_contents)
+
+    def test_unknown_column_type(self):
+        """
+        Make sure unknown column types are handled gracefully.
+        """
+
+        class Box(Column):
+            """
+            A column type which isn't supported by Piccolo officially yet.
+            """
+
+            pass
+
+        MegaTable.alter().add_column("box", Box()).run_sync()
+
+        output_schema: OutputSchema = run_sync(get_output_schema())
+
+        # Make sure there's a warning.
+        self.assertEqual(output_schema.warnings, ["mega_table.box ['box']"])
+
+        # Make sure the column type of the generated table is just ``Column``.
+        for table in output_schema.tables:
+            if table.__name__ == "MegaTable":
+                self.assertEqual(
+                    output_schema.tables[1].box.__class__.__name__, "Column"
+                )