diff options
author | 2018-04-07 10:10:47 +0100 | |
---|---|---|
committer | 2018-04-07 10:10:47 +0100 | |
commit | 2db3ce6c2e91b81a64916eba47ae34f4d9182702 (patch) | |
tree | 26e22969e56caae7d18d03be3606bb9f3fbdce4a /pysite/database.py | |
parent | Flake8 (diff) |
Don't create tables on route load by default.
* If "FLASK_DEBUG" is in your env vars, tables will be created on route load
* If you run a query against a table not declared in database.py, a warning is emitted
Diffstat (limited to 'pysite/database.py')
-rw-r--r-- | pysite/database.py | 27 |
1 files changed, 25 insertions, 2 deletions
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, |