1from __future__ import print_function
2
3from copy import copy, deepcopy
4import datetime
5import inspect
6import sys
7import traceback
8
9from django.core.management import call_command
10from django.core.management.commands import loaddata
11from django.db import models
12from django import VERSION as DJANGO_VERSION
13
14import south.db
15from south import exceptions
16from south.db import DEFAULT_DB_ALIAS
17from south.models import MigrationHistory
18from south.signals import ran_migration
19from south.utils.py3 import StringIO, iteritems
20
21
22class Migrator(object):
23    def __init__(self, verbosity=0, interactive=False):
24        self.verbosity = int(verbosity)
25        self.interactive = bool(interactive)
26
27    @staticmethod
28    def title(target):
29        raise NotImplementedError()
30
31    def print_title(self, target):
32        if self.verbosity:
33            print(self.title(target))
34
35    @staticmethod
36    def status(target):
37        raise NotImplementedError()
38
39    def print_status(self, migration):
40        status = self.status(migration)
41        if self.verbosity and status:
42            print(status)
43
44    @staticmethod
45    def orm(migration):
46        raise NotImplementedError()
47
48    def backwards(self, migration):
49        return self._wrap_direction(migration.backwards(), migration.prev_orm())
50
51    def direction(self, migration):
52        raise NotImplementedError()
53
54    @staticmethod
55    def _wrap_direction(direction, orm):
56        args = inspect.getargspec(direction)
57        if len(args[0]) == 1:
58            # Old migration, no ORM should be passed in
59            return direction
60        return (lambda: direction(orm))
61
62    @staticmethod
63    def record(migration, database):
64        raise NotImplementedError()
65
66    def run_migration_error(self, migration, extra_info=''):
67        return (
68            ' ! Error found during real run of migration! Aborting.\n'
69            '\n'
70            ' ! Since you have a database that does not support running\n'
71            ' ! schema-altering statements in transactions, we have had \n'
72            ' ! to leave it in an interim state between migrations.\n'
73            '%s\n'
74            ' ! The South developers regret this has happened, and would\n'
75            ' ! like to gently persuade you to consider a slightly\n'
76            ' ! easier-to-deal-with DBMS (one that supports DDL transactions)\n'
77            ' ! NOTE: The error which caused the migration to fail is further up.'
78        ) % extra_info
79
80    def run_migration(self, migration, database):
81        migration_function = self.direction(migration)
82        south.db.db.start_transaction()
83        try:
84            migration_function()
85            south.db.db.execute_deferred_sql()
86            if not isinstance(getattr(self, '_wrapper', self), DryRunMigrator):
87                # record us as having done this in the same transaction,
88                # since we're not in a dry run
89                self.record(migration, database)
90        except:
91            south.db.db.rollback_transaction()
92            if not south.db.db.has_ddl_transactions:
93                print(self.run_migration_error(migration))
94            print("Error in migration: %s" % migration)
95            raise
96        else:
97            try:
98                south.db.db.commit_transaction()
99            except:
100                print("Error during commit in migration: %s" % migration)
101                raise
102
103
104    def run(self, migration, database):
105        # Get the correct ORM.
106        south.db.db.current_orm = self.orm(migration)
107        # If we're not already in a dry run, and the database doesn't support
108        # running DDL inside a transaction, *cough*MySQL*cough* then do a dry
109        # run first.
110        if not isinstance(getattr(self, '_wrapper', self), DryRunMigrator):
111            if not south.db.db.has_ddl_transactions:
112                dry_run = DryRunMigrator(migrator=self, ignore_fail=False)
113                dry_run.run_migration(migration, database)
114        return self.run_migration(migration, database)
115
116
117    def send_ran_migration(self, migration, database):
118        ran_migration.send(None,
119                           app=migration.app_label(),
120                           migration=migration,
121                           method=self.__class__.__name__.lower(),
122                           verbosity=self.verbosity,
123                           interactive=self.interactive,
124                           db=database)
125
126    def migrate(self, migration, database):
127        """
128        Runs the specified migration forwards/backwards, in order.
129        """
130        app = migration.migrations._migrations
131        migration_name = migration.name()
132        self.print_status(migration)
133        result = self.run(migration, database)
134        self.send_ran_migration(migration, database)
135        return result
136
137    def migrate_many(self, target, migrations, database):
138        raise NotImplementedError()
139
140
141class MigratorWrapper(object):
142    def __init__(self, migrator, *args, **kwargs):
143        self._migrator = copy(migrator)
144        attributes = dict([(k, getattr(self, k))
145                           for k in self.__class__.__dict__
146                           if not k.startswith('__')])
147        self._migrator.__dict__.update(attributes)
148        self._migrator.__dict__['_wrapper'] = self
149
150    def __getattr__(self, name):
151        return getattr(self._migrator, name)
152
153
154class DryRunMigrator(MigratorWrapper):
155    def __init__(self, ignore_fail=True, *args, **kwargs):
156        super(DryRunMigrator, self).__init__(*args, **kwargs)
157        self._ignore_fail = ignore_fail
158
159    def _run_migration(self, migration):
160        if migration.no_dry_run():
161            if self.verbosity:
162                print(" - Migration '%s' is marked for no-dry-run." % migration)
163            return
164        for name, db in iteritems(south.db.dbs):
165            south.db.dbs[name].dry_run = True
166        # preserve the constraint cache as it can be mutated by the dry run
167        constraint_cache = deepcopy(south.db.db._constraint_cache)
168        if self._ignore_fail:
169            south.db.db.debug, old_debug = False, south.db.db.debug
170        pending_creates = south.db.db.get_pending_creates()
171        south.db.db.start_transaction()
172        migration_function = self.direction(migration)
173        try:
174            try:
175                migration_function()
176                south.db.db.execute_deferred_sql()
177            except:
178                raise exceptions.FailedDryRun(migration, sys.exc_info())
179        finally:
180            south.db.db.rollback_transactions_dry_run()
181            if self._ignore_fail:
182                south.db.db.debug = old_debug
183            south.db.db.clear_run_data(pending_creates)
184            for name, db in iteritems(south.db.dbs):
185                south.db.dbs[name].dry_run = False
186            # restore the preserved constraint cache from before dry run was
187            # executed
188            south.db.db._constraint_cache = constraint_cache
189
190    def run_migration(self, migration, database):
191        try:
192            self._run_migration(migration)
193        except exceptions.FailedDryRun:
194            if self._ignore_fail:
195                return False
196            raise
197
198    def send_ran_migration(self, *args, **kwargs):
199        pass
200
201
202class FakeMigrator(MigratorWrapper):
203    def run(self, migration, database):
204        # Don't actually run, just record as if ran
205        self.record(migration, database)
206        if self.verbosity:
207            print('   (faked)')
208
209    def send_ran_migration(self, *args, **kwargs):
210        pass
211
212
213class LoadInitialDataMigrator(MigratorWrapper):
214
215    def load_initial_data(self, target, db='default'):
216        if target is None or target != target.migrations[-1]:
217            return
218        # Load initial data, if we ended up at target
219        if self.verbosity:
220            print(" - Loading initial data for %s." % target.app_label())
221        if DJANGO_VERSION < (1, 6):
222            self.pre_1_6(target, db)
223        else:
224            self.post_1_6(target, db)
225
226    def pre_1_6(self, target, db):
227        # Override Django's get_apps call temporarily to only load from the
228        # current app
229        old_get_apps = models.get_apps
230        new_get_apps = lambda: [models.get_app(target.app_label())]
231        models.get_apps = new_get_apps
232        loaddata.get_apps = new_get_apps
233        try:
234            call_command('loaddata', 'initial_data', verbosity=self.verbosity, database=db)
235        finally:
236            models.get_apps = old_get_apps
237            loaddata.get_apps = old_get_apps
238
239    def post_1_6(self, target, db):
240        import django.db.models.loading
241        ## build a new 'AppCache' object with just the app we care about.
242        old_cache = django.db.models.loading.cache
243        new_cache = django.db.models.loading.AppCache()
244        new_cache.get_apps = lambda: [new_cache.get_app(target.app_label())]
245
246        ## monkeypatch
247        django.db.models.loading.cache = new_cache
248        try:
249            call_command('loaddata', 'initial_data', verbosity=self.verbosity, database=db)
250        finally:
251            ## unmonkeypatch
252            django.db.models.loading.cache = old_cache
253
254    def migrate_many(self, target, migrations, database):
255        migrator = self._migrator
256        result = migrator.__class__.migrate_many(migrator, target, migrations, database)
257        if result:
258            self.load_initial_data(target, db=database)
259        return True
260
261
262class Forwards(Migrator):
263    """
264    Runs the specified migration forwards, in order.
265    """
266    torun = 'forwards'
267
268    @staticmethod
269    def title(target):
270        if target is not None:
271            return " - Migrating forwards to %s." % target.name()
272        else:
273            assert False, "You cannot migrate forwards to zero."
274
275    @staticmethod
276    def status(migration):
277        return ' > %s' % migration
278
279    @staticmethod
280    def orm(migration):
281        return migration.orm()
282
283    def forwards(self, migration):
284        return self._wrap_direction(migration.forwards(), migration.orm())
285
286    direction = forwards
287
288    @staticmethod
289    def record(migration, database):
290        # Record us as having done this
291        record = MigrationHistory.for_migration(migration, database)
292        try:
293            from django.utils.timezone import now
294            record.applied = now()
295        except ImportError:
296            record.applied = datetime.datetime.utcnow()
297        if database != DEFAULT_DB_ALIAS:
298            record.save(using=database)
299        else:
300            # Django 1.1 and below always go down this branch.
301            record.save()
302
303    def format_backwards(self, migration):
304        if migration.no_dry_run():
305            return "   (migration cannot be dry-run; cannot discover commands)"
306        old_debug, old_dry_run = south.db.db.debug, south.db.db.dry_run
307        south.db.db.debug = south.db.db.dry_run = True
308        stdout = sys.stdout
309        sys.stdout = StringIO()
310        try:
311            try:
312                self.backwards(migration)()
313                return sys.stdout.getvalue()
314            except:
315                raise
316        finally:
317            south.db.db.debug, south.db.db.dry_run = old_debug, old_dry_run
318            sys.stdout = stdout
319
320    def run_migration_error(self, migration, extra_info=''):
321        extra_info = ('\n'
322                      '! You *might* be able to recover with:'
323                      '%s'
324                      '%s' %
325                      (self.format_backwards(migration), extra_info))
326        return super(Forwards, self).run_migration_error(migration, extra_info)
327
328    def migrate_many(self, target, migrations, database):
329        try:
330            for migration in migrations:
331                result = self.migrate(migration, database)
332                if result is False: # The migrations errored, but nicely.
333                    return False
334        finally:
335            # Call any pending post_syncdb signals
336            south.db.db.send_pending_create_signals(verbosity=self.verbosity,
337                                                    interactive=self.interactive)
338        return True
339
340
341class Backwards(Migrator):
342    """
343    Runs the specified migration backwards, in order.
344    """
345    torun = 'backwards'
346
347    @staticmethod
348    def title(target):
349        if target is None:
350            return " - Migrating backwards to zero state."
351        else:
352            return " - Migrating backwards to just after %s." % target.name()
353
354    @staticmethod
355    def status(migration):
356        return ' < %s' % migration
357
358    @staticmethod
359    def orm(migration):
360        return migration.prev_orm()
361
362    direction = Migrator.backwards
363
364    @staticmethod
365    def record(migration, database):
366        # Record us as having not done this
367        record = MigrationHistory.for_migration(migration, database)
368        if record.id is not None:
369            if database != DEFAULT_DB_ALIAS:
370                record.delete(using=database)
371            else:
372                # Django 1.1 always goes down here
373                record.delete()
374
375    def migrate_many(self, target, migrations, database):
376        for migration in migrations:
377            self.migrate(migration, database)
378        return True
379
380
381
382