~fabricematrat/charmworld/redirect-charm

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
# 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.search import ElasticSearchClient
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 run_migration(self, db, index_client, module_name):
        module = imp.load_source(
            module_name.strip('.py'),
            join(self.versions_dir, module_name))
        getattr(module, 'upgrade')(db, index_client)

    def upgrade(self, datastore, index_client, 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]
            self.run_migration(datastore.db, index_client, module_name)
            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)
        index_client = ElasticSearchClient.from_settings(ini)
        migrations = Versions(ini['migrations'])
        new_version = migrations.upgrade(ds, index_client, 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()