Fix migrations on Python 2.

[mediagoblin.git] / mediagoblin / db / migrations.py

diff --git a/mediagoblin/db/migrations.py b/mediagoblin/db/migrations.py
index 6ca10b57ae63950d809ac4fd3cda8ed7543d1a5d..b5b5a026c79518dc251c4de6f3767c1481ac67eb 100644 (file)
--- a/mediagoblin/db/migrations.py
+++ b/mediagoblin/db/migrations.py
@@ -17,13 +17,18 @@
 import datetime
 import uuid
 
+import six
+
+if six.PY2:
+    import migrate
+
 from sqlalchemy import (MetaData, Table, Column, Boolean, SmallInteger,
                         Integer, Unicode, UnicodeText, DateTime,
                         ForeignKey, Date, Index)
 from sqlalchemy.exc import ProgrammingError
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.sql import and_
-from migrate.changeset.constraint import UniqueConstraint
+from sqlalchemy.schema import UniqueConstraint
 
 from mediagoblin.db.extratypes import JSONEncoded, MutationDict
 from mediagoblin.db.migration_tools import (
@@ -249,7 +254,7 @@ def mediaentry_new_slug_era(db):
     for row in db.execute(media_table.select()):
         # no slug, try setting to an id
         if not row.slug:
-            append_garbage_till_unique(row, unicode(row.id))
+            append_garbage_till_unique(row, six.text_type(row.id))
         # has "=" or ":" in it... we're getting rid of those
         elif u"=" in row.slug or u":" in row.slug:
             append_garbage_till_unique(
@@ -278,7 +283,7 @@ def unique_collections_slug(db):
                 existing_slugs[row.creator].append(row.slug)
 
     for row_id in slugs_to_change:
-        new_slug = unicode(uuid.uuid4())
+        new_slug = six.text_type(uuid.uuid4())
         db.execute(collection_table.update().
                    where(collection_table.c.id == row_id).
                    values(slug=new_slug))
@@ -837,10 +842,16 @@ def revert_username_index(db):
     """
     metadata = MetaData(bind=db.bind)
     user_table = inspect_table(metadata, "core__users")
-    indexes = {index.name: index for index in user_table.indexes}
+    indexes = dict(
+        [(index.name, index) for index in user_table.indexes])
+
+    # index from unnecessary migration
+    users_uploader_index = indexes.get(u'ix_core__users_uploader')
+    # index created from models.py after (unique=True, index=True)
+    # was set in models.py
+    users_username_index = indexes.get(u'ix_core__users_username')
 
-    if not (u'ix_core__users_uploader' in indexes or
-            u'ix_core__users_username' in indexes):
+    if users_uploader_index is None and users_username_index is None:
         # We don't need to do anything.
         # The database isn't in a state where it needs fixing
         #
@@ -864,12 +875,14 @@ def revert_username_index(db):
         # table copying.
 
         # Remove whichever of the not-used indexes are in place
-        if u'ix_core__users_uploader' in indexes:
-            index = indexes[u'ix_core__users_uploader']
-            index.drop()
-        if u'ix_core__users_username' in indexes:
-            index = indexes[u'ix_core__users_username']
-            index.drop()
+        if users_uploader_index is not None:
+            users_uploader_index.drop()
+        if users_username_index is not None:
+            users_username_index.drop()
+
+        # Given we're removing indexes then adding a unique constraint
+        # which *we know might fail*, thus probably rolling back the
+        # session, let's commit here.
         db.commit()
 
         try:
@@ -879,6 +892,6 @@ def revert_username_index(db):
             constraint.create()
         except ProgrammingError:
             # constraint already exists, no need to add
-            pass
+            db.rollback()
 
     db.commit()