From e419ec63abb56e622698e5db0616afb8a04ea7c2 Mon Sep 17 00:00:00 2001
From: Joao Pereira and Sarah McAlear <pair+jpereira+smcalear@pivotal.io>
Date: Wed, 15 Mar 2017 14:39:25 -0400
Subject: [PATCH] Switch to Alembic and Flask-migration db migration system

---
 README                                   |  16 ++
 requirements.txt                         |   1 +
 web/migrations/alembic.ini               |  45 ++++
 web/migrations/env.py                    |  99 +++++++
 web/migrations/script.py.mako            |  24 ++
 web/migrations/versions/09d53fca90c7_.py | 233 +++++++++++++++++
 web/migrations/versions/fdc58d9bd449_.py | 119 +++++++++
 web/pgadmin/__init__.py                  |  28 +-
 web/pgadmin/setup/__init__.py            |   3 +
 web/pgadmin/setup/db_upgrade.py          |  16 ++
 web/pgadmin/setup/db_version.py          |  19 ++
 web/pgadmin/setup/user_info.py           |  58 +++++
 web/setup.py                             | 427 +------------------------------
 13 files changed, 641 insertions(+), 447 deletions(-)
 create mode 100644 web/migrations/alembic.ini
 create mode 100755 web/migrations/env.py
 create mode 100755 web/migrations/script.py.mako
 create mode 100644 web/migrations/versions/09d53fca90c7_.py
 create mode 100644 web/migrations/versions/fdc58d9bd449_.py
 create mode 100644 web/pgadmin/setup/__init__.py
 create mode 100644 web/pgadmin/setup/db_upgrade.py
 create mode 100644 web/pgadmin/setup/db_version.py
 create mode 100644 web/pgadmin/setup/user_info.py

diff --git a/README b/README
index b906e0d6..973fdec1 100644
--- a/README
+++ b/README
@@ -180,6 +180,22 @@ http://www.tylerbutler.com/2012/05/how-to-install-python-pip-and-virtualenv-on-w
 Once a virtual environment has been created and enabled, setup can continue
 from step 4 above.
 
+Create Database Migrations
+--------------------------
+
+In order to make changes to the SQLite DB, navigate to the 'web' directory:
+
+(pgadmin4) $ cd $PGADMIN4_SRC/web
+
+Create a migration file with the following command:
+
+(pgadmin4) $ FLASK_APP=pgAdmin4.py flask db revision
+
+This will create a file in: $PGADMIN4_SRC/web/migrations/versions/ .
+Add any changes to the 'upgrade' function.
+
+There is no need to increment the SETTINGS_SCHEMA_VERSION.
+
 Configuring the Runtime
 -----------------------
 
diff --git a/requirements.txt b/requirements.txt
index 071069d3..f4bbe5b8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,6 +13,7 @@ Flask-Gravatar==0.4.2
 Flask-HTMLmin==1.2
 Flask-Login==0.3.2
 Flask-Mail==0.9.1
+Flask-Migrate==2.0.3
 Flask-Principal==0.4.0
 Flask-Security==1.7.5
 Flask-SQLAlchemy==2.1
