pgadmin4/web/setup.py

#########################################################################
#
# pgAdmin 4 - PostgreSQL Tools
#
# Copyright (C) 2013 - 2015, The pgAdmin Development Team
# This software is released under the PostgreSQL Licence
#
##########################################################################

"""Perform the initial setup of the application, by creating the auth
and settings database."""

from flask import Flask
from flask.ext.sqlalchemy import SQLAlchemy
from flask.ext.security import Security, SQLAlchemyUserDatastore
from flask.ext.security.utils import encrypt_password
from pgadmin.settings.settings_model import db, Role, User, Server, ServerGroup, Version

import getpass, os, random, sys, string

# Configuration settings
import config

def do_setup(app):
    """Create a new settings database from scratch"""
    if config.SERVER_MODE is False:
        print("NOTE: Configuring authentication for DESKTOP mode.")
        email = config.DESKTOP_USER
        p1 = ''.join([random.choice(string.ascii_letters + string.digits) for n in xrange(32)])

    else:
        print("NOTE: Configuring authentication for SERVER mode.\n")

        # Prompt the user for their default username and password.
        print("Enter the email address and password to use for the initial pgAdmin user account:\n")
        email = ''
        while email == '':
            email = raw_input("Email address: ")

        pprompt = lambda: (getpass.getpass(), getpass.getpass('Retype password: '))

        p1, p2 = pprompt()
        while p1 != p2:
            print('Passwords do not match. Try again')
            p1, p2 = pprompt()

    # Setup Flask-Security
    user_datastore = SQLAlchemyUserDatastore(db, User, Role)
    security = Security(app, user_datastore)

    with app.app_context():
        password = encrypt_password(p1)

        db.create_all()
        user_datastore.create_role(name='Administrators', description='pgAdmin Administrators Role')
        user_datastore.create_user(email=email, password=password)
        user_datastore.add_role_to_user(email, 'Administrators')

        # 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()

    # Done!
    print("")
    print("The configuration database has been created at %s" % config.SQLITE_PATH)


def do_upgrade(app, datastore, security, version):
    """Upgrade an existing settings database"""
    #######################################################################
    # Run whatever is required to update the database schema to the current
    # version.
    #######################################################################

    # Changes introduced in schema version 2
    if int(version.value) < 2:
        # Create the 'server' table
        db.metadata.create_all(db.engine, tables=[Server.__table__])
    elif int(version.value) < 3:
        db.engine.execute('ALTER TABLE server ADD COLUMN comment TEXT(1024)');

    # Finally, update the schema version
    version.value = config.SETTINGS_SCHEMA_VERSION
    db.session.merge(version)

    db.session.commit()

    # Done!
    print("")
    print("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__)
    app.config.from_object(config)
    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + config.SQLITE_PATH.replace('\\', '/')
    db.init_app(app)

    print("pgAdmin 4 - Application Initialisation")
    print("======================================\n")

    local_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'config_local.py')
    if not os.path.isfile(local_config):
        print("The configuration file %s does not exist.\n" % local_config)
        print("Before running this application, ensure that config_local.py has been created")
        print("and sets values for SECRET_KEY, SECURITY_PASSWORD_SALT and CSRF_SESSION_KEY")
        print("at bare minimum. See config.py for more information and a complete list of")
        print("settings. Exiting...")
        sys.exit(1)

    # Check if the database exists. If it does, tell the user and exit.
    if os.path.isfile(config.SQLITE_PATH):
        print("The configuration database %s already exists.\nEntering upgrade mode...\n" % config.SQLITE_PATH)

        # Setup Flask-Security
        user_datastore = SQLAlchemyUserDatastore(db, User, Role)
        security = Security(app, user_datastore)

        # 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("The database schema version is %d, whilst the version required by the software is %d.\nExiting..."
                      % (version.value, config.SETTINGS_SCHEMA_VERSION))
                sys.exit(1)
            elif int(version.value) == int(config.SETTINGS_SCHEMA_VERSION):
                print("The database schema version is %d as required.\nExiting..." % (version.value))
                sys.exit(1)

            print("NOTE: Upgrading database schema from version %d to %d." % (version.value, config.SETTINGS_SCHEMA_VERSION))
            do_upgrade(app, user_datastore, security, version)
    else:
        print("The configuration database %s does not exist.\nEntering initial setup mode...\n" % config.SQLITE_PATH)
        do_setup(app)