diff --git a/piccolo/table.py b/piccolo/table.py index e4ddd7daf..cda091064 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -16,7 +16,6 @@ Email, ForeignKey, ReferencedTable, - Secret, Serial, ) from piccolo.columns.defaults.base import Default @@ -84,7 +83,7 @@ class TableMeta: foreign_key_columns: list[ForeignKey] = field(default_factory=list) primary_key: Column = field(default_factory=Column) json_columns: list[Union[JSON, JSONB]] = field(default_factory=list) - secret_columns: list[Secret] = field(default_factory=list) + secret_columns: list[Column] = field(default_factory=list) auto_update_columns: list[Column] = field(default_factory=list) tags: list[str] = field(default_factory=list) help_text: Optional[str] = None @@ -274,7 +273,7 @@ def __init_subclass__( non_default_columns: list[Column] = [] array_columns: list[Array] = [] foreign_key_columns: list[ForeignKey] = [] - secret_columns: list[Secret] = [] + secret_columns: list[Column] = [] json_columns: list[Union[JSON, JSONB]] = [] email_columns: list[Email] = [] auto_update_columns: list[Column] = [] @@ -315,15 +314,15 @@ def __init_subclass__( if isinstance(column, Email): email_columns.append(column) - if isinstance(column, Secret): - secret_columns.append(column) - if isinstance(column, ForeignKey): foreign_key_columns.append(column) if isinstance(column, (JSON, JSONB)): json_columns.append(column) + if column._meta.secret: + secret_columns.append(column) + if column._meta.auto_update is not ...: auto_update_columns.append(column) diff --git a/tests/table/test_metaclass.py b/tests/table/test_metaclass.py index 4af11d1aa..7ff186d73 100644 --- a/tests/table/test_metaclass.py +++ b/tests/table/test_metaclass.py @@ -1,13 +1,13 @@ from unittest import TestCase from unittest.mock import MagicMock, patch -from piccolo.columns import Secret from piccolo.columns.column_types import ( JSON, JSONB, Array, Email, ForeignKey, + Secret, Varchar, ) from piccolo.table import TABLENAME_WARNING, Table @@ -99,14 +99,18 @@ class TableB(Table): def test_secret_columns(self): """ - Make sure TableMeta.secret_columns are setup correctly. + Make sure TableMeta.secret_columns are setup correctly with the + ``secret=True`` argument and ``Secret`` column type. """ class Classified(Table): top_secret = Secret() + confidential = Varchar(secret=True) + public = Varchar() self.assertEqual( - Classified._meta.secret_columns, [Classified.top_secret] + Classified._meta.secret_columns, + [Classified.top_secret, Classified.confidential], ) def test_json_columns(self):