libs/dmigrations/mysql/migrations.py

   1 """
   2 These classes represent possible migrations. A migration is simply an object
   3 with an up() method and a down() method - the down() method is allowed to
   4 raise an IrreversibleMigration exception. These objects are instances of
   5 subclasses of BaseMigration. Migration classes will be provided for stuff
   6 ranging from basic SQL migrations to more specialised things such as adding
   7 or removing an index.
   8 """
   9 from dmigrations.migrations import BaseMigration
  10 import re
  11
  12 class IrreversibleMigrationError(Exception):
  13     pass
  14
  15 class Migration(BaseMigration):
  16     "Explict SQL migration, with sql for migrating both up and down"
  17
  18     def __init__(self, sql_up, sql_down=None):
  19         self.sql_up = sql_up
  20         self.sql_down = sql_down
  21
  22     def up(self):
  23         self.execute_sql(self.sql_up)
  24
  25     def down(self):
  26         if self.sql_down:
  27             self.execute_sql(self.sql_down)
  28         else:
  29             raise IrreversibleMigrationError, 'No sql_down provided'
  30
  31     def execute_sql(self, sql, return_rows=False):
  32         "Executes sql, which can be a string or a list of strings"
  33         statements = []
  34         if isinstance(sql, basestring):
  35             # Split string in to multiple statements
  36             statements_re = re.compile(r";[ \t]*$", re.M)
  37             statements = [s for s in statements_re.split(sql) if s.strip()]
  38         elif isinstance(sql, (list, tuple)):
  39             # Assume each item in the list is already an individual statement
  40             statements = sql
  41         else:
  42             assert False, 'sql argument must be string or list/tuple'
  43
  44         from django.db import connection
  45         cursor = connection.cursor()
  46
  47         for statement in statements:
  48             # Escape % due to format strings
  49             cursor.execute(statement.replace('%', '%%'))
  50
  51         if return_rows:
  52             return cursor.fetchall()
  53
  54     def __str__(self):
  55         return 'Migration, up: %r, down: %r' % (self.sql_up, self.sql_down)
  56
  57 class Compound(BaseMigration):
  58     """
  59     A migration that is composed of one or more other migrations. DO NOT USE.
  60     """
  61     def __init__(self, migrations=[]):
  62         self.migrations = migrations
  63
  64     def up(self):
  65         for migration in self.migrations:
  66             migration.up()
  67
  68     def down(self):
  69         for migration in reversed(self.migrations):
  70             migration.down()
  71
  72     def __str__(self):
  73         return 'Compound Migration: %s' % self.migrations
  74
  75
  76 class AddColumn(Migration):
  77     "A migration that adds a database column"
  78
  79     add_column_sql = 'ALTER TABLE `%s_%s` ADD COLUMN `%s` %s;'
  80     drop_column_sql = 'ALTER TABLE `%s_%s` DROP COLUMN `%s`;'
  81     constrain_to_table_sql = 'ALTER TABLE `%s_%s` ADD CONSTRAINT %s FOREIGN KEY (`%s`) REFERENCES `%s` (`id`);'
  82     constrain_to_table_down_sql = 'ALTER TABLE `%s_%s` DROP FOREIGN KEY `%s`;'
  83
  84     def __init__(self, app, model, column, spec, constrain_to_table=None):
  85         self.app, self.model, self.column, self.spec = app, model, column, spec
  86         if constrain_to_table:
  87             # this can only be used for ForeignKeys that link to another table's
  88             # id field. It is not for arbitrary relationships across tables!
  89             # Note also that this will create the ForeignKey field as allowing
  90             # nulls. Even if you don't want it to. This is because if it doesn't
  91             # allow null then the migration will blow up, because we're adding
  92             # the column without adding data to it. So you have to write another
  93             # migration later to change it from NULL to NOT NULL if you need to,
  94             # after you've populated it.
  95
  96             # add the FK constraint
  97             constraint_name = "%s_refs_id_%x" % (column, abs(hash((model,constrain_to_table))))
  98             sql_up = [self.constrain_to_table_sql % (app, model, constraint_name, "%s_id" % column, constrain_to_table)]
  99
 100             sql_up.insert(0,self.add_column_sql % (app, model, "%s_id" % column, spec))
 101             sql_down = [self.drop_column_sql % (app, model, "%s_id" % column)]
 102             # if add_column_sql has NOT NULL in it, bin it
 103             sql_up[0] = sql_up[0].replace(" NOT NULL", "")
 104             # drop FK on sql_down
 105             sql_down.insert(0,self.constrain_to_table_down_sql % (app, model, constraint_name))
 106
 107         else:
 108             sql_up = [self.add_column_sql % (app, model, column, spec)]
 109             sql_down = [self.drop_column_sql % (app, model, column)]
 110
 111         super(AddColumn, self).__init__(
 112             sql_up,
 113             sql_down,
 114         )
 115
 116     def __str__(self):
 117         return "AddColumn: app: %s, model: %s, column: %s, spec: %s" % (
 118             self.app, self.model, self.column, self.spec
 119         )
 120
 121 class DropColumn(AddColumn):
 122     """
 123     A migration that drops a database column. Needs the full column spec so
 124     it can correctly create the down() method.
 125     """
 126     def __init__(self, *args, **kwargs):
 127         super(DropColumn, self).__init__(*args, **kwargs)
 128         # Now swap over the sql_up and sql_down properties
 129         self.sql_up, self.sql_down = self.sql_down, self.sql_up
 130
 131     def __str__(self):
 132         return super(DropColumn, self).replace('AddColumn', 'DropColumn')
 133
 134 class AddIndex(Migration):
 135     "A migration that adds an index (and removes it on down())"
 136
 137     add_index_sql = 'CREATE INDEX `%s` ON `%s_%s` (`%s`);'
 138     drop_index_sql = 'ALTER TABLE %s_%s DROP INDEX `%s`;'
 139
 140     def __init__(self, app, model, column):
 141         self.app, self.model, self.column = app, model, column
 142         index_name = '%s_%s_%s' % (app, model, column)
 143         super(AddIndex, self).__init__(
 144             sql_up = [self.add_index_sql % (index_name, app, model, column)],
 145             sql_down = [self.drop_index_sql % (app, model, index_name)],
 146         )
 147
 148     def __str__(self):
 149         return "AddIndex: app: %s, model: %s, column: %s" % (
 150             self.app, self.model, self.column
 151         )
 152
 153 class DropIndex(AddIndex):
 154     "Drops an index"
 155     def __init__(self, app, model, column):
 156         super(DropIndex, self).__init__(app, model, column)
 157         self.sql_up, self.sql_down = self.sql_down, self.sql_up
 158
 159     def __str__(self):
 160         return super(DropIndex, self).replace('AddIndex', 'DropIndex')
 161
 162 class InsertRows(Migration):
 163     "Inserts some rows in to a table"
 164
 165     insert_row_sql = 'INSERT INTO `%s` (%s) VALUES (%s)'
 166     delete_rows_sql = 'DELETE FROM `%s` WHERE id IN (%s)'
 167
 168     def __init__(self, table_name, columns, insert_rows, delete_ids):
 169         self.table_name = table_name
 170         sql_up = []
 171         from django.db import connection # so we can use escape_string
 172         connection.cursor() # Opens connection if not already open
 173
 174         def escape(v):
 175             if v is None:
 176                 return 'null'
 177             v = unicode(v) # In case v is an integer or long
 178             # escape_string wants a bytestring
 179             escaped = connection.connection.escape_string(v.encode('utf8'))
 180             # We get bugs if we use bytestrings elsewhere, so convert back to unicode
 181             # http://sourceforge.net/forum/forum.php?thread_id=1609278&forum_id=70461
 182             return u"'%s'" % escaped.decode('utf8')
 183
 184         for row in insert_rows:
 185             values = ', '.join(map(escape, row))
 186             sql_up.append(
 187                 self.insert_row_sql % (
 188                     table_name, ', '.join(map(str, columns)), values
 189                 )
 190             )
 191
 192         if delete_ids:
 193             sql_down = [self.delete_rows_sql % (table_name, ', '.join(map(str, delete_ids)))]
 194         else:
 195             sql_down = ["SELECT 1"]
 196
 197         super(InsertRows, self).__init__(
 198             sql_up = ["BEGIN"] + sql_up + ["COMMIT"],
 199             sql_down = ["BEGIN"] + sql_down + ["COMMIT"],
 200         )