Improve notification model

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

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)

db/migrate/20161203164520_add_from_account_id_to_notifications.rb

@@ -1,6 +1,6 @@
 class AddFromAccountIdToNotifications < ActiveRecord::Migration[5.0]
   def up
-    add_column :notifications, :from_account_id, :integer, null: false, default: 1
+    add_column :notifications, :from_account_id, :integer
 
     Notification.where(activity_type: 'Status').update_all('from_account_id = (SELECT statuses.account_id FROM notifications AS notifications1 INNER JOIN statuses ON notifications1.activity_id = statuses.id WHERE notifications1.activity_type = \'Status\' AND notifications1.id = notifications.id)')
     Notification.where(activity_type: 'Mention').update_all('from_account_id = (SELECT statuses.account_id FROM notifications AS notifications1 INNER JOIN mentions ON notifications1.activity_id = mentions.id INNER JOIN statuses ON mentions.status_id = statuses.id WHERE notifications1.activity_type = \'Mention\' AND notifications1.id = notifications.id)')

db/schema.rb

@@ -100,9 +100,9 @@ ActiveRecord::Schema.define(version: 20161203164520) do
     t.integer  "account_id"
     t.integer  "activity_id"
     t.string   "activity_type"
-    t.datetime "created_at",                  null: false
-    t.datetime "updated_at",                  null: false
-    t.integer  "from_account_id", default: 1, null: false
+    t.datetime "created_at",      null: false
+    t.datetime "updated_at",      null: false
+    t.integer  "from_account_id"
     t.index ["account_id", "activity_id", "activity_type"], name: "account_activity", unique: true, using: :btree
   end
 

spec/controllers/admin/accounts_controller_spec.rb

@@ -13,8 +13,10 @@ RSpec.describe Admin::AccountsController, type: :controller do
   end
 
   describe 'GET #show' do
+    let(:account) { Fabricate(:account, username: 'bob') }
+
     it 'returns http success' do
-      get :show, params: { id: 1 }
+      get :show, params: { id: account.id }
       expect(response).to have_http_status(:success)
     end
   end