Improve create_db.py

create_db.py

@@ -8,47 +8,71 @@ import os.path
 import time
 import sys
 
-# Create schema
-if os.environ.get('WAITFOR_DB') is None:
-    db.create_all()
-else:
-    for i in range(0, int(os.environ.get('WAITFOR_DB'))):
-        try:
-            db.create_all()
-            break
-        except:
-            print("INFO: Wait for database server")
-            sys.stdout.flush()
-            time.sleep(1)
-            continue
+def start():
+    wait_time = get_waittime_from_env()
+
+    if not connect_db(wait_time):
         print("ERROR: Couldn't connect to database server")
         exit(1)
 
-# create initial user roles and turn off maintenance mode
-admin_role = Role('Administrator', 'Administrator')
-user_role = Role('User', 'User')
-maintenance_setting = Setting('maintenance', 'False')
-fullscreen_layout_setting = Setting('fullscreen_layout', 'True')
-record_helper_setting = Setting('record_helper', 'True')
-default_table_size_setting = Setting('default_record_table_size', '15')
+    init_records()
 
-# Check if record already exists
-if not db.session.query(Role).filter_by(name="Administrator").first():
-    db.session.add(admin_role)
-if not db.session.query(Role).filter_by(name="User").first():
-    db.session.add(user_role)
-if not db.session.query(Setting).filter_by(name="maintenance").first():
-    db.session.add(maintenance_setting)
-if not db.session.query(Setting).filter_by(name="fullscreen_layout").first():
-    db.session.add(fullscreen_layout_setting)
-if not db.session.query(Setting).filter_by(name="record_helper").first():
-    db.session.add(record_helper_setting)
-if not db.session.query(Setting).filter_by(name="default_record_table_size").first():
-    db.session.add(default_table_size_setting)
-db_commit = db.session.commit()
+def get_waittime_from_env():
+    return int(os.environ.get('WAITFOR_DB', 1))
 
-if not os.path.exists(SQLALCHEMY_MIGRATE_REPO):
-    api.create(SQLALCHEMY_MIGRATE_REPO, 'database repository')
-    api.version_control(SQLALCHEMY_DATABASE_URI, SQLALCHEMY_MIGRATE_REPO)
-elif db_commit is not None:
-    api.version_control(SQLALCHEMY_DATABASE_URI, SQLALCHEMY_MIGRATE_REPO, api.version(SQLALCHEMY_MIGRATE_REPO))
+def connect_db(wait_time):
+    for i in xrange(0, wait_time):
+        print("INFO: Wait for database server")
+        sys.stdout.flush()
+        try:
+            db.create_all()
+            return True
+        except:
+            time.sleep(1)
+
+    return False
+
+def init_roles(db, role_names):
+    name_of_roles = map(lambda r: r.name, role_names)
+    rows = db.session.query(Role).filter(Role.name.in_(name_of_roles)).all()
+    name_of_rows = map(lambda r: r.name, rows)
+    roles = filter(lambda r: r.name not in rows, role_names)
+
+    for role in roles:
+        db.session.add(role)
+
+def init_settings(db, setting_names):
+    name_of_settings = map(lambda r: r.name, setting_names)
+    rows = db.session.query(Setting).filter(Setting.name.in_(name_of_settings)).all()
+    name_of_rows = map(lambda r: r.name, rows)
+
+    settings = filter(lambda r: r.name not in rows, setting_names)
+
+    for setting in settings:
+        db.session.add(setting)
+
+def init_records():
+    # create initial user roles and turn off maintenance mode
+    admin_role = Role('Administrator', 'Administrator')
+    user_role = Role('User', 'User')
+    maintenance_setting = Setting('maintenance', 'False')
+    fullscreen_layout_setting = Setting('fullscreen_layout', 'True')
+    record_helper_setting = Setting('record_helper', 'True')
+    default_table_size_setting = Setting('default_record_table_size', '15')
+
+    # Check if record already exists
+    init_roles(db, [admin_role, user_role])
+    init_settings(db, [maintenance_setting, fullscreen_layout_setting, record_helper_setting, default_table_size_setting])
+
+    db_commit = db.session.commit()
+    commit_version_control(db_commit)
+
+def commit_version_control(db_commit):
+    if not os.path.exists(SQLALCHEMY_MIGRATE_REPO):
+        api.create(SQLALCHEMY_MIGRATE_REPO, 'database repository')
+        api.version_control(SQLALCHEMY_DATABASE_URI, SQLALCHEMY_MIGRATE_REPO)
+    elif db_commit is not None:
+        api.version_control(SQLALCHEMY_DATABASE_URI, SQLALCHEMY_MIGRATE_REPO, api.version(SQLALCHEMY_MIGRATE_REPO))
+
+if __name__ == '__main__':
+    start()