~jcsackett/charmworld/bac-tag-constraints

# Copyright 2012, 2013 Canonical Ltd.  This software is licensed under the

# GNU Affero General Public License version 3 (see the file LICENSE).

"""Migrate data for the mongodb instance.

Loads migrations from the latest available in the versions directory.

"""

import argparse

from datetime import datetime

import imp

from os import listdir

from os.path import abspath

from os.path import dirname

from os.path import join

import re

from charmworld.models import getconnection

from charmworld.models import getdb

from charmworld.utils import get_ini

SCRIPT_TEMPLATE = """

# {description}

def upgrade(db):

    \"\"\"Complete this function with work to be done for the migration/update.

    db is the pymongo db instance for our datastore. Charms are in db.charms

    for instance.

    \"\"\"

"""

def str_to_filename(s):

    """Replaces spaces, quotes, and double underscores to underscores."""

    s = s.replace(' ', '_')\

         .replace('"', '_')\

         .replace("'", '_')\

         .replace(".", "_")

    while '__' in s:

        s = s.replace('__', '_')

    return s

class DataStore(object):

    """Communicate with the data store to determine version status."""

    def __init__(self, db):

        """Talk to the data store

        :param db: The mongo db connection.

        """

        self.db = db

    @property

    def current_version(self):

        """Return the current version of the data store."""

        found = self.db.migration_version.find_one({

            '_id': 'version'

        })

        if found:

            return found['version']

        else:

            return None

    def update_version(self, version):

        """Update the version number in the data store."""

        self.db.migration_version.save({

            '_id': 'version',

            'version': version,

            'date': datetime.now(),

        })

    def version_datastore(self):

        """Init the data to track the current version in the datastore.

        The data store starts out with version 0.

        Raises exception if the store is already versioned.

        """

        if self.current_version is not None:

            raise Exception('Data store is already versioned: {0}'.format(

                self.current_version))

        self.db.migration_version.insert({

            '_id': 'version',

            'version': 0,

            'date': datetime.now(),

        })

class Versions(object):

    def __init__(self, path):

        """Collect version scripts and store them in self.versions

        """

        FILENAME_WITH_VERSION = re.compile(r'^(\d{3,}).*')

        self.versions_dir = path

        # Create temporary list of files, allowing skipped version numbers.

        files = listdir(path)

        versions = {}

        for name in files:

            if not name.endswith('.py'):

                # Skip .pyc and the like

                continue

            match = FILENAME_WITH_VERSION.match(name)

            if match:

                num = int(match.group(1))

                versions[num] = name

            else:

                pass  # Must be a helper file or something, let's ignore it.

        self.versions = {}

        version_numbers = versions.keys()

        version_numbers.sort()

        self.version_indexes = version_numbers

        for idx in version_numbers:

            self.versions[idx] = versions[idx]

    def create_new_version_file(self, description):

        """Create Python files for new version"""

        version = self.latest + 1

        if not description:

            raise ValueError('Please provide a short migration description.')

        filename = "{0}_{1}.py".format(

            str(version).zfill(3),

            str_to_filename(description)

        )

        filepath = join(self.versions_dir, filename)

        with open(filepath, 'w') as new_migration:

            new_migration.write(

                SCRIPT_TEMPLATE.format(description=description))

        # Update our current store of version data.

        self.versions[version] = filename

        self.version_indexes.append(version)

        return filename

    @property

    def latest(self):

        """:returns: Latest version in Collection"""

        return self.version_indexes[-1] if len(self.version_indexes) > 0 else 0

    def upgrade(self, datastore, init):

        """Run `upgrade` methods for required version files.

        :param datastore: An instance of DataStore

        """

        current_version = datastore.current_version

        if current_version is None:

            if init:

                # We want to auto init the database anyway.

                datastore.version_datastore()

                current_version = 0

            else:

                raise Exception('Data store is not versioned')

        if current_version >= self.latest:

            # Nothing to do here. All migrations processed already.

            return None

        while current_version < self.latest:

            # Let's get processing.

            next_version = current_version + 1

            module_name = self.versions[next_version]

            module = imp.load_source(

                module_name.strip('.py'),

                join(self.versions_dir, module_name))

            getattr(module, 'upgrade')(datastore.db)

            current_version = next_version

            datastore.update_version(current_version)

        return current_version

def parse_args():

    desc = "Mongo migration tool."

    parser = argparse.ArgumentParser(description=desc)

    subparsers = parser.add_subparsers(help='sub-command help')

    parser_current = subparsers.add_parser('current')

    parser_current.set_defaults(func=Commands.current)

    parser_latest = subparsers.add_parser('latest')

    parser_latest.set_defaults(func=Commands.latest)

    parser_new = subparsers.add_parser('new')

    parser_new.add_argument(

        '-d', '--description', action='store',

        required=True,

        help='The description of the new migration.')

    parser_new.set_defaults(func=Commands.new)

    parser_upgrade = subparsers.add_parser('upgrade')

    parser_upgrade.set_defaults(func=Commands.upgrade)

    parser_upgrade.add_argument(

        '-i', '--init', action='store_true',

        default=False,

        help='Auto init the database if not already init.')

    args = parser.parse_args()

    return args

class Commands(object):

    """Container for the various commands the cli can execute"""

    @staticmethod

    def get_datastore(ini):

        connection = getconnection(ini)

        db = getdb(connection, ini.get('mongo.database'))

        ds = DataStore(db)

        return ds

    @classmethod

    def current(cls, ini, args):

        """Fetch the current migration version of the data store."""

        ds = cls.get_datastore(ini)

        version = ds.current_version

        if version is None:

            print 'Data store is not currently versioned.'

        else:

            print version

    @staticmethod

    def latest(ini, args):

        """Determine the latest migration version available."""

        migrations = Versions(ini['migrations'])

        print migrations.latest

    @staticmethod

    def new(ini, args):

        """Generate a new migration file to be completed."""

        migrations = Versions(ini['migrations'])

        filename = migrations.create_new_version_file(args.description)

        print "Created new migration: {0}".format(filename)

    @classmethod

    def upgrade(cls, ini, args):

        """Upgrade the data store to the latest available migration."""

        ds = cls.get_datastore(ini)

        migrations = Versions(ini['migrations'])

        new_version = migrations.upgrade(ds, args.init)

        if new_version is None:

            print 'There are no new migrations to apply'

        else:

            print "Updated the datastore to version: {0}".format(new_version)

def main():

    """Target for the console entry point."""

    args = parse_args()

    ini = get_ini()

    # Add the migration path to the ini.

    migration_path = join(abspath(dirname(__file__)), 'versions')

    ini['migrations'] = migration_path

    args.func(ini, args)

if __name__ == '__main__':

    main()