diff --git a/web/migrations/alembic.ini b/web/migrations/alembic.ini
new file mode 100644
index 00000000..f8ed4801
--- /dev/null
+++ b/web/migrations/alembic.ini
@@ -0,0 +1,45 @@
+# A generic, single database configuration.
+
+[alembic]
+# template used to generate migration files
+# file_template = %%(rev)s_%%(slug)s
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S
diff --git a/web/migrations/env.py b/web/migrations/env.py
new file mode 100755
index 00000000..251d9f42
--- /dev/null
+++ b/web/migrations/env.py
@@ -0,0 +1,99 @@
+from __future__ import with_statement
+from alembic import context
+from sqlalchemy import engine_from_config, pool
+from logging.config import fileConfig
+import logging
+
+# Save the logger configuration to restore after migration run
+root = logging.getLogger()
+logging_handlers = root.handlers[:]
+logging_filters = root.filters[:]
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+fileConfig(config.config_file_name)
+logger = logging.getLogger('alembic.env')
+
+# add your model's MetaData object here
+# for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
+from flask import current_app
+config.set_main_option('sqlalchemy.url',
+                       current_app.config.get('SQLALCHEMY_DATABASE_URI'))
+target_metadata = current_app.extensions['migrate'].db.metadata
+
+# other values from the config, defined by the needs of env.py,
+# can be acquired:
+# my_important_option = config.get_main_option("my_important_option")
+# ... etc.
+
+
+def run_migrations_offline():
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(url=url)
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online():
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+
+    """
+
+    # this callback is used to prevent an auto-migration from being generated
+    # when there are no changes to the schema
+    # reference: http://alembic.readthedocs.org/en/latest/cookbook.html
+    def process_revision_directives(context, revision, directives):
+        if getattr(config.cmd_opts, 'autogenerate', False):
+            script = directives[0]
+            if script.upgrade_ops.is_empty():
+                directives[:] = []
+                logger.info('No changes in schema detected.')
+
+    engine = engine_from_config(config.get_section(config.config_ini_section),
+                                prefix='sqlalchemy.',
+                                poolclass=pool.NullPool)
+
+    connection = engine.connect()
+    context.configure(connection=connection,
+                      target_metadata=target_metadata,
+                      process_revision_directives=process_revision_directives,
+                      **current_app.extensions['migrate'].configure_args)
+
+    try:
+        with context.begin_transaction():
+            context.run_migrations()
+    finally:
+        connection.close()
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()
+
+# Remove the logging configuration applied by the migration and restore the application logging configuration
+root = logging.getLogger()
+map(root.removeHandler, root.handlers[:])
+map(root.removeFilter, root.filters[:])
+map(root.addHandler, logging_handlers[:])
+map(root.addFilter, logging_filters[:])
diff --git a/web/migrations/script.py.mako b/web/migrations/script.py.mako
new file mode 100755
index 00000000..2c015630
--- /dev/null
+++ b/web/migrations/script.py.mako
@@ -0,0 +1,24 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision = ${repr(up_revision)}
+down_revision = ${repr(down_revision)}
+branch_labels = ${repr(branch_labels)}
+depends_on = ${repr(depends_on)}
+
+
+def upgrade():
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade():
+    ${downgrades if downgrades else "pass"}
diff --git a/web/migrations/versions/09d53fca90c7_.py b/web/migrations/versions/09d53fca90c7_.py
new file mode 100644
index 00000000..80332b0a
--- /dev/null
+++ b/web/migrations/versions/09d53fca90c7_.py
@@ -0,0 +1,233 @@
+"""Update DB to version 14
+
+Revision ID: 09d53fca90c7
+Revises: fdc58d9bd449
+Create Date: 2017-03-13 12:27:30.543908
+
+"""
+import base64
+
+import sys
+from alembic import op
+from pgadmin.model import db, Server
+import config
+import os
+from pgadmin.setup import get_version
+
+# revision identifiers, used by Alembic.
+
+revision = '09d53fca90c7'
+down_revision = 'fdc58d9bd449'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    version = get_version()
+    # Changes introduced in schema version 2
+    if version < 2:
+        # Create the 'server' table
+        db.metadata.create_all(db.engine, tables=[Server.__table__])
+    if version < 3:
+        db.engine.execute(
+            'ALTER TABLE server ADD COLUMN comment TEXT(1024)'
+        )
+    if version < 4:
+        db.engine.execute(
+            'ALTER TABLE server ADD COLUMN password TEXT(64)'
+        )
+    if version < 5:
+        db.engine.execute('ALTER TABLE server ADD COLUMN role text(64)')
+    if version < 6:
+        db.engine.execute("ALTER TABLE server RENAME TO server_old")
+        db.engine.execute("""
+    CREATE TABLE server (
+        id INTEGER NOT NULL,
+        user_id INTEGER NOT NULL,
+        servergroup_id INTEGER NOT NULL,
+        name VARCHAR(128) NOT NULL,
+        host VARCHAR(128) NOT NULL,
+        port INTEGER NOT NULL CHECK (port >= 1024 AND port <= 65534),
+        maintenance_db VARCHAR(64) NOT NULL,
+        username VARCHAR(64) NOT NULL,
+        ssl_mode VARCHAR(16) NOT NULL CHECK (
+            ssl_mode IN (
+                'allow', 'prefer', 'require', 'disable', 'verify-ca', 'verify-full'
+                )),
+        comment VARCHAR(1024), password TEXT(64), role text(64),
+        PRIMARY KEY (id),
+        FOREIGN KEY(user_id) REFERENCES user (id),
+        FOREIGN KEY(servergroup_id) REFERENCES servergroup (id)
+    )""")
+        db.engine.execute("""
+    INSERT INTO server (
+        id, user_id, servergroup_id, name, host, port, maintenance_db, username,
+        ssl_mode, comment, password, role
+    ) SELECT
+        id, user_id, servergroup_id, name, host, port, maintenance_db, username,
+        ssl_mode, comment, password, role
+    FROM server_old""")
+        db.engine.execute("DROP TABLE server_old")
+
+    if version < 8:
+        db.engine.execute("""
+    CREATE TABLE module_preference(
+        id INTEGER PRIMARY KEY,
+        name VARCHAR(256) NOT NULL
+        )""")
+
+        db.engine.execute("""
+    CREATE TABLE preference_category(
+        id INTEGER PRIMARY KEY,
+        mid INTEGER,
+        name VARCHAR(256) NOT NULL,
+
+        FOREIGN KEY(mid) REFERENCES module_preference(id)
+        )""")
+
+        db.engine.execute("""
+    CREATE TABLE preferences (
+
+        id INTEGER PRIMARY KEY,
+        cid INTEGER NOT NULL,
+        name VARCHAR(256) NOT NULL,
+
+        FOREIGN KEY(cid) REFERENCES preference_category (id)
+        )""")
+
+        db.engine.execute("""
+    CREATE TABLE user_preferences (
+
+        pid INTEGER,
+        uid INTEGER,
+        value VARCHAR(1024) NOT NULL,
+
+        PRIMARY KEY (pid, uid),
+        FOREIGN KEY(pid) REFERENCES preferences (pid),
+        FOREIGN KEY(uid) REFERENCES user (id)
+        )""")
+
+    if version < 9:
+        db.engine.execute("""
+    CREATE TABLE IF NOT EXISTS debugger_function_arguments (
+        server_id INTEGER ,
+        database_id INTEGER ,
+        schema_id INTEGER ,
+        function_id INTEGER ,
+        arg_id INTEGER ,
+        is_null INTEGER NOT NULL CHECK (is_null >= 0 AND is_null <= 1) ,
+        is_expression INTEGER NOT NULL CHECK (is_expression >= 0 AND is_expression <= 1) ,
+        use_default INTEGER NOT NULL CHECK (use_default >= 0 AND use_default <= 1) ,
+        value TEXT,
+        PRIMARY KEY (server_id, database_id, schema_id, function_id, arg_id)
+        )""")
+
+    if version < 10:
+        db.engine.execute("""
+    CREATE TABLE process(
+        user_id INTEGER NOT NULL,
+        pid TEXT NOT NULL,
+        desc TEXT NOT NULL,
+        command TEXT NOT NULL,
+        arguments TEXT,
+        start_time TEXT,
+        end_time TEXT,
+        logdir TEXT,
+        exit_code INTEGER,
+        acknowledge TEXT,
+        PRIMARY KEY(pid),
+        FOREIGN KEY(user_id) REFERENCES user (id)
+        )""")
+
+    if version < 11:
+        db.engine.execute("""
+    UPDATE role
+        SET name = 'Administrator',
+        description = 'pgAdmin Administrator Role'
+        WHERE name = 'Administrators'
+        """)
+
+        db.engine.execute("""
+    INSERT INTO role ( name, description )
+                VALUES ('User', 'pgAdmin User Role')
+        """)
+
+    if version < 12:
+        db.engine.execute("ALTER TABLE server RENAME TO server_old")
+        db.engine.execute("""
+    CREATE TABLE server (
+        id INTEGER NOT NULL,
+        user_id INTEGER NOT NULL,
+        servergroup_id INTEGER NOT NULL,
+        name VARCHAR(128) NOT NULL,
+        host VARCHAR(128) NOT NULL,
+        port INTEGER NOT NULL CHECK (port >= 1024 AND port <= 65535),
+        maintenance_db VARCHAR(64) NOT NULL,
+        username VARCHAR(64) NOT NULL,
+        ssl_mode VARCHAR(16) NOT NULL CHECK (
+            ssl_mode IN (
+                'allow', 'prefer', 'require', 'disable', 'verify-ca', 'verify-full'
+                )),
+        comment VARCHAR(1024), password TEXT(64), role text(64),
+        PRIMARY KEY (id),
+        FOREIGN KEY(user_id) REFERENCES user (id),
+        FOREIGN KEY(servergroup_id) REFERENCES servergroup (id)
+    )""")
+        db.engine.execute("""
+    INSERT INTO server (
+        id, user_id, servergroup_id, name, host, port, maintenance_db, username,
+        ssl_mode, comment, password, role
+    ) SELECT
+        id, user_id, servergroup_id, name, host, port, maintenance_db, username,
+        ssl_mode, comment, password, role
+    FROM server_old""")
+        db.engine.execute("DROP TABLE server_old")
+
+    if version < 13:
+        db.engine.execute("""
+    ALTER TABLE SERVER
+        ADD COLUMN discovery_id TEXT
+        """)
+
+    if version < 14:
+        db.engine.execute("""
+    CREATE TABLE keys (
+        name TEST NOT NULL,
+        value TEXT NOT NULL,
+        PRIMARY KEY (name))
+                    """)
+
+        sql = "INSERT INTO keys (name, value) VALUES ('CSRF_SESSION_KEY', '%s')" % base64.urlsafe_b64encode(
+            os.urandom(32)).decode()
+        db.engine.execute(sql)
+
+        sql = "INSERT INTO keys (name, value) VALUES ('SECRET_KEY', '%s')" % base64.urlsafe_b64encode(
+            os.urandom(32)).decode()
+        db.engine.execute(sql)
+
+        # If SECURITY_PASSWORD_SALT is not in the config, but we're upgrading, then it must (unless the
+        # user edited the main config - which they shouldn't have done) have been at it's default
+        # value, so we'll use that. Otherwise, use whatever we can find in the config.
+        if hasattr(config, 'SECURITY_PASSWORD_SALT'):
+            sql = "INSERT INTO keys (name, value) VALUES ('SECURITY_PASSWORD_SALT', '%s')" % config.SECURITY_PASSWORD_SALT
+        else:
+            sql = "INSERT INTO keys (name, value) VALUES ('SECURITY_PASSWORD_SALT', 'SuperSecret3')"
+        db.engine.execute(sql)
+
+        # Finally, update the schema version
+
+    # version.value = config.SETTINGS_SCHEMA_VERSION
+
+    db.engine.execute(
+        'UPDATE version set value="%s" WHERE name = "ConfigDB"' % config.SETTINGS_SCHEMA_VERSION
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    print(u"""
+    Cannot downgrade from this version
+    Exiting...""")
+    sys.exit(1)
+    # ### end Alembic commands ###
diff --git a/web/migrations/versions/fdc58d9bd449_.py b/web/migrations/versions/fdc58d9bd449_.py
new file mode 100644
index 00000000..d43bdfac
--- /dev/null
+++ b/web/migrations/versions/fdc58d9bd449_.py
@@ -0,0 +1,119 @@
+"""Initial database creation
+
+Revision ID: fdc58d9bd449
+Revises: 
+Create Date: 2017-03-13 11:15:16.401139
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from pgadmin.model import db
+from pgadmin.setup import get_version
+
+from pgadmin.setup import user_info
+
+# revision identifiers, used by Alembic.
+revision = 'fdc58d9bd449'
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    if get_version() != -1:
+        return
+
+    op.create_table('version',
+                    sa.Column('name', sa.String(length=32), nullable=False),
+                    sa.Column('value', sa.Integer(), nullable=False),
+                    sa.PrimaryKeyConstraint('name')
+                    )
+    op.create_table('user',
+                    sa.Column('id', sa.Integer(), nullable=False),
+                    sa.Column('email', sa.String(length=256), nullable=False),
+                    sa.Column('password', sa.String(length=256), nullable=True),
+                    sa.Column('active', sa.Boolean(), nullable=False),
+                    sa.Column('confirmed_at', sa.DateTime(), nullable=True),
+                    sa.PrimaryKeyConstraint('id'),
+                    sa.UniqueConstraint('email')
+                    )
+    op.create_table('role',
+                    sa.Column('id', sa.Integer(), nullable=False),
+                    sa.Column('name', sa.String(length=128), nullable=False),
+                    sa.Column('description', sa.String(length=256), nullable=False),
+                    sa.PrimaryKeyConstraint('id'),
+                    sa.UniqueConstraint('name')
+                    )
+    op.create_table('setting',
+                    sa.Column('user_id', sa.Integer(), nullable=False),
+                    sa.Column('setting', sa.String(length=256), nullable=False),
+                    sa.Column('value', sa.String(length=1024), nullable=True),
+                    sa.ForeignKeyConstraint(['user_id'], ['user.id'], ),
+                    sa.PrimaryKeyConstraint('user_id', 'setting')
+                    )
+    op.create_table('roles_users',
+                    sa.Column('user_id', sa.Integer(), nullable=True),
+                    sa.Column('role_id', sa.Integer(), nullable=True),
+                    sa.ForeignKeyConstraint(['role_id'], ['role.id'], ),
+                    sa.ForeignKeyConstraint(['user_id'], ['user.id'], )
+                    )
+    op.create_table('servergroup',
+                    sa.Column('id', sa.Integer(), nullable=False),
+                    sa.Column('user_id', sa.Integer(), nullable=False),
+                    sa.Column('name', sa.String(length=128), nullable=False),
+                    sa.ForeignKeyConstraint(['user_id'], ['user.id'], ),
+                    sa.PrimaryKeyConstraint('id'),
+                    sa.UniqueConstraint('user_id', 'name')
+                    )
+    op.create_table('server',
+                    sa.Column('id', sa.Integer(), nullable=False),
+                    sa.Column('user_id', sa.Integer(), nullable=False),
+                    sa.Column('servergroup_id', sa.Integer(), nullable=False),
+                    sa.Column('name', sa.String(length=128), nullable=False),
+                    sa.Column('host', sa.String(length=128), nullable=False),
+                    sa.Column('port', sa.Integer(), nullable=False),
+                    sa.Column('maintenance_db', sa.String(length=64), nullable=False),
+                    sa.Column('username', sa.String(length=64), nullable=False),
+                    sa.Column('ssl_mode', sa.String(length=16), nullable=False),
+                    sa.ForeignKeyConstraint(['servergroup_id'], ['servergroup.id'], ),
+                    sa.ForeignKeyConstraint(['user_id'], ['user.id'], ),
+                    sa.PrimaryKeyConstraint('id')
+                    )
+    email, password = user_info()
+    db.engine.execute("""
+INSERT INTO "user"
+    VALUES(1, '%s',
+           '%s',
+           1, NULL)
+    """ % (email, password))
+    db.engine.execute("""
+INSERT INTO "version"
+VALUES('ConfigDB', 2);
+    """)
+    db.engine.execute("""
+INSERT INTO "role"
+VALUES(1, 'Administrators', 'pgAdmin Administrators Role')
+        """)
+    db.engine.execute("""
+INSERT INTO "roles_users"
+VALUES(1, 1);
+    """)
+    db.engine.execute("""
+INSERT INTO "servergroup"
+VALUES(1, 1, 'Servers')
+""")
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('server')
+    op.drop_table('servergroup')
+    op.drop_table('roles_users')
+    op.drop_table('setting')
+    op.drop_table('role')
+    op.drop_table('user')
+    op.drop_table('version')
+    # ### end Alembic commands ###
diff --git a/web/pgadmin/__init__.py b/web/pgadmin/__init__.py
index 1a143251..8f110dde 100644
--- a/web/pgadmin/__init__.py
+++ b/web/pgadmin/__init__.py
@@ -10,7 +10,7 @@
 """The main pgAdmin module. This handles the application initialisation tasks,
 such as setup of logging, dynamic loading of modules etc."""
 import logging
