Merge pull request piccolo-orm#97 from piccolo-orm/user_model_enhancements

dantownsend · web-flow · commit 8a486c36956a · 2021-05-07T14:02:08.000+01:00
User model enhancements
diff --git a/piccolo/apps/migrations/commands/check.py b/piccolo/apps/migrations/commands/check.py
@@ -1,23 +1,28 @@
+import dataclasses
+import typing as t
+
 from piccolo.apps.migrations.commands.base import BaseMigrationManager
 from piccolo.apps.migrations.tables import Migration
 from piccolo.utils.printing import get_fixed_length_string
 
 
+@dataclasses.dataclass
+class MigrationStatus:
+    app_name: str
+    migration_id: str
+    has_ran: bool
+
+
 class CheckMigrationManager(BaseMigrationManager):
     def __init__(self, app_name: str):
         self.app_name = app_name
         super().__init__()
 
-    async def run(self):
-        print("Listing migrations ...")
-
+    async def get_migration_statuses(self) -> t.List[MigrationStatus]:
         # Make sure the migration table exists, otherwise we'll get an error.
         await self.create_migration_table()
 
-        print(
-            f'{get_fixed_length_string("APP NAME")} | '
-            f'{get_fixed_length_string("MIGRATION_ID")} | RAN'
-        )
+        migration_statuses: t.List[MigrationStatus] = []
 
         app_modules = self.get_app_modules()
 
@@ -29,8 +34,6 @@ async def run(self):
             if (self.app_name != "all") and (self.app_name != app_name):
                 continue
 
-            fixed_length_app_name = get_fixed_length_string(app_name)
-
             migration_modules = self.get_migration_modules(
                 app_config.migrations_folder_path
             )
@@ -44,11 +47,53 @@ async def run(self):
                     )
                     .run()
                 )
