diff options
Diffstat (limited to 'pysite')
| -rw-r--r-- | pysite/route_manager.py | 43 |
1 files changed, 42 insertions, 1 deletions
diff --git a/pysite/route_manager.py b/pysite/route_manager.py index f922be6a..813816b7 100644 --- a/pysite/route_manager.py +++ b/pysite/route_manager.py @@ -2,13 +2,19 @@ import importlib import inspect import os +import rethinkdb -from flask import Flask +from flask import Flask, g, abort from pysite.base_route import BaseView, ErrorView, RouteView __author__ = "Gareth Coles" +DB_HOST = os.environ["RETHINKDB_HOST"] +DB_PORT = os.environ["RETHINKDB_PORT"] +DB_DATABASE = os.environ["RETHINKDB_DATABASE"] +DB_TABLE = os.environ["RETHINKDB_TABLE"] + class RouteManager: def __init__(self): @@ -40,3 +46,38 @@ class RouteManager: ): 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 |