-import os, sys, time
+import os, sys
 from collections import defaultdict
 from importlib import import_module
 
@@ -221,30 +221,8 @@ def create_app(app_name=None):
 
     # Upgrade the schema (if required)
     with app.app_context():
-        try:
-            version = Version.query.filter_by(name='ConfigDB').first()
-        except:
-            backup_file = config.SQLITE_PATH + '.' + time.strftime("%Y%m%d%H%M%S")
-            app.logger.error(
-                """The configuration database ({0}) appears to be corrupt.\n\n"""
-                """The database will be moved to {1}.\n"""
-                """Please restart {2} to create a new configuration database.\n""".format(
-                    config.SQLITE_PATH, backup_file, config.APP_NAME
-                )
-            )
-
-            os.rename(config.SQLITE_PATH, backup_file)
-            exit(1)
-
-        # Pre-flight checks
-        if int(version.value) < int(config.SETTINGS_SCHEMA_VERSION):
-            app.logger.info(
-                """Upgrading the database schema from version {0} to {1}.""".format(
-                    version.value, config.SETTINGS_SCHEMA_VERSION
-                )
-            )
-            from setup import do_upgrade
-            do_upgrade(app, user_datastore, version)
+        from setup import db_upgrade
+        db_upgrade(app)
 
     ##########################################################################
     # Setup security
