From 2db3ce6c2e91b81a64916eba47ae34f4d9182702 Mon Sep 17 00:00:00 2001 From: Gareth Coles Date: Sat, 7 Apr 2018 10:10:47 +0100 Subject: 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 --- pysite/database.py | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) (limited to 'pysite/database.py') 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, -- cgit v1.2.3