Merge remote-tracking branch 'origin/master'

1 files changed, 55 insertions, 4 deletions

diff --git a/pysite/route_manager.py b/pysite/route_manager.py
index 5fe22ba0..09202d0c 100644
--- a/pysite/route_manager.py
+++ b/pysite/route_manager.py
@@ -3,20 +3,35 @@ import importlib
 import inspect
 import os
 
-from flask import Flask
+from flask import Flask, abort, g
 
 from pysite.base_route import BaseView, ErrorView
+import rethinkdb
+
+from pysite.base_route import BaseView, ErrorView, RouteView
+
+DB_HOST = os.environ.get("RETHINKDB_HOST")
+DB_PORT = os.environ.get("RETHINKDB_PORT")
+DB_DATABASE = os.environ.get("RETHINKDB_DATABASE")
+DB_TABLE = os.environ.get("RETHINKDB_TABLE")
+
+TEMPLATES_PATH = "../templates"
+STATIC_PATH = "../static"
 
 
 class RouteManager:
     def __init__(self):
-        self.app = Flask(__name__)
+        self.app = Flask(
+            __name__, template_folder=TEMPLATES_PATH, static_folder=STATIC_PATH, static_url_path="/static",
+        )
         self.app.secret_key = os.environ.get("WEBPAGE_SECRET_KEY")
 
         self.load_views()
 
     def run(self):
-        self.app.run(port=int(os.environ.get("WEBPAGE_PORT")), debug=False)
+        self.app.run(
+            port=int(os.environ.get("WEBPAGE_PORT")), debug="FLASK_DEBUG" in os.environ
+        )
 
     def load_views(self, location="pysite/views"):
         for filename in os.listdir(location):
@@ -33,7 +48,43 @@ class RouteManager:
                             inspect.isclass(cls) and
                             cls is not BaseView and
                             cls is not ErrorView and
-                            (BaseView in cls.__mro__ or ErrorView in cls.__mro__)
+                            cls is not RouteView and
+                            BaseView in cls.__mro__
                     ):
                         cls.setup(self.app)
                         print(f"View loaded: {cls.name: <25} ({module.__name__}.{cls_name})")
+
+    def setup_db(self):
+        connection = self.get_db_connection(connect_database=False)
+
+        try:
+            rethinkdb.db_create(DB_DATABASE).run(connection)
+            rethinkdb.db(DB_DATABASE).table_create(DB_TABLE).run(connection)
+            print("Database created")
+        except rethinkdb.RqlRuntimeError:
+            print("Database found")
+        finally:
+            connection.close()
+
+        self.app.before_request(self.db_before_request)
+        self.app.teardown_request(self.db_teardown_request)
+
+    def get_db_connection(self, connect_database=True):
+        if connect_database:
+            return rethinkdb.connect(host=DB_HOST, port=DB_PORT, db=DB_DATABASE)
+        else:
+            return rethinkdb.connect(host=DB_HOST, port=DB_PORT)
+
+    def db_before_request(self):
+        try:
+            # g is the Flask global context object
+            g.rdb_conn = self.get_db_connection()
+        except rethinkdb.RqlDriverError:
+            abort(503, "Database connection could be established.")
+
+    def db_teardown_request(self, _):
+        try:
+            # g is the Flask global context object
+            g.rdb_conn.close()
+        except AttributeError:
+            pass