Add orion db upgrade command

docs/src/install/database.rst

@@ -265,3 +265,22 @@ tests fail because of insufficient user access rights on the database.
    Check if database supports read operation... Success
    Check if database supports count operation... Success
    Check if database supports delete operation... Success
+
+
+Upgrade Database
+================
+
+Database scheme may change from one version of Oríon to another. If such change happens, you will
+get the following error after upgrading Oríon.
+
+.. code-block:: sh
+
+   The database is outdated. You can upgrade it with the command `orion db upgrade`.
+
+Make sure to create a backup of your database before upgrading it. You should also make sure that no
+process writes to the database during the upgrade otherwise the latter could fail. When ready,
+simply run the upgrade command.
+
+.. code-block:: sh
+
+   orion db upgrade

src/orion/core/cli/db/upgrade.py

@@ -0,0 +1,183 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+:mod:`orion.core.cli.db.upgrade` -- Module to upgrade DB schemes
+================================================================
+
+.. module:: test_db
+   :platform: Unix
+   :synopsis: Upgrade the scheme of the databases
+
+"""
+import argparse
+import logging
+import sys
+
+from orion.core.io.database.ephemeraldb import EphemeralCollection
+from orion.core.io.database.mongodb import MongoDB
+from orion.core.io.database.pickleddb import PickledDB
+from orion.core.io.experiment_builder import ExperimentBuilder
+import orion.core.utils.backward as backward
+from orion.storage.base import get_storage
+from orion.storage.legacy import Legacy
+
+
+log = logging.getLogger(__name__)
+
+
+# TODO: Move somewhere else to share with `db setup`.
+def ask_question(question, default=None):
+    """Ask a question to the user and receive an answer.
+
+    Parameters
+    ----------
+    question: str
+        The question to be asked.
+    default: str
+        The default value to use if the user enters nothing.
+
+    Returns
+    -------
+    str
+        The answer provided by the user.
+
+    """
+    if default is not None:
+        question = question + " (default: {}) ".format(default)
+
+    answer = input(question)
+
+    if answer.strip() == "":
+        return default
+
+    return answer
+
+
+def add_subparser(parser):
+    """Add the subparser that needs to be used for this command"""
+    upgrade_db_parser = parser.add_parser('upgrade', help='Upgrade the database scheme')
+
+    upgrade_db_parser.add_argument('-c', '--config', type=argparse.FileType('r'),
+                                   metavar='path-to-config', help="user provided "
+                                   "orion configuration file")
+
+    upgrade_db_parser.add_argument('-f', '--force', action='store_true',
+                                   help="Don't prompt user")
+
+    upgrade_db_parser.set_defaults(func=main)
+
+    return upgrade_db_parser
+
+
+def main(args):
+    """Upgrade the databases for current version"""
+    print("Upgrading your database may damage your data. Make sure to make a backup before the "
+          "upgrade and stop any other process that may read/write the database during the upgrade.")
+
+    if not args.get('force'):
+        action = ''
+        while action not in ['y', 'yes', 'no', 'n']:
+            action = ask_question("Do you wish to proceed? (y/N)", "N").lower()
+
+        if action in ['no', 'n']:
+            sys.exit(0)
+
+    experiment_builder = ExperimentBuilder()
+    local_config = experiment_builder.fetch_full_config(args, use_db=False)
+    local_config['protocol'] = {'type': 'legacy', 'setup': False}
+
+    experiment_builder.setup_storage(local_config)
+
+    storage = get_storage()
+
+    upgrade_db_specifics(storage)
+
+    print('Updating documents...')
+    upgrade_documents(storage)
+    print('Database upgrade completed successfully')
+
+
+def upgrade_db_specifics(storage):
+    """Make upgrades that are specific to some backends"""
+    if isinstance(storage, Legacy):
+        database = storage._db  # pylint: disable=protected-access
+        print('Updating indexes...')
+        update_indexes(database)
+        if isinstance(database, PickledDB):
+            print('Updating pickledb scheme...')
+            upgrade_pickledb(database)
+        elif isinstance(database, MongoDB):
+            print('Updating mongodb scheme...')
+            upgrade_mongodb(database)
+
+
+def upgrade_documents(storage):
+    """Upgrade scheme of the documents"""
+    for experiment in storage.fetch_experiments({}):
+        add_version(experiment)
+        add_priors(experiment)
+        storage.update_experiment(uid=experiment.pop('_id'), **experiment)
+
+
+def add_version(experiment):
+    """Add version 1 if not present"""
+    experiment.setdefault('version', 1)
+
+
+def add_priors(experiment):
+    """Add priors to metadata if not present"""
+    backward.populate_priors(experiment['metadata'])
+
+
+def update_indexes(database):
+    """Remove user from unique indices.
+
+    This is required for migration to v0.1.6+
+    """
+    # For backward compatibility
+    index_info = database.index_information('experiments')
+    deprecated_indices = [('name', 'metadata.user'), ('name', 'metadata.user', 'version'),
+                          'name_1_metadata.user_1', 'name_1_metadata.user_1_version_1']
+
+    for deprecated_idx in deprecated_indices:
+        if deprecated_idx in index_info:
+            database.drop_index('experiments', deprecated_idx)
+
+
+# pylint: disable=unused-argument
+def upgrade_mongodb(database):
+    """Update mongo specific db scheme."""
+    pass
+
+
+def upgrade_pickledb(database):
+    """Update pickledb specific db scheme."""
+    # pylint: disable=protected-access
+    def upgrade_state(self, state):
+        """Set state while ensuring backward compatibility"""
+        self._documents = state['_documents']
+
+        # if indexes are from <=v0.1.6
+        if state['_indexes'] and isinstance(next(iter(state['_indexes'].keys())), tuple):
+            self._indexes = dict()
+            for keys, values in state['_indexes'].items():
+                if isinstance(keys, str):
+                    self._indexes[keys] = values
+                # Convert keys that were registered with old index signature
+                else:
+                    keys = [(key, None) for key in keys]
+                    self.create_index(keys, unique=True)
+        else:
+            self._indexes = state['_indexes']
+
+    old_setstate = getattr(EphemeralCollection, '__setstate__', None)
+    EphemeralCollection.__setstate__ = upgrade_state
+
+    document = database.read('experiments', {})[0]
+    # One document update is enough to fix all collections
+    database.write('experiments', document, query={'_id': document['_id']})
+
+    if old_setstate is not None:
+        EphemeralCollection.__setstate__ = old_setstate
+    else:
+        del EphemeralCollection.__setstate__

src/orion/core/io/database/__init__.py

@@ -305,6 +305,12 @@ class DuplicateKeyError(DatabaseError):
     pass
 
 
+class OutdatedDatabaseError(DatabaseError):
+    """Exception type used when the database is outdated."""
+
+    pass
+
+
 # pylint: disable=too-few-public-methods,abstract-method
 class Database(AbstractDB, metaclass=SingletonFactory):
     """Class used to inject dependency on a database framework."""

src/orion/core/io/database/ephemeraldb.py

@@ -149,20 +149,6 @@ def __init__(self):
         self._indexes = dict()
         self.create_index('_id', unique=True)
 
-    def __setstate__(self, state):
-        """Set state while ensuring backward compatibility"""
-        self._documents = state['_documents']
-
-        # if indexes are from <=v0.1.6
-        if state['_indexes'] and isinstance(next(iter(state['_indexes'].keys())), tuple):
-            self._indexes = dict()
-            for keys in state['_indexes'].keys():
-                # Re-introduce fake ordering
-                keys = [(key, None) for key in keys]
-                self.create_index(keys, unique=True)
-        else:
-            self._indexes = state['_indexes']
-
     def create_index(self, keys, unique=False):
         """Create given indexes if they do not already exist for this collection.
 

