aboutsummaryrefslogtreecommitdiffstats
path: root/pysite/database.py
diff options
context:
space:
mode:
authorGravatar Gareth Coles <[email protected]>2018-04-07 10:10:47 +0100
committerGravatar Gareth Coles <[email protected]>2018-04-07 10:10:47 +0100
commit2db3ce6c2e91b81a64916eba47ae34f4d9182702 (patch)
tree26e22969e56caae7d18d03be3606bb9f3fbdce4a /pysite/database.py
parentFlake8 (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.py27
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,