-
Notifications
You must be signed in to change notification settings - Fork 0
/
db.py
77 lines (62 loc) · 2.44 KB
/
db.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.ext.declarative import declarative_base
import click
import os
import logging
ENV = os.getenv('ENV')
if ENV == 'DEV':
import config_dev as config
else:
import config_prod as config
engine = create_engine(
'sqlite:///' + os.path.join(config._DATABASE_FOLDER, config._DATABASE_NAME))
db_session = scoped_session(sessionmaker(autocommit=False,
autoflush=False,
bind=engine))
Base = declarative_base()
Base.query = db_session.query_property()
def init_db():
import app.models
Base.metadata.create_all(bind=engine)
# Check if the database already has content
# Create roles/scopes
admin_role = app.models.Role('admin')
wizard_role = app.models.Role('wizard')
db_session.add(admin_role, wizard_role)
# Create first admin, and a normal user
marduk_password = 'universe' if ENV == 'DEV' else config.ADMIN_PROD_PASSWORD
marduk = app.models.User('marduk', marduk_password,
name='Marduk', surname='Babylonian', roles=[admin_role, wizard_role])
db_session.add(marduk)
# Create pregenerated client
# First the grants
basic_grant = app.models.OAuth2Grant('password')
refresh_token_grant = app.models.OAuth2Grant('refresh_token')
db_session.add(basic_grant)
db_session.add(refresh_token_grant)
pregen_client = app.models.OAuth2Client(config.OAUTH2_PREGEN_CLIENT_ID, [
admin_role, wizard_role], [basic_grant, refresh_token_grant])
db_session.add(pregen_client)
if ENV == 'DEV':
human = app.models.User('human', 'earth',
name='Human', surname='Earthling', roles=[wizard_role])
db_session.add(human)
db_session.commit()
def close_db(exception=None):
db_session.remove()
@click.command('init-db')
def init_db_command():
"""Create database from scratch"""
init_db()
click.echo('Created the database')
def init_app(app):
app.teardown_appcontext(close_db)
app.cli.add_command(init_db_command)
if not engine.has_table('user'):
logging.info("Database has not been initialized")
logging.info("Creating database...")
init_db()
logging.info('Created the database')
else:
logging.info('Database exists already, not regenerating')