added initial tests

Author: dantownsend

piccolo/apps/schema/commands/generate.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import dataclasses
 import typing as t
 
@@ -6,11 +8,17 @@
 
 from piccolo.columns.base import Column
 from piccolo.columns.column_types import (
+    JSON,
+    JSONB,
     UUID,
     BigInt,
     Boolean,
+    Bytea,
+    Date,
     Integer,
+    Interval,
     Numeric,
+    Real,
     SmallInt,
     Text,
     Timestamp,
@@ -50,16 +58,47 @@ def get_column_name_str(cls) -> str:
         return ", ".join([i.name for i in dataclasses.fields(cls)])
 
 
+@dataclasses.dataclass
+class OutputSchema:
+    """
+    Represents the schema which will printed output.
+
+    :param imports:
+        e.g. ["from piccolo.table import Table"]
+    :param warnings:
+        e.g. ["some_table.some_column unrecognised_type"]
+    :param tables:
+        e.g. ["class MyTable(Table): ..."]
+
+    """
+
+    imports: t.List[str]
+    warnings: t.List[str]
+    tables: t.List[t.Type[Table]]
+
+    def get_table_with_name(self, name: str) -> t.Type[Table]:
+        """
+        Just used by unit tests.
+        """
+        return next(table for table in self.tables if table.__name__ == name)
+
+
 COLUMN_TYPE_MAP = {
     "bigint": BigInt,
     "boolean": Boolean,
+    "bytea": Bytea,
     "character varying": Varchar,
+    "date": Date,
     "integer": Integer,
+    "interval": Interval,
+    "json": JSON,
+    "jsonb": JSONB,
     "numeric": Numeric,
+    "real": Real,
     "smallint": SmallInt,
     "text": Text,
-    "timestamp without time zone": Timestamp,
     "timestamp with time zone": Timestamptz,
+    "timestamp without time zone": Timestamp,
     "uuid": UUID,
 }
 
@@ -129,15 +168,7 @@ async def get_table_schema(
     return [PostgresRowMeta(**row_meta) for row_meta in row_meta_list]
 
 
-# This is currently a beta version, and can be improved. However, having
-# something working is still useful for people migrating large schemas to
-# Piccolo.
-async def generate(schema_name: str = "public"):
-    """
-    Automatically generates Piccolo Table classes by introspecting the
-    database. Please check the generated code in case there are errors.
-
-    """
+async def get_output_schema(schema_name: str = "public") -> OutputSchema:
     engine: t.Optional[Engine] = engine_finder()
 
     if engine is None:
@@ -152,11 +183,15 @@ async def generate(schema_name: str = "public"):
         )
 
     class Schema(Table, db=engine):
+        """
+        Just used for making raw queries on the db.
+        """
+
         pass
 
     tablenames = await get_tablenames(Schema, schema_name=schema_name)
 
-    output: t.List[str] = []
+    tables: t.List[t.Type[Table]] = []
     imports: t.Set[str] = {"from piccolo.table import Table"}
     warnings: t.List[str] = []
 
@@ -191,12 +226,30 @@ class Schema(Table, db=engine):
             class_kwargs={"tablename": tablename},
             class_members=columns,
         )
-        output.append(table._table_str())
+        tables.append(table)
+
+    return OutputSchema(
+        imports=sorted(list(imports)), warnings=warnings, tables=tables
+    )
+
+
+# This is currently a beta version, and can be improved. However, having
+# something working is still useful for people migrating large schemas to
+# Piccolo.
+async def generate(schema_name: str = "public"):
+    """
+    Automatically generates Piccolo Table classes by introspecting the
+    database. Please check the generated code in case there are errors.
+
+    """
+    output_schema = await get_output_schema(schema_name=schema_name)
 
-    output = sorted(list(imports)) + output
+    output = output_schema.imports + [
+        i._table_str() for i in output_schema.tables
+    ]
 
-    if warnings:
-        warning_str = "\n".join(warnings)
+    if output_schema.warnings:
+        warning_str = "\n".join(output_schema.warnings)
 
         output.append('"""')
         output.append(

tests/apps/schema/__init__.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import typing as t
+from unittest import TestCase
+
+from piccolo.apps.schema.commands.generate import (
+    OutputSchema,
+    get_output_schema,
+)
+from piccolo.columns.column_types import (
+    JSON,
+    JSONB,
+    UUID,
+    BigInt,
+    Boolean,
+    Bytea,
+    Date,
+    ForeignKey,
+    Integer,
+    Interval,
+    Numeric,
+    Real,
+    SmallInt,
+    Text,
+    Timestamp,
+    Timestamptz,
+    Varchar,
+)
+from piccolo.table import Table
+from piccolo.utils.sync import run_sync
+
+
+class SmallTable(Table):
+    varchar_col = Varchar()
+
+
+class MegaTable(Table):
+    """
+    A table containing all of the column types, and different column kwargs.
+    """
+
+    bigint_col = BigInt()
+    boolean_col = Boolean()
+    bytea_col = Bytea()
+    date_col = Date()
+    foreignkey_col = ForeignKey(SmallTable)
+    integer_col = Integer()
+    interval_col = Interval()
+    json_col = JSON()
+    jsonb_col = JSONB()
+    numeric_col = Numeric()
+    real_col = Real()
+    smallint_col = SmallInt()
+    text_col = Text()
+    timestamp_col = Timestamp()
+    timestamptz_col = Timestamptz()
+    uuid_col = UUID()
+    varchar_col = Varchar()
+
+    unique_col = Varchar(unique=True)
+    null_col = Varchar(null=True)
+    not_null_col = Varchar(null=False)
+
+
+class TestGenerate(TestCase):
+    def setUp(self):
+        for table_class in (SmallTable, MegaTable):
+            table_class.create_table().run_sync()
+
+    def tearDown(self):
+        for table_class in (MegaTable, SmallTable):
+            table_class.alter().drop_table().run_sync()
+
+    def _compare_table_columns(
+        self, table_1: t.Type[Table], table_2: t.Type[Table]
+    ):
+        """
+        Make sure that for each column in table_1, there is a corresponding
+        column in table_2 of the same type.
+        """
+        column_names = [
+            column._meta.name for column in table_1._meta.non_default_columns
+        ]
+        for column_name in column_names:
+            self.assertTrue(
+                type(table_1._meta.get_column_by_name(column_name)),
+                type(table_2._meta.get_column_by_name(column_name)),
+            )
+
+    def test_get_output_schema(self):
+        """
+        Make sure that the a Piccolo schema can be generated from the database.
+        """
+        output_schema: OutputSchema = run_sync(get_output_schema())
+
+        self.assertTrue(len(output_schema.warnings) == 0)
+        self.assertTrue(len(output_schema.tables) == 2)
+        self.assertTrue(len(output_schema.imports) > 0)
+
+        MegaTable_ = output_schema.get_table_with_name("MegaTable")
+        self._compare_table_columns(MegaTable, MegaTable_)
+
+        SmallTable_ = output_schema.get_table_with_name("SmallTable")
+        self._compare_table_columns(SmallTable, SmallTable_)

tests/apps/schema/commands/test_generate.py

@@ -18,6 +18,8 @@ async def drop_tables():
         "my_table",
         "recording_studio",
         "shirt",
+        "mega_table",
+        "small_table",
     ]:
         await ENGINE._run_in_new_connection(f"DROP TABLE IF EXISTS {table}")