1 files changed, 72 insertions, 0 deletions
diff --git a/mediagoblin/db/migrations/versions/4066b9f8b84a_use_comment_link_ids_notifications.py b/mediagoblin/db/migrations/versions/4066b9f8b84a_use_comment_link_ids_notifications.py
new file mode 100644
index 00000000..0c023c59
--- /dev/null
+++ b/mediagoblin/db/migrations/versions/4066b9f8b84a_use_comment_link_ids_notifications.py
@@ -0,0 +1,72 @@
+"""use_comment_link_ids_notifications
+
+Revision ID: 4066b9f8b84a
+Revises: 8429e33fdf7
+Create Date: 2016-02-29 11:46:13.511318
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = '4066b9f8b84a'
+down_revision = '8429e33fdf7'
+
+from alembic import op
+from sqlalchemy import MetaData
+from mediagoblin.db.migration_tools import inspect_table
+
+def upgrade():
+    """"
+    This replaces the Notification.obj with the ID of the Comment (i.e. comment
+    link) ID instead of the TextComment object.
+    """
+    db = op.get_bind()
+    metadata = MetaData(bind=db)
+    notification_table = inspect_table(metadata, "core__notifications")
+    comment_table = inspect_table(metadata, "core__comment_links")
+
+    # Get the notifications.
+    notifications = list(db.execute(notification_table.select()))
+
+    # Iterate through all the notifications
+    for notification in notifications:
+        # Lookup the Comment link object from the notification's ID
+        comment_link = db.execute(comment_table.select().where(
+            comment_table.c.comment_id == notification.object_id
+        )).first()
+
+        # Okay now we need to update the notification with the ID of the link
+        # rather than the ID of TextComment object.
+        db.execute(notification_table.update().values(
+            object_id=comment_link.id
+        ).where(
+            notification_table.c.id == notification.id
+        ))
+
+
+def downgrade():
+    """
+    This puts back the TextComment ID for the notification.object_id field
+    where we're using the Comment object (i.e. the comment link ID)
+    """
+    db = op.get_bind()
+    metadata = MetaData(bind=db)
+    notification_table = inspect_table(metadata, "core__notifications")
+    comment_table = inspect_table(metadata, "core__comment_links")
+
+    # Notificaitons
+    notifications = list(db.execute(notification_table.select()))
+
+    # Iterate through all the notifications
+    for notification in notifications:
+        # Lookup the Comment link object from the notification's ID
+        comment_link = db.execute(comment_table.select().where(
+            comment_table.c.id == notification.object_id
+        )).first()
+
+        # Update the notification with the TextComment (i.e. the comment object)
+        db.execute(notification_table.update().values(
+            object_id=comment_link.comment_id
+        ).where(
+            notification_table.c.id == notification.id
+        ))
+