diff --git a/web/pgadmin/setup/__init__.py b/web/pgadmin/setup/__init__.py
new file mode 100644
index 00000000..8d7e149b
--- /dev/null
+++ b/web/pgadmin/setup/__init__.py
@@ -0,0 +1,3 @@
+from user_info import user_info
+from db_version import get_version
+from db_upgrade import db_upgrade
diff --git a/web/pgadmin/setup/db_upgrade.py b/web/pgadmin/setup/db_upgrade.py
new file mode 100644
index 00000000..ee7e3120
--- /dev/null
+++ b/web/pgadmin/setup/db_upgrade.py
@@ -0,0 +1,16 @@
+import os
+import flask_migrate
+
+from pgadmin import db
+
+
+def db_upgrade(app):
+    from pgadmin.utils import u, fs_encoding
+    with app.app_context():
+        flask_migrate.Migrate(app, db)
+        migration_folder = os.path.join(
+                os.path.dirname(os.path.realpath(u(__file__, fs_encoding))),
+                os.pardir, os.pardir,
+                u'migrations'
+            )
+        flask_migrate.upgrade(migration_folder)
diff --git a/web/pgadmin/setup/db_version.py b/web/pgadmin/setup/db_version.py
new file mode 100644
index 00000000..7f5acf31
--- /dev/null
+++ b/web/pgadmin/setup/db_version.py
@@ -0,0 +1,19 @@
+from pgadmin.model import Version
+import config
+import sys
+
+
+def get_version():
+    try:
+        version = Version.query.filter_by(name='ConfigDB').first()
+    except Exception:
+        return -1
+
+    if int(version.value) > int(config.SETTINGS_SCHEMA_VERSION):
+        print(u"""
+    The database schema version is {0}, whilst the version required by the \
+    software is {1}.
+    Exiting...""".format(version.value, config.SETTINGS_SCHEMA_VERSION))
+        sys.exit(1)
+
+    return version.value
diff --git a/web/pgadmin/setup/user_info.py b/web/pgadmin/setup/user_info.py
new file mode 100644
index 00000000..ffd1065d
--- /dev/null
+++ b/web/pgadmin/setup/user_info.py
@@ -0,0 +1,58 @@
+import config
+import string
+import random
+import os
+import re
+import getpass
+
+
+def user_info():
+    if config.SERVER_MODE is False:
+        print(u"NOTE: Configuring authentication for DESKTOP mode.")
+        email = config.DESKTOP_USER
+        p1 = ''.join([
+            random.choice(string.ascii_letters + string.digits)
+            for n in range(32)
+        ])
+
+    else:
+        print(u"NOTE: Configuring authentication for SERVER mode.\n")
+
+        if all(value in os.environ for value in
+               ['PGADMIN_SETUP_EMAIL', 'PGADMIN_SETUP_PASSWORD']):
+            email = ''
+            p1 = ''
+            if os.environ['PGADMIN_SETUP_EMAIL'] and os.environ[
+                'PGADMIN_SETUP_PASSWORD']:
+                email = os.environ['PGADMIN_SETUP_EMAIL']
+                p1 = os.environ['PGADMIN_SETUP_PASSWORD']
+        else:
+            # Prompt the user for their default username and password.
+            print(
+                u"Enter the email address and password to use for the initial "
+                u"pgAdmin user account:\n"
+            )
+
+            email_filter = re.compile(
+                "^[a-zA-Z0-9.!#$%&'*+\/=?^_`{|}~-]+@[a-zA-Z0-9]"
+                "(?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(?:\.[a-zA-Z0-9]"
+                "(?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)*$")
+
+            email = input("Email address: ")
+            while email == '' or not email_filter.match(email):
+                print(u'Invalid email address. Please try again.')
+                email = input("Email address: ")
+
+            def pprompt():
+                return getpass.getpass(), getpass.getpass('Retype password:')
+
+            p1, p2 = pprompt()
+            while p1 != p2 or len(p1) < 6:
+                if p1 != p2:
+                    print(u'Passwords do not match. Please try again.')
+                else:
+                    print(
+                        u'Password must be at least 6 characters. Please try again.'
+                    )
+                p1, p2 = pprompt()
+    return email, p1
diff --git a/web/setup.py b/web/setup.py
index c5a329c5..ce9f6daf 100755
--- a/web/setup.py
+++ b/web/setup.py
@@ -10,33 +10,22 @@
 """Perform the initial setup of the application, by creating the auth
 and settings database."""
 
