diff options
| -rw-r--r-- | gunicorn_config.py | 12 | ||||
| -rw-r--r-- | pysite/database.py | 27 | ||||
| -rw-r--r-- | pysite/mixins.py | 5 |
3 files changed, 41 insertions, 3 deletions
diff --git a/gunicorn_config.py b/gunicorn_config.py new file mode 100644 index 00000000..7e17b0cc --- /dev/null +++ b/gunicorn_config.py @@ -0,0 +1,12 @@ +# coding=utf-8 +def when_ready(_server): + """ server hook that only runs when the gunicorn master process loads """ + + print("Creating tables...") + + from pysite.database import RethinkDB + + db = RethinkDB(loop_type=None) + created = db.create_tables() + + print(f"Created {created} tables.") diff --git a/pysite/database.py b/pysite/database.py index 4c2153fe..f53b9e60 100644 --- a/pysite/database.py +++ b/pysite/database.py @@ -9,16 +9,27 @@ from rethinkdb.ast import RqlMethodQuery, Table, UserError from rethinkdb.net import DefaultConnection +ALL_TABLES = { + # table: primary_key + + "oauth_data": "id", + "tags": "tag_name", + "users": "user_id", + "wiki": "slug", +} + + class RethinkDB: - def __init__(self, loop_type: str = "gevent"): + def __init__(self, loop_type: Optional[str] = "gevent"): self.host = os.environ.get("RETHINKDB_HOST", "127.0.0.1") self.port = os.environ.get("RETHINKDB_PORT", "28016") self.database = os.environ.get("RETHINKDB_DATABASE", "pythondiscord") self.log = logging.getLogger() self.conn = None - rethinkdb.set_loop_type(loop_type) + if loop_type: + rethinkdb.set_loop_type(loop_type) with self.get_connection(connect_database=False) as conn: try: @@ -27,6 +38,15 @@ class RethinkDB: except rethinkdb.RqlRuntimeError: self.log.debug(f"Database found: '{self.database}'") + def create_tables(self) -> int: + created = 0 + + for table, primary_key in ALL_TABLES.values(): + if self.create_table(table, primary_key): + created += 1 + + return created + def get_connection(self, connect_database: bool=True) -> DefaultConnection: """ Grab a connection to the RethinkDB server, optionally without selecting a database @@ -170,6 +190,9 @@ class RethinkDB: :return: The RethinkDB table object for the table """ + if table_name not in ALL_TABLES: + self.log.warning(f"Table not declared in database.py: {table_name}") + return rethinkdb.table(table_name) def run(self, query: RqlMethodQuery, *, new_connection: bool=False, diff --git a/pysite/mixins.py b/pysite/mixins.py index a2730ae4..c57ca85f 100644 --- a/pysite/mixins.py +++ b/pysite/mixins.py @@ -1,4 +1,5 @@ # coding=utf-8 +import os from weakref import ref from flask import Blueprint @@ -51,7 +52,9 @@ class DBMixin: raise RuntimeError("Routes using DBViewMixin must define `table_name`") cls._db = ref(manager.db) - manager.db.create_table(cls.table_name, primary_key=cls.table_primary_key) + + if "FLASK_DEBUG" in os.environ: + manager.db.create_table(cls.table_name, primary_key=cls.table_primary_key) @property def table(self) -> Table: |