-                fixed_length_id = get_fixed_length_string(_id)
-                print(
-                    f"{fixed_length_app_name} | {fixed_length_id} | {has_ran}"
+                migration_statuses.append(
+                    MigrationStatus(
+                        app_name=app_name, migration_id=_id, has_ran=has_ran
+                    )
                 )
 
+        return migration_statuses
+
+    async def have_ran_count(self) -> int:
+        """
+        :returns:
+            The number of migrations which have been ran.
+        """
+        migration_statuses = await self.get_migration_statuses()
+        return len([i for i in migration_statuses if i.has_ran])
+
+    async def havent_ran_count(self) -> int:
+        """
+        :returns:
+            The number of migrations which haven't been ran.
+        """
+        migration_statuses = await self.get_migration_statuses()
+        return len([i for i in migration_statuses if not i.has_ran])
+
+    async def run(self):
+        """
+        Prints out the migrations which have and haven't ran.
+        """
+        print("Listing migrations ...")
+
+        print(
+            f'{get_fixed_length_string("APP NAME")} | '
+            f'{get_fixed_length_string("MIGRATION_ID")} | RAN'
+        )
+
+        migration_statuses = await self.get_migration_statuses()
+
+        for migration_status in migration_statuses:
+            fixed_length_app_name = get_fixed_length_string(
+                migration_status.app_name
+            )
+            fixed_length_id = get_fixed_length_string(
+                migration_status.migration_id
+            )
+            has_ran = migration_status.has_ran
+            print(f"{fixed_length_app_name} | {fixed_length_id} | {has_ran}")
+
 
 async def check(app_name: str = "all"):
     """
diff --git a/piccolo/apps/user/commands/create.py b/piccolo/apps/user/commands/create.py
@@ -33,6 +33,17 @@ def get_is_admin() -> bool:
     return admin == "y"
 
 
+def get_is_superuser() -> bool:
+    while True:
+        superuser = input("Superuser? Enter y or n:\n")
+        if superuser in ("y", "n"):
+            break
+        else:
+            print("Unrecognised option")
+
+    return superuser == "y"
+
+
 def create():
     """
     Create a new user.
@@ -49,13 +60,15 @@ def create():
         sys.exit("The password is too short")
 
     is_admin = get_is_admin()
+    is_superuser = get_is_superuser()
 
     user = BaseUser(
         username=username,
         password=password,
         admin=is_admin,
         email=email,
         active=True,
+        superuser=is_superuser,
     )
     user.save().run_sync()
 
diff --git a/piccolo/apps/user/piccolo_migrations/2021-04-30T16-14-15.py b/piccolo/apps/user/piccolo_migrations/2021-04-30T16-14-15.py
@@ -0,0 +1,48 @@
+from piccolo.apps.migrations.auto import MigrationManager
+from piccolo.columns.column_types import Boolean
+from piccolo.columns.column_types import Timestamp
+from piccolo.columns.indexes import IndexMethod
+
+
+ID = "2021-04-30T16:14:15"
+VERSION = "0.18.2"
+
+
+async def forwards():
+    manager = MigrationManager(migration_id=ID, app_name="user")
+
+    manager.add_column(
+        table_class_name="BaseUser",
+        tablename="piccolo_user",
+        column_name="superuser",
+        column_class_name="Boolean",
+        column_class=Boolean,
+        params={
+            "default": False,
+            "null": False,
+            "primary": False,
+            "key": False,
+            "unique": False,
+            "index": False,
+            "index_method": IndexMethod.btree,
+        },
+    )
+
+    manager.add_column(
+        table_class_name="BaseUser",
+        tablename="piccolo_user",
+        column_name="last_login",
+        column_class_name="Timestamp",
+        column_class=Timestamp,
+        params={
+            "default": None,
+            "null": True,
+            "primary": False,
+            "key": False,
+            "unique": False,
+            "index": False,
+            "index_method": IndexMethod.btree,
+        },
+    )
+
+    return manager
diff --git a/piccolo/apps/user/tables.py b/piccolo/apps/user/tables.py
@@ -6,7 +6,7 @@
 import secrets
 import typing as t
 
-from piccolo.columns import Varchar, Boolean, Secret
+from piccolo.columns import Varchar, Boolean, Secret, Timestamp
 from piccolo.columns.readable import Readable
 from piccolo.table import Table
 from piccolo.utils.sync import run_sync
@@ -23,7 +23,19 @@ class BaseUser(Table, tablename="piccolo_user"):
     last_name = Varchar(null=True)
     email = Varchar(length=255, unique=True)
     active = Boolean(default=False)
-    admin = Boolean(default=False)
+    admin = Boolean(
+        default=False, help_text="An admin can log into the Piccolo admin GUI."
+    )
+    superuser = Boolean(
+        default=False,
+        help_text=(
+            "If True, this user can manage other users's passwords in the "
+            "Piccolo admin GUI."
+        ),
+    )
+    last_login = Timestamp(
+        null=True, default=None, help_text="When this user last logged in."
+    )
 
     def __init__(self, **kwargs):
         """
diff --git a/piccolo/engine/postgres.py b/piccolo/engine/postgres.py
@@ -14,7 +14,7 @@
 from piccolo.query.base import Query
 from piccolo.querystring import QueryString
 from piccolo.utils.sync import run_sync
-from piccolo.utils.warnings import colored_warning
+from piccolo.utils.warnings import colored_warning, colored_string, Level
 
 
 @dataclass
@@ -284,11 +284,14 @@ async def prep_database(self):
                 )
             except InsufficientPrivilegeError:
                 print(
-                    f"Unable to create {extension} extension - some "
-                    "functionality may not behave as expected. Make sure your "
-                    "database user has permission to create extensions, or "
-                    "add it manually using "
-                    f'`CREATE EXTENSION "{extension}";`'
+                    colored_string(
+                        f"=> Unable to create {extension} extension - some "
+                        "functionality may not behave as expected. Make sure "
+                        "your database user has permission to create "
+                        "extensions, or add it manually using "
+                        f'`CREATE EXTENSION "{extension}";`',
+                        level=Level.medium,
+                    )
                 )
 
     ###########################################################################
diff --git a/piccolo/main.py b/piccolo/main.py
@@ -3,16 +3,19 @@
 
 from targ import CLI  # type: ignore
 
-from piccolo.conf.apps import AppRegistry, Finder
 from piccolo.apps.app.piccolo_app import APP_CONFIG as app_config
 from piccolo.apps.asgi.piccolo_app import APP_CONFIG as asgi_config
 from piccolo.apps.meta.piccolo_app import APP_CONFIG as meta_config
+from piccolo.apps.migrations.commands.check import CheckMigrationManager
 from piccolo.apps.migrations.piccolo_app import APP_CONFIG as migrations_config
 from piccolo.apps.playground.piccolo_app import APP_CONFIG as playground_config
 from piccolo.apps.project.piccolo_app import APP_CONFIG as project_config
 from piccolo.apps.shell.piccolo_app import APP_CONFIG as shell_config
 from piccolo.apps.sql_shell.piccolo_app import APP_CONFIG as sql_shell_config
 from piccolo.apps.user.piccolo_app import APP_CONFIG as user_config
+from piccolo.conf.apps import AppRegistry, Finder
+from piccolo.utils.sync import run_sync
+from piccolo.utils.warnings import colored_string, Level
 
 
 DIAGNOSE_FLAG = "--diagnose"
@@ -81,6 +84,37 @@ def main():
                     continue
                 cli.register(command, group_name=app_name)
 
+        if "migrations" not in sys.argv:
+            # Show a warning if any migrations haven't been run.
+            # Don't run it if it looks like the user is running a migration
+            # command, as this information is redundant.
+
+            try:
+                havent_ran_count = run_sync(
+                    CheckMigrationManager(app_name="all").havent_ran_count()
+                )
+                if havent_ran_count:
+                    message = (
+                        f"{havent_ran_count} migration hasn't"
+                        if havent_ran_count == 1
+                        else f"{havent_ran_count} migrations haven't"
+                    )
+                    print(
+                        colored_string(
+                            message=(
+                                "=> {} been run - the app "
+                                "might not behave as expected.\n"
+                                "To check which use:\n"
+                                "    piccolo migrations check\n"
+                                "To run all migrations:\n"
+                                "    piccolo migrations forwards all\n"
+                            ).format(message),
+                            level=Level.high,
+                        )
+                    )
+            except Exception:
+                pass
+
     ###########################################################################
 
     cli.run()
diff --git a/piccolo/utils/warnings.py b/piccolo/utils/warnings.py
@@ -15,6 +15,10 @@ class Level(Enum):
     high = colorama.Fore.RED
 
 
+def colored_string(message: str, level: Level = Level.medium) -> str:
+    return level.value + message + colorama.Fore.RESET
+
+
 def colored_warning(
     message: str,
     category: t.Type[Warning] = Warning,
@@ -36,5 +40,5 @@ def colored_warning(
     :level:
         Used to determine the colour of the text displayed to the user.
     """
-    colored_message = level.value + message + colorama.Fore.RESET
+    colored_message = colored_string(message=message, level=level)
     warnings.warn(colored_message, category=category, stacklevel=stacklevel)
diff --git a/piccolo_conf.py b/piccolo_conf.py
@@ -0,0 +1,20 @@
+"""
+This piccolo_conf file is just here so migrations can be made for Piccolo's own
+internal apps.
+
+For example:
+
+python -m piccolo.main migration new user --auto
+
+"""
+
+from piccolo.conf.apps import AppRegistry
+from piccolo.engine.postgres import PostgresEngine
+
+
+DB = PostgresEngine(config={})
+
+
+# A list of paths to piccolo apps
+# e.g. ['blog.piccolo_app']
+APP_REGISTRY = AppRegistry(apps=["piccolo.apps.user.piccolo_app"])
diff --git a/tests/apps/migrations/test_migration.py b/tests/apps/migrations/test_migration.py
@@ -5,6 +5,6 @@
 
 class TestMigrationTable(TestCase):
     def test_migration_table(self):
-        Migration.create_table().run_sync()
+        Migration.create_table(if_not_exists=True).run_sync()
         Migration.select().run_sync()
         Migration.alter().drop_table().run_sync()
diff --git a/tests/apps/user/commands/test_create.py b/tests/apps/user/commands/test_create.py
@@ -31,6 +31,10 @@ def tearDown(self):
     @patch(
         "piccolo.apps.user.commands.create.get_is_admin", return_value=True,
     )
+    @patch(
+        "piccolo.apps.user.commands.create.get_is_superuser",
+        return_value=True,
+    )
     def test_create(self, *args, **kwargs):
         create()
 
@@ -40,6 +44,7 @@ def test_create(self, *args, **kwargs):
                 (BaseUser.admin == True)  # noqa: E712
                 & (BaseUser.username == "bob123")
                 & (BaseUser.email == "bob@test.com")
+                & (BaseUser.superuser == True)
             )
             .run_sync()
         )

Original file line number	Diff line number	Diff line change
`@@ -31,6 +31,10 @@ def tearDown(self):`
`31`	`31`	`@patch(`
`32`	`32`	`"piccolo.apps.user.commands.create.get_is_admin", return_value=True,`
`33`	`33`	`)`
	`34`	`+ @patch(`
	`35`	`+ "piccolo.apps.user.commands.create.get_is_superuser",`
	`36`	`+ return_value=True,`
	`37`	`+ )`
`34`	`38`	`def test_create(self, args, *kwargs):`
`35`	`39`	`create()`
`36`	`40`
`@@ -40,6 +44,7 @@ def test_create(self, args, *kwargs):`
`40`	`44`	`(BaseUser.admin == True) # noqa: E712`
`41`	`45`	`& (BaseUser.username == "bob123")`
`42`	`46`	`& (BaseUser.email == "[email protected]")`
	`47`	`+ & (BaseUser.superuser == True)`
`43`	`48`	`)`
`44`	`49`	`.run_sync()`
`45`	`50`	`)`