Improve notification model

2 files changed, 20 insertions, 4 deletions

diff --git a/app/models/notification.rb b/app/models/notification.rb
index 2d48abce5..9d076ad41 100644
--- a/app/models/notification.rb
+++ b/app/models/notification.rb
@@ -17,10 +17,12 @@ class Notification < ApplicationRecord
 
   STATUS_INCLUDES = [:account, :stream_entry, :media_attachments, :tags, mentions: :account, reblog: [:stream_entry, :account, :media_attachments, :tags, mentions: :account]].freeze
 
+  scope :cache_ids, -> { select(:id, :updated_at, :activity_type, :activity_id) }
+
   cache_associated :from_account, status: STATUS_INCLUDES, mention: [status: STATUS_INCLUDES], favourite: [:account, status: STATUS_INCLUDES], follow: :account
 
-  def activity
-    send(activity_type.downcase)
+  def activity(eager_loaded = true)
+    eager_loaded ? send(activity_type.downcase) : super
   end
 
   def type
@@ -51,4 +53,18 @@ class Notification < ApplicationRecord
       end
     end
   end
+
+  after_initialize :set_from_account
+  before_validation :set_from_account
+
+  private
+
+  def set_from_account
+    case activity_type
+    when 'Status', 'Follow', 'Favourite'
+      self.from_account_id = activity(false)&.account_id
+    when 'Mention'
+      self.from_account_id = activity(false)&.status&.account_id
+    end
+  end
 end
diff --git a/app/models/status.rb b/app/models/status.rb
index f6d8fc9bb..1f5cf9b46 100644
--- a/app/models/status.rb
+++ b/app/models/status.rb
@@ -94,11 +94,11 @@ class Status < ApplicationRecord
 
   class << self
     def as_home_timeline(account)
-      where(account: [account] + account.following).with_includes
+      where(account: [account] + account.following)
     end
 
     def as_mentions_timeline(account)
-      where(id: Mention.where(account: account).pluck(:status_id)).with_includes
+      where(id: Mention.where(account: account).select(:status_id))
     end
 
     def as_public_timeline(account = nil)