diff --git a/piccolo/utils/pydantic.py b/piccolo/utils/pydantic.py index 7a456db6a..9f88ce35a 100644 --- a/piccolo/utils/pydantic.py +++ b/piccolo/utils/pydantic.py @@ -18,6 +18,7 @@ ForeignKey, Numeric, Text, + Timestamptz, Varchar, ) from piccolo.table import Table @@ -297,6 +298,8 @@ def create_pydantic_model( extra["widget"] = "text-area" elif isinstance(column, (JSON, JSONB)): extra["widget"] = "json" + elif isinstance(column, Timestamptz): + extra["widget"] = "timestamptz" field = pydantic.Field( json_schema_extra={"extra": extra}, diff --git a/tests/utils/test_pydantic.py b/tests/utils/test_pydantic.py index 3ac66afa7..35aff0fa4 100644 --- a/tests/utils/test_pydantic.py +++ b/tests/utils/test_pydantic.py @@ -18,6 +18,8 @@ Secret, Text, Time, + Timestamp, + Timestamptz, Varchar, ) from piccolo.columns.column_types import ForeignKey @@ -201,6 +203,29 @@ class Concert(Table): ) +class TestTimestamptzColumn(TestCase): + def test_timestamptz_widget(self): + """ + Make sure that we indicate that `Timestamptz` columns require a special + widget in Piccolo Admin. + """ + + class Concert(Table): + starts_on_1 = Timestamptz() + starts_on_2 = Timestamp() + + pydantic_model = create_pydantic_model(table=Concert) + + properties = pydantic_model.model_json_schema()["properties"] + + self.assertEqual( + properties["starts_on_1"]["extra"]["widget"], + "timestamptz", + ) + + self.assertIsNone(properties["starts_on_2"]["extra"].get("widget")) + + class TestUUIDColumn(TestCase): class Ticket(Table): code = UUID()