From 96702e7f6727d9c688b2c6e2527c4a5bfefff886 Mon Sep 17 00:00:00 2001
From: Eugen Rochko <eugen@zeonfederated.com>
Date: Sun, 18 Aug 2019 14:55:03 +0200
Subject: [PATCH] Add `tootctl cache recount` command (#11597)

---
 app/models/concerns/account_counters.rb |  3 +-
 app/models/status.rb                    |  7 ++--
 lib/mastodon/cache_cli.rb               | 45 +++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/app/models/concerns/account_counters.rb b/app/models/concerns/account_counters.rb
index 3581df8dd..6e25e1905 100644
--- a/app/models/concerns/account_counters.rb
+++ b/app/models/concerns/account_counters.rb
@@ -26,7 +26,8 @@ module AccountCounters
   private
 
   def save_account_stat
-    return unless account_stat&.changed?
+    return unless association(:account_stat).loaded? && account_stat&.changed?
+
     account_stat.save
   end
 end
diff --git a/app/models/status.rb b/app/models/status.rb
index 23682c84b..0538c4e9e 100644
--- a/app/models/status.rb
+++ b/app/models/status.rb
@@ -388,13 +388,16 @@ class Status < ApplicationRecord
     end
   end
 
+  def status_stat
+    super || build_status_stat
+  end
+
   private
 
   def update_status_stat!(attrs)
     return if marked_for_destruction? || destroyed?
 
-    record = status_stat || build_status_stat
-    record.update(attrs)
+    status_stat.update(attrs)
   end
 
   def store_uri
diff --git a/lib/mastodon/cache_cli.rb b/lib/mastodon/cache_cli.rb
index e9b6667b3..5b0eea91b 100644
--- a/lib/mastodon/cache_cli.rb
+++ b/lib/mastodon/cache_cli.rb
@@ -15,5 +15,50 @@ module Mastodon
       Rails.cache.clear
       say('OK', :green)
     end
+
+    desc 'recount TYPE', 'Update hard-cached counters'
+    long_desc <<~LONG_DESC
+      Update hard-cached counters of TYPE by counting referenced
+      records from scratch. TYPE can be "accounts" or "statuses".
+
+      It may take a very long time to finish, depending on the
+      size of the database.
+    LONG_DESC
+    def recount(type)
+      processed = 0
+
+      case type
+      when 'accounts'
+        Account.local.includes(:account_stat).find_each do |account|
+          account_stat                 = account.account_stat
+          account_stat.following_count = account.active_relationships.count
+          account_stat.followers_count = account.passive_relationships.count
+          account_stat.statuses_count  = account.statuses.where.not(visibility: :direct).count
+
+          account_stat.save if account_stat.changed?
+
+          processed += 1
+          say('.', :green, false)
+        end
+      when 'statuses'
+        Status.includes(:status_stat).find_each do |status|
+          status_stat                  = status.status_stat
+          status_stat.replies_count    = status.replies.where.not(visibility: :direct).count
+          status_stat.reblogs_count    = status.reblogs.count
+          status_stat.favourites_count = status.favourites.count
+
+          status_stat.save if status_stat.changed?
+
+          processed += 1
+          say('.', :green, false)
+        end
+      else
+        say("Unknown type: #{type}", :red)
+        exit(1)
+      end
+
+      say
+      say("OK, recounted #{processed} records", :green)
+    end
   end
 end