-import base64
-import getpass
 import os
-import random
-import re
-import string
 import sys
 
 from flask import Flask
-from flask_security import Security, SQLAlchemyUserDatastore
-from flask_security.utils import encrypt_password
 
 # We need to include the root directory in sys.path to ensure that we can
 # find everything we need when running in the standalone runtime.
+from pgadmin.setup import db_upgrade
+
 root = os.path.dirname(os.path.realpath(__file__))
 if sys.path[0] != root:
     sys.path.insert(0, root)
 
-
 # Configuration settings
 import config
-
-# Get the config database schema version. We store this in pgadmin.model
-# as it turns out that putting it in the config files isn't a great idea
-from pgadmin.model import db, Role, User, Server, ServerGroup, Version, Keys, \
-    SCHEMA_VERSION
-from pgadmin.utils.versioned_template_loader import VersionedTemplateLoader
+from pgadmin.model import db, SCHEMA_VERSION
 
 config.SETTINGS_SCHEMA_VERSION = SCHEMA_VERSION
 
@@ -45,352 +34,6 @@ if hasattr(__builtins__, 'raw_input'):
     input = raw_input
     range = xrange
 
-
-def do_setup(app):
-    """Create a new settings database from scratch"""
-
-    if config.SERVER_MODE is False:
-        print(u"NOTE: Configuring authentication for DESKTOP mode.")
-        email = config.DESKTOP_USER
-        p1 = ''.join([
-                         random.choice(string.ascii_letters + string.digits)
-                         for n in range(32)
-                         ])
-
-    else:
-        print(u"NOTE: Configuring authentication for SERVER mode.\n")
-
-        if all(value in os.environ for value in
-               ['PGADMIN_SETUP_EMAIL', 'PGADMIN_SETUP_PASSWORD']):
-            email = ''
-            p1 = ''
-            if os.environ['PGADMIN_SETUP_EMAIL'] and os.environ[
-                'PGADMIN_SETUP_PASSWORD']:
-                email = os.environ['PGADMIN_SETUP_EMAIL']
-                p1 = os.environ['PGADMIN_SETUP_PASSWORD']
-        else:
-            # Prompt the user for their default username and password.
-            print(
-                u"Enter the email address and password to use for the initial "
-                u"pgAdmin user account:\n"
-            )
-
-            email_filter = re.compile(
-                "^[a-zA-Z0-9.!#$%&'*+\/=?^_`{|}~-]+@[a-zA-Z0-9]"
-                "(?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(?:\.[a-zA-Z0-9]"
-                "(?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)*$")
-
-            email = input("Email address: ")
-            while email == '' or not email_filter.match(email):
-                print(u'Invalid email address. Please try again.')
-                email = input("Email address: ")
-
-            def pprompt():
-                return getpass.getpass(), getpass.getpass('Retype password:')
-
-            p1, p2 = pprompt()
-            while p1 != p2 or len(p1) < 6:
-                if p1 != p2:
-                    print(u'Passwords do not match. Please try again.')
-                else:
-                    print(
-                        u'Password must be at least 6 characters. Please try again.'
-                    )
-                p1, p2 = pprompt()
-
-    # Setup Flask-Security
-    user_datastore = SQLAlchemyUserDatastore(db, User, Role)
-    Security(app, user_datastore)
-
-    with app.app_context():
-        password = encrypt_password(p1)
-
-        db.create_all()
-        user_datastore.create_role(
-            name='Administrator',
-            description='pgAdmin Administrator Role'
-        )
-        user_datastore.create_role(
-            name='User',
-            description='pgAdmin User Role'
-        )
-        user_datastore.create_user(email=email, password=password)
-        db.session.flush()
-        user_datastore.add_role_to_user(email, 'Administrator')
-
-        # Get the user's ID and create the default server group
-        user = User.query.filter_by(email=email).first()
-        server_group = ServerGroup(user_id=user.id, name="Servers")
-        db.session.merge(server_group)
-
-        # Set the schema version
-        version = Version(
-            name='ConfigDB', value=config.SETTINGS_SCHEMA_VERSION
-        )
-        db.session.merge(version)
-        db.session.commit()
-
-        # Create the keys
-        key = Keys(name='CSRF_SESSION_KEY', value=config.CSRF_SESSION_KEY)
-        db.session.merge(key)
-
-        key = Keys(name='SECRET_KEY', value=config.SECRET_KEY)
-        db.session.merge(key)
-
-        key = Keys(name='SECURITY_PASSWORD_SALT', value=config.SECURITY_PASSWORD_SALT)
-        db.session.merge(key)
-
-        db.session.commit()
-
-    # Done!
-    print(u"")
-    print(
-        u"The configuration database has been created at {0}".format(
-            config.SQLITE_PATH
-        )
-    )
-
-
-def do_upgrade(app, datastore, version):
-    """Upgrade an existing settings database"""
-    #######################################################################
-    # Run whatever is required to update the database schema to the current
-    # version.
-    #######################################################################
-
-    with app.app_context():
-        version = Version.query.filter_by(name='ConfigDB').first()
-
-        # Pre-flight checks
-        if int(version.value) > int(config.SETTINGS_SCHEMA_VERSION):
-            print(u"""
-The database schema version is {0}, whilst the version required by the \
-software is {1}.
-Exiting...""".format(version.value, config.SETTINGS_SCHEMA_VERSION))
-            sys.exit(1)
-        elif int(version.value) == int(config.SETTINGS_SCHEMA_VERSION):
-            print(u"""
-The database schema version is {0} as required.
-Exiting...""".format(version.value))
-            sys.exit(1)
-
-        app.logger.info(
-            u"NOTE: Upgrading database schema from version %d to %d." %
-            (version.value, config.SETTINGS_SCHEMA_VERSION)
-        )
-
-        #######################################################################
-        # Run whatever is required to update the database schema to the current
-        # version. Always use "< REQUIRED_VERSION" as the test for readability
-        #######################################################################
-
-        # Changes introduced in schema version 2
-        if int(version.value) < 2:
-            # Create the 'server' table
-            db.metadata.create_all(db.engine, tables=[Server.__table__])
-        if int(version.value) < 3:
-            db.engine.execute(
-                'ALTER TABLE server ADD COLUMN comment TEXT(1024)'
-            )
-        if int(version.value) < 4:
-            db.engine.execute(
-                'ALTER TABLE server ADD COLUMN password TEXT(64)'
-            )
-        if int(version.value) < 5:
-            db.engine.execute('ALTER TABLE server ADD COLUMN role text(64)')
-        if int(version.value) < 6:
-            db.engine.execute("ALTER TABLE server RENAME TO server_old")
-            db.engine.execute("""
-CREATE TABLE server (
-    id INTEGER NOT NULL,
-    user_id INTEGER NOT NULL,
-    servergroup_id INTEGER NOT NULL,
-    name VARCHAR(128) NOT NULL,
-    host VARCHAR(128) NOT NULL,
-    port INTEGER NOT NULL CHECK (port >= 1024 AND port <= 65534),
-    maintenance_db VARCHAR(64) NOT NULL,
-    username VARCHAR(64) NOT NULL,
-    ssl_mode VARCHAR(16) NOT NULL CHECK (
-        ssl_mode IN (
-            'allow', 'prefer', 'require', 'disable', 'verify-ca', 'verify-full'
-            )),
-    comment VARCHAR(1024), password TEXT(64), role text(64),
-    PRIMARY KEY (id),
-    FOREIGN KEY(user_id) REFERENCES user (id),
-    FOREIGN KEY(servergroup_id) REFERENCES servergroup (id)
-)""")
-            db.engine.execute("""
-INSERT INTO server (
-    id, user_id, servergroup_id, name, host, port, maintenance_db, username,
-    ssl_mode, comment, password, role
-) SELECT
-    id, user_id, servergroup_id, name, host, port, maintenance_db, username,
-    ssl_mode, comment, password, role
-FROM server_old""")
-            db.engine.execute("DROP TABLE server_old")
-
-        if int(version.value) < 8:
-            app.logger.info(
-                "Creating the preferences tables..."
-            )
-            db.engine.execute("""
-CREATE TABLE module_preference(
-    id INTEGER PRIMARY KEY,
-    name VARCHAR(256) NOT NULL
-    )""")
-
-            db.engine.execute("""
-CREATE TABLE preference_category(
-    id INTEGER PRIMARY KEY,
-    mid INTEGER,
-    name VARCHAR(256) NOT NULL,
-
-    FOREIGN KEY(mid) REFERENCES module_preference(id)
-    )""")
-
-            db.engine.execute("""
-CREATE TABLE preferences (
-
-    id INTEGER PRIMARY KEY,
-    cid INTEGER NOT NULL,
-    name VARCHAR(256) NOT NULL,
-
-    FOREIGN KEY(cid) REFERENCES preference_category (id)
-    )""")
-
-            db.engine.execute("""
-CREATE TABLE user_preferences (
-
-    pid INTEGER,
-    uid INTEGER,
-    value VARCHAR(1024) NOT NULL,
-
-    PRIMARY KEY (pid, uid),
-    FOREIGN KEY(pid) REFERENCES preferences (pid),
-    FOREIGN KEY(uid) REFERENCES user (id)
-    )""")
-
-        if int(version.value) < 9:
-            db.engine.execute("""
-CREATE TABLE IF NOT EXISTS debugger_function_arguments (
-    server_id INTEGER ,
-    database_id INTEGER ,
-    schema_id INTEGER ,
-    function_id INTEGER ,
-    arg_id INTEGER ,
-    is_null INTEGER NOT NULL CHECK (is_null >= 0 AND is_null <= 1) ,
-    is_expression INTEGER NOT NULL CHECK (is_expression >= 0 AND is_expression <= 1) ,
-    use_default INTEGER NOT NULL CHECK (use_default >= 0 AND use_default <= 1) ,
-    value TEXT,
-    PRIMARY KEY (server_id, database_id, schema_id, function_id, arg_id)
-    )""")
-
-        if int(version.value) < 10:
-            db.engine.execute("""
-CREATE TABLE process(
-    user_id INTEGER NOT NULL,
-    pid TEXT NOT NULL,
-    desc TEXT NOT NULL,
-    command TEXT NOT NULL,
-    arguments TEXT,
-    start_time TEXT,
-    end_time TEXT,
-    logdir TEXT,
-    exit_code INTEGER,
-    acknowledge TEXT,
-    PRIMARY KEY(pid),
-    FOREIGN KEY(user_id) REFERENCES user (id)
-    )""")
-
-        if int(version.value) < 11:
-            db.engine.execute("""
-UPDATE role
-    SET name = 'Administrator',
-    description = 'pgAdmin Administrator Role'
-    WHERE name = 'Administrators'
-    """)
-
-            db.engine.execute("""
-INSERT INTO role ( name, description )
-            VALUES ('User', 'pgAdmin User Role')
-    """)
-
-        if int(version.value) < 12:
-            db.engine.execute("ALTER TABLE server RENAME TO server_old")
-            db.engine.execute("""
-CREATE TABLE server (
-    id INTEGER NOT NULL,
-    user_id INTEGER NOT NULL,
-    servergroup_id INTEGER NOT NULL,
-    name VARCHAR(128) NOT NULL,
-    host VARCHAR(128) NOT NULL,
-    port INTEGER NOT NULL CHECK (port >= 1024 AND port <= 65535),
-    maintenance_db VARCHAR(64) NOT NULL,
-    username VARCHAR(64) NOT NULL,
-    ssl_mode VARCHAR(16) NOT NULL CHECK (
-        ssl_mode IN (
-            'allow', 'prefer', 'require', 'disable', 'verify-ca', 'verify-full'
-            )),
-    comment VARCHAR(1024), password TEXT(64), role text(64),
-    PRIMARY KEY (id),
-    FOREIGN KEY(user_id) REFERENCES user (id),
-    FOREIGN KEY(servergroup_id) REFERENCES servergroup (id)
-)""")
-            db.engine.execute("""
-INSERT INTO server (
-    id, user_id, servergroup_id, name, host, port, maintenance_db, username,
-    ssl_mode, comment, password, role
-) SELECT
-    id, user_id, servergroup_id, name, host, port, maintenance_db, username,
-    ssl_mode, comment, password, role
-FROM server_old""")
-            db.engine.execute("DROP TABLE server_old")
-
-        if int(version.value) < 13:
-            db.engine.execute("""
-ALTER TABLE SERVER
-    ADD COLUMN discovery_id TEXT
-    """)
-
-        if int(version.value) < 14:
-            db.engine.execute("""
-CREATE TABLE keys (
-    name TEST NOT NULL,
-    value TEXT NOT NULL,
-    PRIMARY KEY (name))
-                """)
-
-            sql = "INSERT INTO keys (name, value) VALUES ('CSRF_SESSION_KEY', '%s')" % base64.urlsafe_b64encode(os.urandom(32)).decode()
-            db.engine.execute(sql)
-
-            sql = "INSERT INTO keys (name, value) VALUES ('SECRET_KEY', '%s')" % base64.urlsafe_b64encode(os.urandom(32)).decode()
-            db.engine.execute(sql)
-
-            # If SECURITY_PASSWORD_SALT is not in the config, but we're upgrading, then it must (unless the
-            # user edited the main config - which they shouldn't have done) have been at it's default
-            # value, so we'll use that. Otherwise, use whatever we can find in the config.
-            if hasattr(config, 'SECURITY_PASSWORD_SALT'):
-                sql = "INSERT INTO keys (name, value) VALUES ('SECURITY_PASSWORD_SALT', '%s')" % config.SECURITY_PASSWORD_SALT
-            else:
-                sql = "INSERT INTO keys (name, value) VALUES ('SECURITY_PASSWORD_SALT', 'SuperSecret3')"
-            db.engine.execute(sql)
-
-    # Finally, update the schema version
-    version.value = config.SETTINGS_SCHEMA_VERSION
-    db.session.merge(version)
-
-    db.session.commit()
-
-    # Done!
-    app.logger.info(
-        "The configuration database %s has been upgraded to version %d" %
-        (config.SQLITE_PATH, config.SETTINGS_SCHEMA_VERSION)
-    )
-
-
-###############################################################################
-# Do stuff!
-###############################################################################
 if __name__ == '__main__':
     app = Flask(__name__)
 
