Fix CASCADE in some foreignkeys, and fix migration script

kittystore/sa/alembic/versions/d1992a75f51_user_id_uuid.py

@@ -1,5 +1,7 @@
 """The User.id field is now a proper UUID
 
+Also add ONUPDATE=CASCADE to some foreign keys.
+
 Revision ID: d1992a75f51
 Revises: 31959b2e8f44
 Create Date: 2014-10-27 19:01:30.222710
@@ -17,6 +19,53 @@
 from kittystore.sa.model import Base
 
 
+FKEYS_CASCADE = ( # Foreign keys to add ONUPDATE=CASCADE to.
+    {"from_t": "email", "from_c": ["list_name"],
+     "to_t": "list", "to_c": ["name"]},
+    {"from_t": "email", "from_c": ["list_name", "thread_id"],
+     "to_t": "thread", "to_c": ["list_name", "thread_id"]},
+    {"from_t": "email_full", "from_c": ["list_name", "message_id"],
+     "to_t": "email", "to_c": ["list_name", "message_id"]},
+    {"from_t": "sender", "from_c": ["user_id"],
+     "to_t": "user", "to_c": ["id"]},
+    {"from_t": "vote", "from_c": ["user_id"],
+     "to_t": "user", "to_c": ["id"]},
+)
+
+
+def drop_user_id_fkeys():
+    op.drop_constraint("sender_user_id_fkey", "sender")
+    op.drop_constraint("vote_user_id_fkey", "vote")
+
+def create_user_id_fkeys(cascade):
+    op.create_foreign_key("sender_user_id_fkey",
+        "sender", "user", ["user_id"], ["id"],
+        onupdate=cascade, ondelete=cascade)
+    op.create_foreign_key("vote_user_id_fkey",
+        "vote", "user", ["user_id"], ["id"],
+        onupdate=cascade, ondelete=cascade)
+
+def rebuild_fkeys(cascade):
+    # Add or remove onupdate=CASCADE on some foreign keys.
+    # We need to be online or we can't reflect the constraint names.
+    if (context.is_offline_mode()
+        or op.get_context().dialect.name != 'postgresql'):
+        return
+    connection = op.get_bind()
+    md = sa.MetaData()
+    md.reflect(bind=connection)
+    for fkey in FKEYS_CASCADE:
+        keyname = None
+        for existing_fk in md.tables[fkey["from_t"]].foreign_keys:
+            if existing_fk.constraint.columns == fkey["from_c"]:
+                keyname = existing_fk.name
+        assert keyname is not None
+        op.drop_constraint(keyname, fkey["from_t"])
+        op.create_foreign_key(keyname,
+            fkey["from_t"], fkey["to_t"], fkey["from_c"], fkey["to_c"],
+            onupdate=cascade, ondelete=cascade)
+
+
 def upgrade():
     # Convert existing data into UUID strings
     if not context.is_offline_mode():
@@ -26,42 +75,92 @@ def upgrade():
         metadata = sa.MetaData()
         metadata.bind = connection
         User = Base.metadata.tables["user"].tometadata(metadata)
-        Sender = Base.metadata.tables["sender"].tometadata(metadata)
-        Vote = Base.metadata.tables["vote"].tometadata(metadata)
-        User = sa.Table("user", metadata, sa.Column("id", sa.Unicode(255), primary_key=True), extend_existing=True)
-        Sender = sa.Table("sender", metadata, sa.Column("user_id", sa.Unicode(255)), extend_existing=True)
-        Vote = sa.Table("vote", metadata, sa.Column("user_id", sa.Unicode(255), primary_key=True), extend_existing=True)
+        User = sa.Table("user", metadata,
+            sa.Column("id", sa.Unicode(255), primary_key=True),
+            extend_existing=True)
+        if connection.dialect.name != "sqlite":
+            drop_user_id_fkeys()
+            create_user_id_fkeys("CASCADE")
         transaction = connection.begin()
         for user in User.select().execute():
             try:
-                new_user_id = str(UUID(int=int(user.id)))
+                new_user_id = unicode(UUID(int=int(user.id)))
             except ValueError:
                 continue # Already converted
-            Sender.update().where(
-                Sender.c.user_id == user.id
-            ).values(user_id=new_user_id).execute()
-            Vote.update().where(
-                Vote.c.user_id == user.id
-            ).values(user_id=new_user_id).execute()
             User.update().where(
                 User.c.id == user.id
             ).values(id=new_user_id).execute()
         transaction.commit()
     # Convert to UUID for PostreSQL or to CHAR(32) for others
     if op.get_context().dialect.name == 'sqlite':
         pass # No difference between varchar and char in SQLite
+    elif op.get_context().dialect.name == 'postgresql':
+        drop_user_id_fkeys()
+        for table, col in ( ("user", "id"),
+                            ("sender", "user_id"),
+                            ("vote", "user_id") ):
+            op.execute('''
+                ALTER TABLE "{table}"
+                    ALTER COLUMN {col} TYPE UUID USING {col}::uuid
+                '''.format(table=table, col=col))
+        create_user_id_fkeys("CASCADE")
     else:
-        # This fails on PostgreSQL because it requires a 'USING' clause, and I
-        # can't find a way to generate the correct SQL statement that will not
-        # violate the foreign key constraints.
-        #for table, col in ( ("user", "id"),
-        #                    ("sender", "user_id"),
-        #                    ("votes", "user_id") ):
-        #    op.alter_column(table, col, type_=types.UUID,
-        #                    existing_type=sa.Unicode(255),
-        #                    existing_nullable=False)
-        pass
+        # Untested on other engines
+        for table, col in ( ("user", "id"),
+                            ("sender", "user_id"),
+                            ("vote", "user_id") ):
+            op.alter_column(table, col, type_=types.UUID,
+                            existing_type=sa.Unicode(255),
+                            existing_nullable=False)
+    # Now add onupdate=CASCADE to some foreign keys.
+    rebuild_fkeys("CASCADE")
 
 
 def downgrade():
-    raise RuntimeError("Downgrades are unsupported")
+    # Convert to UUID for PostreSQL or to CHAR(32) for others
+    if op.get_context().dialect.name == 'sqlite':
+        pass # No difference between varchar and char in SQLite
+    elif op.get_context().dialect.name == 'postgresql':
+        drop_user_id_fkeys()
+        for table, col in ( ("user", "id"),
+                            ("sender", "user_id"),
+                            ("vote", "user_id") ):
+            op.alter_column(table, col, type_=sa.Unicode(255),
+                            existing_type=types.UUID,
+                            existing_nullable=False)
+        # Need cascade for data conversion below, it will be removed by the
+        # last operation (or the loop on FKEYS_CASCADE if offline).
+        create_user_id_fkeys("CASCADE")
+    else:
+        # Untested on other engines
+        for table, col in ( ("user", "id"),
+                            ("sender", "user_id"),
+                            ("vote", "user_id") ):
+            op.alter_column(table, col, type_=sa.Unicode(255),
+                            existing_type=types.UUID,
+                            existing_nullable=False)
+    if not context.is_offline_mode():
+        connection = op.get_bind()
+        # Create a new MetaData instance here because the data is UUIDs and we
+        # want to convert to simple strings
+        metadata = sa.MetaData()
+        metadata.bind = connection
+        User = Base.metadata.tables["user"].tometadata(metadata)
+        User = sa.Table("user", metadata,
+            sa.Column("id", sa.Unicode(255), primary_key=True),
+            extend_existing=True)
+        transaction = connection.begin()
+        for user in User.select().execute():
+            try:
+                new_user_id = UUID(user.id).int
+            except ValueError:
+                continue # Already converted
+            User.update().where(
+                User.c.id == user.id
+            ).values(id=new_user_id).execute()
+        transaction.commit()
+        if connection.dialect.name != "sqlite":
+            drop_user_id_fkeys()
+            create_user_id_fkeys(None)
+    # Now remove onupdate=CASCADE from some foreign keys
+    rebuild_fkeys(None)

kittystore/sa/model.py

@@ -177,7 +177,9 @@ class Sender(Base):
     # TODO: rename "email" to "address"
     email = Column(Unicode(255), primary_key=True, nullable=False)
     name = Column(Unicode(255))
-    user_id = Column(UUID, ForeignKey("user.id"), index=True)
+    user_id = Column(UUID,
+        ForeignKey("user.id", onupdate="CASCADE", ondelete="CASCADE"),
+        index=True)
     emails = relationship("Email", backref="sender", cascade="all, delete-orphan")
 
 
@@ -192,8 +194,8 @@ class Email(Base):
     __tablename__ = "email"
 
     list_name = Column(Unicode(255),
-                       ForeignKey("list.name", ondelete="CASCADE"),
-                       primary_key=True, nullable=False, index=True)
+        ForeignKey("list.name", onupdate="CASCADE", ondelete="CASCADE"),
+        primary_key=True, nullable=False, index=True)
     message_id = Column(Unicode(255), primary_key=True, nullable=False)
     # TODO: rename to sender_address
     sender_email = Column(Unicode(255), ForeignKey("sender.email"),
@@ -316,7 +318,8 @@ def get_vote_by_user_id(self, user_id):
     ForeignKeyConstraint(
         ["list_name", "thread_id"],
         ["thread.list_name", "thread.thread_id"],
-        ondelete="CASCADE"
+        onupdate="CASCADE",
+        ondelete="CASCADE",
     ))
 # composite indexes
 Index("ix_email_list_name_message_id_hash",
@@ -343,7 +346,8 @@ class EmailFull(Base):
     ForeignKeyConstraint(
         ["list_name", "message_id"],
         ["email.list_name", "email.message_id"],
-        ondelete="CASCADE"
+        onupdate="CASCADE",
+        ondelete="CASCADE",
     ))
 
 
@@ -581,11 +585,12 @@ class Vote(Base):
     __tablename__ = "vote"
 
     list_name = Column(Unicode(255),
-                       ForeignKey("list.name", ondelete="CASCADE"),
-                       nullable=False, primary_key=True)
+        ForeignKey("list.name", onupdate="CASCADE", ondelete="CASCADE"),
+        nullable=False, primary_key=True)
     message_id = Column(Unicode(255), nullable=False, primary_key=True)
-    user_id = Column(UUID, ForeignKey("user.id"),
-                     nullable=False, primary_key=True, index=True)
+    user_id = Column(UUID,
+        ForeignKey("user.id", onupdate="CASCADE", ondelete="CASCADE"),
+        nullable=False, primary_key=True, index=True)
     value = Column(Integer, nullable=False, index=True)
     mlist = relationship("List")