diff --git a/piccolo/columns/base.py b/piccolo/columns/base.py index 8d3fe444e..c85961d38 100644 --- 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 index 28e880988..edaeccb26 100644 --- 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 index ced0b0e77..338dc5620 100644 --- 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 index a0b0d93e2..35e4a6a1f 100644 --- 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 index 8b4c46d9f..1b813ba66 100644 --- 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 index 701bf5916..3bfe7ebb6 100644 --- 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 index 091e1b764..cef974c0c 100644 --- 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 new file mode 100644 index 000000000..e7db0f3a9 --- /dev/null +++ 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 index 331fab081..774718cae 100644 --- 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 index 92ac15ab2..5d7d774e0 100644 --- 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 index f7f296f20..ac3861700 100644 --- 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 ), )