added 'Secret' column, and the ability to automatically remove them from select queries

Author: dantownsend

piccolo/columns/__init__.py

@@ -1,13 +1,14 @@
 from .column_types import (  # noqa
-    Varchar,
+    Boolean,
+    ForeignKey,
     Integer,
-    Serial,
     PrimaryKey,
-    Timestamp,
+    Secret,
+    Serial,
     Text,
-    Boolean,
-    ForeignKey,
+    Timestamp,
     UUID,
+    Varchar,
 )
 from .base import Column, ForeignKeyMeta, Selectable, OnDelete  # noqa
 from .combination import And, Or, Where  # noqa

piccolo/columns/column_types.py

@@ -33,6 +33,17 @@ def column_type(self):
             return "VARCHAR"
 
 
+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.
+    """
+
+    pass
+
+
 class Text(Column):
     """
     Used for text when you don't want any character length limits.

piccolo/extensions/user/tables.py

@@ -8,7 +8,7 @@
 from asgiref.sync import async_to_sync
 
 from piccolo.table import Table
-from piccolo.columns import Varchar, Boolean
+from piccolo.columns import Varchar, Boolean, Secret
 from piccolo.columns.readable import Readable
 
 
@@ -19,7 +19,7 @@ class BaseUser(Table, tablename="piccolo_user"):
     """
 
     username = Varchar(length=100, unique=True)
-    password = Varchar(length=255)
+    password = Secret(length=255)
     email = Varchar(length=255, unique=True)
     active = Boolean(default=False)
     admin = Boolean(default=False)
@@ -61,7 +61,9 @@ async def update_password(cls, user: t.Union[str, int], password: str):
         elif type(user) == int:
             clause = cls.id == user  # type: ignore
         else:
-            raise ValueError("The `user` arg must be a user id, or a username.")
+            raise ValueError(
+                "The `user` arg must be a user id, or a username."
+            )
 
         password = cls.hash_password(password)
         await cls.update().values({cls.password: password}).where(clause).run()

piccolo/query/methods/select.py

@@ -26,6 +26,7 @@ class Select(Query):
 
     __slots__ = (
         "columns_list",
+        "exclude_secrets",
         "columns_delegate",
         "distinct_delegate",
         "limit_delegate",
@@ -36,9 +37,13 @@ class Select(Query):
     )
 
     def __init__(
-        self, table: t.Type[Table], columns_list: t.Iterable[Selectable] = []
+        self,
+        table: t.Type[Table],
+        columns_list: t.Sequence[Selectable] = [],
+        exclude_secrets: bool = False,
     ):
         super().__init__(table)
+        self.exclude_secrets = exclude_secrets
 
         self.columns_delegate = ColumnsDelegate()
         self.distinct_delegate = DistinctDelegate()
@@ -187,6 +192,10 @@ def querystrings(self) -> t.Sequence[QueryString]:
         if len(self.columns_delegate.selected_columns) == 0:
             self.columns_delegate.selected_columns = self.table._meta.columns
 
+        # If secret fields need to be omitted, remove them from the list.
+        if self.exclude_secrets:
+            self.columns_delegate.remove_secret_columns()
+
         engine_type = self.table._meta.db.engine_type
 
         select_strings: t.List[str] = [

piccolo/query/mixins.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass, field
 import typing as t
 
-from piccolo.columns import And, Column, Where, Or
+from piccolo.columns import And, Column, Secret, Where, Or
 from piccolo.custom_types import Combinable
 from piccolo.querystring import QueryString
 
@@ -195,6 +195,11 @@ class ColumnsDelegate:
     def columns(self, *columns: Column):
         self.selected_columns += columns
 
+    def remove_secret_columns(self):
+        self.selected_columns = [
+            i for i in self.selected_columns if not isinstance(i, Secret)
+        ]
+
 
 @dataclass
 class ValuesDelegate:

piccolo/table.py

@@ -395,7 +395,9 @@ def _process_column_args(
         ]
 
     @classmethod
-    def select(cls, *columns: t.Union[Selectable, str]) -> Select:
+    def select(
+        cls, *columns: t.Union[Selectable, str], exclude_secrets=False
+    ) -> Select:
         """
         Get data in the form of a list of dictionaries, with each dictionary
         representing a row.
@@ -405,9 +407,15 @@ def select(cls, *columns: t.Union[Selectable, str]) -> Select:
         await Band.select().columns(Band.name).run()
         await Band.select(Band.name).run()
         await Band.select('name').run()
+
+        :param exclude_secrets: If True, any password fields are omitted from
+        the response. Even though passwords are hashed, you still don't want
+        them being passed over the network if avoidable.
         """
         columns = cls._process_column_args(*columns)
-        return Select(table=cls, columns_list=columns)
+        return Select(
+            table=cls, columns_list=columns, exclude_secrets=exclude_secrets
+        )
 
     @classmethod
     def delete(cls, force=False) -> Delete:

tests/table/test_select.py

@@ -1,3 +1,7 @@
+from unittest import TestCase
+
+from piccolo.extensions.user.tables import BaseUser
+
 from ..base import DBTestCase, postgres_only, sqlite_only
 from ..example_project.tables import Band, Concert
 
@@ -333,3 +337,22 @@ def test_call_chain(self):
         """
         # self.assertEqual(len(Concert.band_1.name._meta.call_chain), 1)
         self.assertEqual(len(Concert.band_1.manager.name._meta.call_chain), 2)
+
+
+class TestSelectSecret(TestCase):
+    def setUp(self):
+        BaseUser.create_table().run_sync()
+
+    def tearDown(self):
+        BaseUser.alter().drop_table().run_sync()
+
+    def test_secret(self):
+        """
+        Make sure that secret fields are omitted from the response when
+        requested.
+        """
+        user = BaseUser(username="piccolo", password="piccolo123")
+        user.save().run_sync()
+
+        user_dict = BaseUser.select(exclude_secrets=True).first().run_sync()
+        self.assertTrue("password" not in user_dict.keys())