@@ -403,67 +46,7 @@ if __name__ == '__main__':
         'sqlite:///' + config.SQLITE_PATH.replace('\\', '/')
     db.init_app(app)
 
+    db_upgrade(app)
+
     print(u"pgAdmin 4 - Application Initialisation")
     print(u"======================================\n")
-
-    from pgadmin.utils import u, fs_encoding, file_quote
-
-    local_config = os.path.join(
-        os.path.dirname(os.path.realpath(u(__file__, fs_encoding))),
-        u'config_local.py'
-    )
-
-    # Check if the database exists. If it does, tell the user and exit.
-    if os.path.isfile(config.SQLITE_PATH):
-        print(
-            u"The configuration database '{0}' already exists.".format(
-                config.SQLITE_PATH
-            )
-        )
-        print(u"Entering upgrade mode...")
-
-        # Setup Flask-Security
-        user_datastore = SQLAlchemyUserDatastore(db, User, Role)
-
-        # Always use "< REQUIRED_VERSION" as the test for readability
-        with app.app_context():
-            version = Version.query.filter_by(name='ConfigDB').first()
-
-            # Pre-flight checks
-            if int(version.value) > int(config.SETTINGS_SCHEMA_VERSION):
-                print(u"""
-The database schema version is %d, whilst the version required by the \
-software is %d.
-Exiting...""" % (version.value, config.SETTINGS_SCHEMA_VERSION))
-                sys.exit(1)
-            elif int(version.value) == int(config.SETTINGS_SCHEMA_VERSION):
-                print(u"""
-The database schema version is %d as required.
-Exiting...""" % (version.value))
-                sys.exit(1)
-
-            print(u"NOTE: Upgrading database schema from version %d to %d." % (
-                version.value, config.SETTINGS_SCHEMA_VERSION
-            ))
-            do_upgrade(app, user_datastore, version)
-    else:
-        # Get some defaults for the various keys
-        config.CSRF_SESSION_KEY = base64.urlsafe_b64encode(os.urandom(32)).decode()
-        config.SECRET_KEY = base64.urlsafe_b64encode(os.urandom(32)).decode()
-        config.SECURITY_PASSWORD_SALT = base64.urlsafe_b64encode(os.urandom(32)).decode()
-
-        app.config.from_object(config)
-
-        directory = os.path.dirname(config.SQLITE_PATH)
-
-        if not os.path.exists(directory):
-            os.makedirs(directory, int('700', 8))
-
-        db_file = os.open(config.SQLITE_PATH, os.O_CREAT, int('600', 8))
-        os.close(db_file)
-
-        print(u"""
-The configuration database - '{0}' does not exist.
-Entering initial setup mode...""".format(config.SQLITE_PATH))
-
-        do_setup(app)
-- 
2.12.0

