script: better database initialization

* BETTER Uses SQLAlchemy and SQLAlchemy-Utils to init the database instead of execute mysql in a python subshell. (closes inveniosoftware#2846) (closes inveniosoftware#2844) * NEW Adds support for PostgreSQL database initialization. Signed-off-by: Leonardo Rossi <[email protected]> Reviewed-by: Jiri Kuncar <[email protected]>
hachreak · Mar 20, 2015 · 2674834 · 2674834
1 parent 66aa5be
commit 2674834
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 56 deletions.
diff --git a/invenio/base/scripts/database.py b/invenio/base/scripts/database.py
@@ -23,18 +23,10 @@
 
 import datetime
 
-import os
-
-import sys
-
-from pipes import quote
-
 from flask import current_app
 
 from invenio.ext.script import Manager, change_command_name, print_progress
 
-from six import iteritems
-
 manager = Manager(usage="Perform database operations")
 
 # Shortcuts for manager options to keep code DRY.
@@ -54,47 +46,53 @@
 @option_yes_i_know
 def init(user='root', password='', yes_i_know=False):
     """Initialize database and user."""
-    from invenio.ext.sqlalchemy import db
+    from invenio.ext.sqlalchemy.utils import initialize_database_user
     from invenio.utils.text import wrap_text_in_a_box, wait_for_user
 
+    from sqlalchemy_utils.functions import database_exists, create_database, \
+        drop_database
+
+    from sqlalchemy.engine.url import URL
+    from sqlalchemy import create_engine
+
     # Step 0: confirm deletion
     wait_for_user(wrap_text_in_a_box(
         "WARNING: You are going to destroy your database tables! Run first"
         " `inveniomanage database drop`."
     ))
 
-    # Step 1: drop database and recreate it
-    if db.engine.name == 'mysql':
-        # FIXME improve escaping
-        args = dict((k, str(v).replace('$', '\$'))
-                    for (k, v) in iteritems(current_app.config)
-                    if k.startswith('CFG_DATABASE'))
-        args = dict(zip(args, map(quote, args.values())))
-        prefix = ('{cmd} -u {user} --password={password} '
-                  '-h {CFG_DATABASE_HOST} -P {CFG_DATABASE_PORT} ')
-        cmd_prefix = prefix.format(cmd='mysql', user=user, password=password,
-                                   **args)
-        cmd_admin_prefix = prefix.format(cmd='mysqladmin', user=user,
-                                         password=password,
-                                         **args)
-        cmds = [
-            cmd_prefix + '-e "DROP DATABASE IF EXISTS {CFG_DATABASE_NAME}"',
-            (cmd_prefix + '-e "CREATE DATABASE IF NOT EXISTS '
-             '{CFG_DATABASE_NAME} DEFAULT CHARACTER SET utf8 '
-             'COLLATE utf8_general_ci"'),
-            # Create user and grant access to database.
-            (cmd_prefix + '-e "GRANT ALL PRIVILEGES ON '
-             '{CFG_DATABASE_NAME}.* TO {CFG_DATABASE_USER}@localhost '
-             'IDENTIFIED BY {CFG_DATABASE_PASS}"'),
-            cmd_admin_prefix + 'flush-privileges'
-        ]
-        for cmd in cmds:
-            cmd = cmd.format(**args)
-            print(cmd)
-            if os.system(cmd):
-                print("ERROR: failed execution of", cmd, file=sys.stderr)
-                sys.exit(1)
-        print('>>> Database has been installed.')
+    # Step 1: create URI to connect admin user
+    cfg = current_app.config
+    SQLALCHEMY_DATABASE_URI = URL(
+        cfg.get('CFG_DATABASE_TYPE', 'mysql'),
+        username=user,
+        password=password,
+        host=cfg.get('CFG_DATABASE_HOST'),
+        database=cfg.get('CFG_DATABASE_NAME'),
+        port=cfg.get('CFG_DATABASE_PORT'),
+    )
+
+    # Step 2: drop the database if already exists
+    if database_exists(SQLALCHEMY_DATABASE_URI):
+        drop_database(SQLALCHEMY_DATABASE_URI)
+        print('>>> Database has been dropped.')
+
+    # Step 3: create the database
+    create_database(SQLALCHEMY_DATABASE_URI, encoding='utf8')
+    print('>>> Database has been created.')
+
+    # Step 4: setup connection with special user
+    engine = create_engine(SQLALCHEMY_DATABASE_URI)
+    engine.connect()
+
+    # Step 5: grant privileges for the user
+    initialize_database_user(
+        engine=engine,
+        database_name=current_app.config['CFG_DATABASE_NAME'],
+        database_user=current_app.config['CFG_DATABASE_USER'],
+        database_pass=current_app.config['CFG_DATABASE_PASS'],
+    )
+    print('>>> Database user has been initialized.')
 
 
 @option_yes_i_know
@@ -148,7 +146,7 @@ def _dropper(items, prefix, dropper):
                 dropper(table)
                 dropped += 1
             except Exception:
-                print('\r', '>>> problem with dropping ', table)
+                print('\r>>> problem with dropping {0}'.format(table))
                 current_app.logger.exception(table)
 
         if dropped == N:
@@ -201,7 +199,7 @@ def _creator(items, prefix, creator):
                 creator(table)
                 created += 1
             except Exception:
-                print('\r', '>>> problem with creating ', table)
+                print('\r>>> problem with creating {0}'.format(table))
                 current_app.logger.exception(table)
 
         if created == N:

diff --git a/invenio/ext/script/__init__.py b/invenio/ext/script/__init__.py
@@ -62,10 +62,11 @@ def generate_secret_key():
 def print_progress(p, L=40, prefix='', suffix=''):
     """Print textual progress bar."""
     bricks = int(p * L)
-    print('\r', prefix, end=' ')
-    print('[{0}{1}] {2}%'.format('#' * bricks, ' ' * (L - bricks),
-                                 int(p * 100)), end=' ')
-    print(suffix, end=' ')
+    print('\r{prefix} [{bricks}{spaces}] {progress}% {suffix}'.format(
+        prefix=prefix, suffix=suffix,
+        bricks='#' * bricks, spaces=' ' * (L - bricks),
+        progress=int(p * 100),
+    ), end=' ')
 
 
 def check_for_software_updates(flash_message=False):

diff --git a/invenio/ext/sqlalchemy/utils.py b/invenio/ext/sqlalchemy/utils.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 #
 # This file is part of Invenio.
-# Copyright (C) 2011, 2012, 2013, 2014 CERN.
+# Copyright (C) 2011, 2012, 2013, 2014, 2015 CERN.
 #
 # Invenio is free software; you can redistribute it and/or
 # modify it under the terms of the GNU General Public License as
@@ -38,14 +38,12 @@ def save(self):
 import sys
 
 from intbitset import intbitset
+
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.declarative import declared_attr
 from sqlalchemy.orm import class_mapper, properties
-from sqlalchemy.orm.collections import (
-    InstrumentedList,
-    attribute_mapped_collection,
-    collection,
-)
+from sqlalchemy.orm.collections import InstrumentedList, collection
+
 
 first_cap_re = re.compile('(.)([A-Z][a-z]+)')
 all_cap_re = re.compile('([a-z0-9])([A-Z])')
@@ -100,7 +98,7 @@ def todict(self):
         def convert_datetime(value):
             try:
                 return value.strftime("%Y-%m-%d %H:%M:%S")
-            except:
+            except Exception:
                 return ''
 
         for c in self.__table__.columns:
@@ -127,7 +125,7 @@ def fromdict(self, args):
         #    name = str(c).split('.')[1]
         #    try:
         #        d = args[name]
-        #    except:
+        #    except Exception:
         #        continue
         #
         #    setattr(self, c.name, d)
@@ -173,7 +171,7 @@ def new_func(self, *a, **k):
             resp = orig_func(self, *a, **k)
             db.session.commit()
             return resp
-        except:
+        except Exception:
             db.session.rollback()
             raise
 
@@ -296,6 +294,7 @@ class OrderedList(InstrumentedList):
     """Implemented ordered instrumented list."""
 
     def append(self, item):
+        """Append item."""
         if self:
             s = sorted(self, key=lambda obj: obj.score)
             item.score = s[-1].score + 1
@@ -304,6 +303,7 @@ def append(self, item):
         InstrumentedList.append(self, item)
 
     def set(self, item, index=0):
+        """Set item."""
         if self:
             s = sorted(self, key=lambda obj: obj.score)
             if index >= len(s):
@@ -322,6 +322,7 @@ def set(self, item, index=0):
         InstrumentedList.append(self, item)
 
     def pop(self, item):
+        """Pop item."""
         # FIXME
         if self:
             obj_list = sorted(self, key=lambda obj: obj.score)
@@ -362,3 +363,50 @@ def __repr__(self):
             return '%s(%r)' % (type(self).__name__, self._data)
 
     return MultiMappedCollection
+
+
+def initialize_database_user(engine, database_name, database_user,
+                             database_pass):
+    """Grant user's privileges.
+
+    :param engine: engine to use to execute queries
+    :param database_name: database name
+    :param database_user: the username that you want to init
+    :param database_pass: password for the user
+    """
+    if engine.name == 'mysql':
+        # create user and grant privileges
+        engine.execute(
+            """GRANT ALL PRIVILEGES ON `%s`.* """
+            """TO %s@'%%' IDENTIFIED BY %s """,
+            database_name, database_user, database_pass)
+    elif engine.name == 'postgresql':
+        # check if already exists
+        res = engine.execute(
+            """SELECT 1 FROM pg_roles WHERE rolname='{0}' """
+            .format(database_user.replace("'", "''")))
+        # if already don't exists, create user
+        if not res.first():
+            engine.execute(
+                """CREATE USER "{0}" WITH PASSWORD '{1}' """
+                .format(
+                    database_user.replace('"', '""'),
+                    database_pass.replace("'", "''")
+                ))
+            # grant privileges for user
+            engine.execute(
+                """grant all privileges on database """
+                """ "{0}" to "{1}" """
+                .format(
+                    database_name.replace('"', '""'),
+                    database_user.replace('"', '""')
+                ))
+    else:
+        raise Exception((
+            """Database engine %(engine)s not supported. """
+            """You need to manually adds privileges to %(database_user)s """
+            """in order to access to the database %(database_name)s.""") % {
+                'engine': engine.name,
+                'database_user': database_user,
+                'database_name': database_name
+            })