Skip to content

Commit 336bdd7

Browse files
committed
added secret_columns to TableMeta
This makes accessing the secret columns for a table easier.
1 parent ce0e12c commit 336bdd7

File tree

2 files changed

+37
-5
lines changed

2 files changed

+37
-5
lines changed

piccolo/table.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,8 @@
55
import typing as t
66

77
from piccolo.engine import Engine, engine_finder
8-
from piccolo.columns import (
9-
Column,
10-
Selectable,
11-
)
12-
from piccolo.columns.column_types import ForeignKey, PrimaryKey
8+
from piccolo.columns import Column, Selectable
9+
from piccolo.columns.column_types import ForeignKey, PrimaryKey, Secret
1310
from piccolo.columns.readable import Readable
1411
from piccolo.columns.reference import (
1512
LazyTableReference,
@@ -51,6 +48,7 @@ class TableMeta:
5148
default_columns: t.List[Column] = field(default_factory=list)
5249
non_default_columns: t.List[Column] = field(default_factory=list)
5350
foreign_key_columns: t.List[ForeignKey] = field(default_factory=list)
51+
secret_columns: t.List[Secret] = field(default_factory=list)
5452
tags: t.List[str] = field(default_factory=list)
5553
help_text: t.Optional[str] = None
5654
_db: t.Optional[Engine] = None
@@ -158,6 +156,7 @@ def __init_subclass__(
158156
default_columns: t.List[Column] = []
159157
non_default_columns: t.List[Column] = []
160158
foreign_key_columns: t.List[ForeignKey] = []
159+
secret_columns: t.List[Secret] = []
161160

162161
cls.id = PrimaryKey()
163162

@@ -190,6 +189,9 @@ def __init_subclass__(
190189
column._meta._name = attribute_name
191190
column._meta._table = cls
192191

192+
if isinstance(column, Secret):
193+
secret_columns.append(column)
194+
193195
if isinstance(column, ForeignKey):
194196
foreign_key_columns.append(column)
195197

@@ -199,6 +201,7 @@ def __init_subclass__(
199201
default_columns=default_columns,
200202
non_default_columns=non_default_columns,
201203
foreign_key_columns=foreign_key_columns,
204+
secret_columns=secret_columns,
202205
tags=tags,
203206
help_text=help_text,
204207
_db=db,

tests/table/test_metaclass.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
1+
from piccolo.columns.column_types import ForeignKey
12
from unittest import TestCase
23

4+
from piccolo.columns import Secret
35
from piccolo.table import Table
46

57
from ..example_app.tables import Band
@@ -34,3 +36,30 @@ class Manager(Table, help_text=help_text):
3436
pass
3537

3638
self.assertEqual(Manager._meta.help_text, help_text)
39+
40+
def test_foreign_key_columns(self):
41+
"""
42+
Make sure TableMeta.foreign_keys and TableMeta.foreign_key_references
43+
are setup correctly.
44+
"""
45+
46+
class TableA(Table):
47+
pass
48+
49+
class TableB(Table):
50+
table_a = ForeignKey(references=TableA)
51+
52+
self.assertEqual(TableB._meta.foreign_key_columns, [TableB.table_a])
53+
self.assertEqual(TableA._meta.foreign_key_references, [TableB.table_a])
54+
55+
def test_secret_columns(self):
56+
"""
57+
Make sure TableMeta.secret_columns are setup correctly.
58+
"""
59+
60+
class Classified(Table):
61+
top_secret = Secret()
62+
63+
self.assertEqual(
64+
Classified._meta.secret_columns, [Classified.top_secret]
65+
)

0 commit comments

Comments
 (0)