3 files changed, 19 insertions, 10 deletions
diff --git a/mediagoblin/notifications/__init__.py b/mediagoblin/notifications/__init__.py
index ea468f8e..8690aae5 100644
--- a/mediagoblin/notifications/__init__.py
+++ b/mediagoblin/notifications/__init__.py
@@ -16,8 +16,8 @@
 
 import logging
 
-from mediagoblin.db.models import Notification, \
-        CommentNotification, CommentSubscription, User
+from mediagoblin.db.models import Notification, CommentSubscription, User, \
+                                  Comment, GenericModelReference
 from mediagoblin.notifications.task import email_notification_task
 from mediagoblin.notifications.tools import generate_comment_message
 
@@ -37,10 +37,10 @@ def trigger_notification(comment, media_entry, request):
         if comment.get_actor == subscription.user:
             continue
 
-        cn = CommentNotification(
+        cn = Notification(
             user_id=subscription.user_id,
-            subject_id=comment.id)
-
+        )
+        cn.obj = comment
         cn.save()
 
         if subscription.send_email:
@@ -61,9 +61,15 @@ def mark_notification_seen(notification):
 
 
 def mark_comment_notification_seen(comment_id, user):
-    notification = CommentNotification.query.filter_by(
+    comment = Comment.query.get(comment_id).comment()
+    comment_gmr = GenericModelReference.query.filter_by(
+        obj_pk=comment.id,
+        model_type=comment.__tablename__
+    ).first()
+    notification = Notification.query.filter_by(
         user_id=user.id,
-        subject_id=comment_id).first()
+        object_id=comment_gmr.id
+    ).first()
 
     _log.debug(u'Marking {0} as seen.'.format(notification))
 
diff --git a/mediagoblin/notifications/task.py b/mediagoblin/notifications/task.py
index d915212a..652b78e2 100644
--- a/mediagoblin/notifications/task.py
+++ b/mediagoblin/notifications/task.py
@@ -20,7 +20,7 @@ from celery import registry
 from celery.task import Task
 
 from mediagoblin.tools.mail import send_email
-from mediagoblin.db.models import CommentNotification
+from mediagoblin.db.models import Notification
 
 
 _log = logging.getLogger(__name__)
@@ -34,7 +34,7 @@ class EmailNotificationTask(Task):
     the web server.
     '''
     def run(self, notification_id, message):
-        cn = CommentNotification.query.filter_by(id=notification_id).first()
+        cn = Notification.query.filter_by(id=notification_id).first()
         _log.info(u'Sending notification email about {0}'.format(cn))
 
         return send_email(
diff --git a/mediagoblin/notifications/views.py b/mediagoblin/notifications/views.py
index cfe66b2e..984b9c9b 100644
--- a/mediagoblin/notifications/views.py
+++ b/mediagoblin/notifications/views.py
@@ -57,7 +57,10 @@ def mark_all_comment_notifications_seen(request):
     Marks all comment notifications seen.
     """
     for comment in get_notifications(request.user.id):
-        mark_comment_notification_seen(comment.subject_id, request.user)
+        mark_comment_notification_seen(
+            comment.obj().get_comment_link().id,
+            request.user
+        )
 
     if request.GET.get('next'):
         return redirect(request, location=request.GET.get('next'))