From f5d448db24fae0ffdcb086ec32ef07c2b37f0ec8 Mon Sep 17 00:00:00 2001
From: d42 <d42@users.noreply.github.com>
Date: Tue, 25 Feb 2020 02:08:42 +0100
Subject: [PATCH] flask shell environment

---
 website/app.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/website/app.py b/website/app.py
index 7d7f6bd..801e94d 100644
--- a/website/app.py
+++ b/website/app.py
@@ -4,7 +4,7 @@ from flask import Flask
 from flask_login import LoginManager
 from flask_migrate import Migrate
 
-from .models import db
+from . import models
 from .routes import bp
 from .ldap import LDAPUser
 from .oauth2 import config_oauth
@@ -18,17 +18,27 @@ def create_app():
     setup_app(app)
     return app
 
+def register_shellcontext(app):
+    """Register shell context objects."""
+
+    def shell_context():
+        """Shell context objects."""
+        return {"db": models.db, "models": models}
+
+    app.shell_context_processor(shell_context)
+
 
 def setup_app(app):
     logging.basicConfig(
         level=logging.DEBUG,
         format="[%(asctime)-15s] %(name)-10s %(levelname)7s: %(message)s",
     )
-    db.init_app(app)
-    migrate = Migrate(app, db)
+    models.db.init_app(app)
+    migrate = Migrate(app, models.db)
     login_manager = LoginManager()
     login_manager.init_app(app)
     login_manager.login_view = "/login"
     login_manager.user_loader(LDAPUser.by_login)
     config_oauth(app)
     app.register_blueprint(bp, url_prefix='')
+    register_shellcontext(app)