added column secret attribute (#347)

* added column secret attribute * tweak docs, and force a `Secret` column to have `_meta.secret=True` * fix linting error Co-authored-by: Daniel Townsend <[email protected]>
piccolo-orm · Nov 14, 2021 · 5135269 · 5135269
1 parent bffd3df
commit 5135269
Show file tree

Hide file tree

Showing 11 changed files with 88 additions and 32 deletions.
diff --git a/piccolo/columns/base.py b/piccolo/columns/base.py
@@ -127,6 +127,7 @@ class ColumnMeta:
     required: bool = False
     help_text: t.Optional[str] = None
     choices: t.Optional[t.Type[Enum]] = None
+    secret: bool = False
 
     # Used for representing the table in migrations and the playground.
     params: t.Dict[str, t.Any] = field(default_factory=dict)
@@ -338,6 +339,20 @@ class MyTable(Table):
         This is an advanced feature which you should only need in niche
         situations.
 
+    :param secret:
+        If ``secret=True`` is specified, it allows a user to automatically
+        omit any fields when doing a select query, to help prevent
+        inadvertent leakage of sensitive data.
+
+        .. code-block:: python
+
+            class Band(Table):
+                name = Varchar()
+                net_worth = Integer(secret=True)
+
+            >>> Property.select(exclude_secrets=True).run_sync()
+            [{'name': 'Pythonistas'}]
+
     """
 
     value_type: t.Type = int
@@ -353,6 +368,7 @@ def __init__(
         help_text: t.Optional[str] = None,
         choices: t.Optional[t.Type[Enum]] = None,
         db_column_name: t.Optional[str] = None,
+        secret: bool = False,
         **kwargs,
     ) -> None:
         # This is for backwards compatibility - originally there were two
@@ -375,6 +391,7 @@ def __init__(
                 "index_method": index_method,
                 "choices": choices,
                 "db_column_name": db_column_name,
+                "secret": secret,
             }
         )
 
@@ -398,6 +415,7 @@ def __init__(
             help_text=help_text,
             choices=choices,
             _db_column_name=db_column_name,
+            secret=secret,
         )
 
         self.alias: t.Optional[str] = None

diff --git a/piccolo/columns/column_types.py b/piccolo/columns/column_types.py
@@ -201,30 +201,13 @@ def __radd__(self, value: t.Union[str, Varchar, Text]) -> QueryString:
 
 class Secret(Varchar):
     """
-    The database treats it the same as a ``Varchar``, but Piccolo may treat it
-    differently internally - for example, allowing a user to automatically
-    omit any secret fields when doing a select query, to help prevent
-    inadvertant leakage. A common use for a ``Secret`` field is a password.
-
-    Uses the ``str`` type for values.
-
-    **Example**
-
-    .. code-block:: python
-
-        class Door(Table):
-            code = Secret(length=100)
-
-        # Create
-        >>> Door(code='123abc').save().run_sync()
-
-        # Query
-        >>> Door.select(Door.code).run_sync()
-        {'code': '123abc'}
-
+    This is just an alias to ``Varchar(secret=True)``. It's here for backwards
+    compatibility.
     """
 
-    pass
+    def __init__(self, *args, **kwargs):
+        kwargs["secret"] = True
+        super().__init__(*args, **kwargs)
 
 
 class Text(Column):

diff --git a/piccolo/query/mixins.py b/piccolo/query/mixins.py
@@ -3,7 +3,7 @@
 import typing as t
 from dataclasses import dataclass, field
 
-from piccolo.columns import And, Column, Or, Secret, Where
+from piccolo.columns import And, Column, Or, Where
 from piccolo.columns.column_types import ForeignKey
 from piccolo.custom_types import Combinable
 from piccolo.querystring import QueryString
@@ -288,7 +288,7 @@ def columns(self, *columns: t.Union[Selectable, t.List[Selectable]]):
 
     def remove_secret_columns(self):
         self.selected_columns = [
-            i for i in self.selected_columns if not isinstance(i, Secret)
+            i for i in self.selected_columns if not i._meta.secret
         ]
 
 

diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py
@@ -38,7 +38,7 @@ def test_add_table(self):
         self.assertTrue(len(new_table_columns.statements) == 1)
         self.assertEqual(
             new_table_columns.statements[0],
-            "manager.add_column(table_class_name='Band', tablename='band', column_name='name', db_column_name='name', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary_key': False, 'unique': False, 'index': False, 'index_method': IndexMethod.btree, 'choices': None, 'db_column_name': None})",  # noqa
+            "manager.add_column(table_class_name='Band', tablename='band', column_name='name', db_column_name='name', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary_key': False, 'unique': False, 'index': False, 'index_method': IndexMethod.btree, 'choices': None, 'db_column_name': None, 'secret': False})",  # noqa
         )
 
     def test_drop_table(self):
@@ -122,7 +122,7 @@ def test_add_column(self):
         self.assertTrue(len(schema_differ.add_columns.statements) == 1)
         self.assertEqual(
             schema_differ.add_columns.statements[0],
-            "manager.add_column(table_class_name='Band', tablename='band', column_name='genre', db_column_name='genre', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary_key': False, 'unique': False, 'index': False, 'index_method': IndexMethod.btree, 'choices': None, 'db_column_name': None})",  # noqa
+            "manager.add_column(table_class_name='Band', tablename='band', column_name='genre', db_column_name='genre', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary_key': False, 'unique': False, 'index': False, 'index_method': IndexMethod.btree, 'choices': None, 'db_column_name': None, 'secret': False})",  # noqa
         )
 
     def test_drop_column(self):

diff --git a/tests/apps/migrations/auto/test_serialisation.py b/tests/apps/migrations/auto/test_serialisation.py
@@ -247,7 +247,7 @@ def test_lazy_table_reference(self):
                     'class Manager(Table, tablename="manager"): '
                     "id = Serial(null=False, primary_key=True, unique=False, "
                     "index=False, index_method=IndexMethod.btree, "
-                    "choices=None, db_column_name='id')"
+                    "choices=None, db_column_name='id', secret=False)"
                 ),
             )
 
@@ -298,7 +298,7 @@ def test_column_instance(self):
 
         self.assertEqual(
             serialised.params["base_column"].__repr__(),
-            "Varchar(length=255, default='', null=False, primary_key=False, unique=False, index=False, index_method=IndexMethod.btree, choices=None, db_column_name=None)",  # noqa: E501
+            "Varchar(length=255, default='', null=False, primary_key=False, unique=False, index=False, index_method=IndexMethod.btree, choices=None, db_column_name=None, secret=False)",  # noqa: E501
         )
 
         self.assertEqual(

diff --git a/tests/apps/migrations/commands/test_forwards_backwards.py b/tests/apps/migrations/commands/test_forwards_backwards.py
@@ -176,6 +176,7 @@ def test_forwards_fake(self):
                 "2020-12-17T18:44:44",
                 "2021-07-25T22:38:48:009306",
                 "2021-09-06T13:58:23:024723",
+                "2021-11-13T14:01:46:114725",
             ],
         )
 

diff --git a/tests/columns/test_base.py b/tests/columns/test_base.py
@@ -34,6 +34,16 @@ def test_help_text(self):
         self.assertTrue(column._meta.help_text == help_text)
 
 
+class TestSecretParameter(TestCase):
+    def test_secret_parameter(self):
+        """
+        Test adding secret parameter to a column.
+        """
+        secret = False
+        column = Varchar(secret=secret)
+        self.assertTrue(column._meta.secret == secret)
+
+
 class TestChoices(TestCase):
     def test_choices(self):
         """

diff --git a/tests/example_apps/music/piccolo_migrations/2021-11-13T14-01-46-114725.py b/tests/example_apps/music/piccolo_migrations/2021-11-13T14-01-46-114725.py
@@ -0,0 +1,24 @@
+from piccolo.apps.migrations.auto import MigrationManager
+from piccolo.columns.column_types import Integer
+
+ID = "2021-11-13T14:01:46:114725"
+VERSION = "0.59.0"
+DESCRIPTION = ""
+
+
+async def forwards():
+    manager = MigrationManager(
+        migration_id=ID, app_name="music", description=DESCRIPTION
+    )
+
+    manager.alter_column(
+        table_class_name="Venue",
+        tablename="venue",
+        column_name="capacity",
+        params={"secret": True},
+        old_params={"secret": False},
+        column_class=Integer,
+        old_column_class=Integer,
+    )
+
+    return manager
diff --git a/tests/example_apps/music/tables.py b/tests/example_apps/music/tables.py
@@ -36,7 +36,7 @@ class Band(Table):
 
 class Venue(Table):
     name = Varchar(length=100)
-    capacity = Integer(default=0)
+    capacity = Integer(default=0, secret=True)
 
 
 class Concert(Table):

diff --git a/tests/table/test_select.py b/tests/table/test_select.py
@@ -3,7 +3,7 @@
 from piccolo.apps.user.tables import BaseUser
 from piccolo.columns.combination import WhereRaw
 from piccolo.query.methods.select import Avg, Count, Max, Min, Sum
-from tests.example_apps.music.tables import Band, Concert, Manager
+from tests.example_apps.music.tables import Band, Concert, Manager, Venue
 
 from ..base import DBTestCase, postgres_only, sqlite_only
 
@@ -965,3 +965,23 @@ def test_secret(self):
 
         user_dict = BaseUser.select(exclude_secrets=True).first().run_sync()
         self.assertTrue("password" not in user_dict.keys())
+
+
+class TestSelectSecretParameter(TestCase):
+    def setUp(self):
+        Venue.create_table().run_sync()
+
+    def tearDown(self):
+        Venue.alter().drop_table().run_sync()
+
+    def test_secret_parameter(self):
+        """
+        Make sure that fields with parameter ``secret=True`` are omitted
+        from the response when requested.
+        """
+        venue = Venue(name="The Garage", capacity=1000)
+        venue.save().run_sync()
+
+        venue_dict = Venue.select(exclude_secrets=True).first().run_sync()
+        self.assertTrue(venue_dict, {"id": 1, "name": "The Garage"})
+        self.assertTrue("capacity" not in venue_dict.keys())
diff --git a/tests/table/test_str.py b/tests/table/test_str.py
@@ -9,8 +9,8 @@ def test_str(self):
             Manager._table_str(),
             (
                 "class Manager(Table, tablename='manager'):\n"
-                "    id = Serial(null=False, primary_key=True, unique=False, index=False, index_method=IndexMethod.btree, choices=None, db_column_name='id')\n"  # noqa: E501
-                "    name = Varchar(length=50, default='', null=False, primary_key=False, unique=False, index=False, index_method=IndexMethod.btree, choices=None, db_column_name=None)\n"  # noqa: E501
+                "    id = Serial(null=False, primary_key=True, unique=False, index=False, index_method=IndexMethod.btree, choices=None, db_column_name='id', secret=False)\n"  # noqa: E501
+                "    name = Varchar(length=50, default='', null=False, primary_key=False, unique=False, index=False, index_method=IndexMethod.btree, choices=None, db_column_name=None, secret=False)\n"  # noqa: E501
             ),
         )