src/orion/core/io/experiment_builder.py

@@ -247,6 +247,8 @@ def build_from(self, cmdargs, handle_racecondition=True):
             if handle_racecondition:
                 experiment = self.build_from(cmdargs, handle_racecondition=False)
 
+            raise
+
         return experiment
 
     def build_from_config(self, config):
@@ -295,12 +297,12 @@ def setup_storage(self, config):
 
         """
         # TODO: Fix this in config refactoring
-        db_opts = config.get('protocol', {'type': 'legacy'})
-        dbtype = db_opts.pop('type')
+        storage_opts = config.get('protocol', {'type': 'legacy'})
+        storage_type = storage_opts.pop('type')
 
-        log.debug("Creating %s database client with args: %s", dbtype, db_opts)
+        log.debug("Creating %s storage client with args: %s", storage_type, storage_opts)
         try:
-            Storage(of_type=dbtype, config=config, **db_opts)
+            Storage(of_type=storage_type, config=config, **storage_opts)
         except ValueError:
-            if Storage().__class__.__name__.lower() != dbtype.lower():
+            if Storage().__class__.__name__.lower() != storage_type.lower():
                 raise

src/orion/core/utils/backward.py

@@ -23,3 +23,12 @@ def populate_priors(metadata):
     parser.parse(metadata["user_args"])
     metadata["parser"] = parser.get_state_dict()
     metadata["priors"] = dict(parser.priors)
+
+
+def db_is_outdated(database):
+    """Return True if the database scheme is outdated."""
+    deprecated_indices = [('name', 'metadata.user'), ('name', 'metadata.user', 'version'),
+                          'name_1_metadata.user_1', 'name_1_metadata.user_1_version_1']
+
+    index_information = database.index_information('experiments')
+    return any(index in deprecated_indices for index in index_information.keys())

src/orion/storage/base.py

@@ -35,13 +35,16 @@ def create_experiment(self, config):
         """Insert a new experiment inside the database"""
         raise NotImplementedError()
 
-    def update_experiment(self, experiment, where=None, **kwargs):
+    def update_experiment(self, experiment=None, uid=None, where=None, **kwargs):
         """Update a the fields of a given trials
 
         Parameters
         ----------
