Merge pull request piccolo-orm#101 from piccolo-orm/secret_columns_meta

added `secret_columns` to `TableMeta`

Author: dantownsend

piccolo/table.py

@@ -5,11 +5,8 @@
 import typing as t
 
 from piccolo.engine import Engine, engine_finder
-from piccolo.columns import (
-    Column,
-    Selectable,
-)
-from piccolo.columns.column_types import ForeignKey, PrimaryKey
+from piccolo.columns import Column, Selectable
+from piccolo.columns.column_types import ForeignKey, PrimaryKey, Secret
 from piccolo.columns.readable import Readable
 from piccolo.columns.reference import (
     LazyTableReference,
@@ -51,6 +48,7 @@ class TableMeta:
     default_columns: t.List[Column] = field(default_factory=list)
     non_default_columns: t.List[Column] = field(default_factory=list)
     foreign_key_columns: t.List[ForeignKey] = field(default_factory=list)
+    secret_columns: t.List[Secret] = field(default_factory=list)
     tags: t.List[str] = field(default_factory=list)
     help_text: t.Optional[str] = None
     _db: t.Optional[Engine] = None
@@ -158,6 +156,7 @@ def __init_subclass__(
         default_columns: t.List[Column] = []
         non_default_columns: t.List[Column] = []
         foreign_key_columns: t.List[ForeignKey] = []
+        secret_columns: t.List[Secret] = []
 
         cls.id = PrimaryKey()
 
@@ -190,6 +189,9 @@ def __init_subclass__(
                 column._meta._name = attribute_name
                 column._meta._table = cls
 
+                if isinstance(column, Secret):
+                    secret_columns.append(column)
+
                 if isinstance(column, ForeignKey):
                     foreign_key_columns.append(column)
 
@@ -199,6 +201,7 @@ def __init_subclass__(
             default_columns=default_columns,
             non_default_columns=non_default_columns,
             foreign_key_columns=foreign_key_columns,
+            secret_columns=secret_columns,
             tags=tags,
             help_text=help_text,
             _db=db,

tests/table/test_metaclass.py

@@ -1,5 +1,7 @@
+from piccolo.columns.column_types import ForeignKey
 from unittest import TestCase
 
+from piccolo.columns import Secret
 from piccolo.table import Table
 
 from ..example_app.tables import Band
@@ -34,3 +36,30 @@ class Manager(Table, help_text=help_text):
             pass
 
         self.assertEqual(Manager._meta.help_text, help_text)
+
+    def test_foreign_key_columns(self):
+        """
+        Make sure TableMeta.foreign_keys and TableMeta.foreign_key_references
+        are setup correctly.
+        """
+
+        class TableA(Table):
+            pass
+
+        class TableB(Table):
+            table_a = ForeignKey(references=TableA)
+
+        self.assertEqual(TableB._meta.foreign_key_columns, [TableB.table_a])
+        self.assertEqual(TableA._meta.foreign_key_references, [TableB.table_a])
+
+    def test_secret_columns(self):
+        """
+        Make sure TableMeta.secret_columns are setup correctly.
+        """
+
+        class Classified(Table):
+            top_secret = Secret()
+
+        self.assertEqual(
+            Classified._meta.secret_columns, [Classified.top_secret]
+        )