-        experiment: Experiment
-            Experiment object to update
+        experiment: Experiment, optional
+           experiment object to retrieve from the database
+
+        uid: str, optional
+            experiment id used to retrieve the trial object
 
         where: Optional[dict]
             constraint experiment must respect
@@ -53,6 +56,14 @@ def update_experiment(self, experiment, where=None, **kwargs):
         -------
         returns true if the underlying storage was updated
 
+        Raises
+        ------
+        UndefinedCall
+            if both experiment and uid are not set
+
+        AssertionError
+            if both experiment and uid are provided and they do not match
+
         """
         raise NotImplementedError()
 

src/orion/storage/legacy.py

@@ -13,7 +13,8 @@
 
 import orion.core
 from orion.core.io.convert import JSONConverter
-from orion.core.io.database import Database
+from orion.core.io.database import Database, OutdatedDatabaseError
+import orion.core.utils.backward as backward
 from orion.core.worker.trial import Trial
 from orion.storage.base import BaseStorageProtocol, FailedUpdate, MissingArguments
 
@@ -52,29 +53,32 @@ class Legacy(BaseStorageProtocol):
         configuration definition passed from experiment_builder
         to storage factory to legacy constructor.
         See `~orion.io.database.Database` for more details
+    setup: bool
+        Setup the database (create indexes)
 
     """
 
-    def __init__(self, config=None):
+    def __init__(self, config=None, setup=True):
         if config is not None:
             setup_database(config)
 
         self._db = Database()
-        self._setup_db()
+
+        if setup:
+            self._setup_db()
 
     def _setup_db(self):
         """Database index setup"""
+        if backward.db_is_outdated(self._db):
+            raise OutdatedDatabaseError("The database is outdated. You can upgrade it with the "
+                                        "command `orion db upgrade`.")
+
+        self._db.index_information('experiment')
         self._db.ensure_index('experiments',
                               [('name', Database.ASCENDING),
                                ('version', Database.ASCENDING)],
                               unique=True)
 
-        # For backward compatibility
-        index_info = self._db.index_information('experiments')
-        for depracated_idx in ['name_1_metadata.user_1', 'name_1_metadata.user_1_version_1']:
-            if depracated_idx in index_info:
-                self._db.drop_index('experiments', depracated_idx)
-
         self._db.ensure_index('experiments', 'metadata.datetime')
 
         self._db.ensure_index('trials', 'experiment')
@@ -87,12 +91,21 @@ def create_experiment(self, config):
         """See :func:`~orion.storage.BaseStorageProtocol.create_experiment`"""
         return self._db.write('experiments', data=config, query=None)
 
-    def update_experiment(self, experiment, where=None, **kwargs):
+    def update_experiment(self, experiment=None, uid=None, where=None, **kwargs):
         """See :func:`~orion.storage.BaseStorageProtocol.update_experiment`"""
+        if experiment is not None and uid is not None:
+            assert experiment._id == uid
+
+        if uid is None:
+            if experiment is None:
+                raise MissingArguments('Either `experiment` or `uid` should be set')
+
+            uid = experiment._id
+
         if where is None:
             where = dict()
 
-        where['_id'] = experiment._id
+        where['_id'] = uid
         return self._db.write('experiments', data=kwargs, query=where)
 
     def fetch_experiments(self, query, selection=None):