From d537bfd4e1e3119188519a4b1839fdbe3f6fad39 Mon Sep 17 00:00:00 2001
From: Egor Kislitsyn <egor@kislitsyn.com>
Date: Wed, 9 Oct 2019 13:11:57 +0700
Subject: [PATCH] Add a task to re-count statuses for all users

---
 CHANGELOG.md                            |  1 +
 lib/mix/tasks/pleroma/count_statuses.ex | 22 ++++++++++++++
 test/tasks/count_statuses_test.exs      | 39 +++++++++++++++++++++++++
 3 files changed, 62 insertions(+)
 create mode 100644 lib/mix/tasks/pleroma/count_statuses.ex
 create mode 100644 test/tasks/count_statuses_test.exs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8b24db7f4..759745dd2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - OAuth: support for hierarchical permissions / [Mastodon 2.4.3 OAuth permissions](https://docs.joinmastodon.org/api/permissions/)
 - Authentication: Added rate limit for password-authorized actions / login existence checks
 - Metadata Link: Atom syndication Feed
+- Mix task to re-count statuses for all users (`mix pleroma.count_statuses`)
 
 ### Changed
 - **Breaking:** Elixir >=1.8 is now required (was >= 1.7)
diff --git a/lib/mix/tasks/pleroma/count_statuses.ex b/lib/mix/tasks/pleroma/count_statuses.ex
new file mode 100644
index 000000000..e1e8195dd
--- /dev/null
+++ b/lib/mix/tasks/pleroma/count_statuses.ex
@@ -0,0 +1,22 @@
+defmodule Mix.Tasks.Pleroma.CountStatuses do
+  @shortdoc "Re-counts statuses for all users"
+
+  use Mix.Task
+  alias Pleroma.User
+  import Ecto.Query
+
+  def run([]) do
+    Mix.Pleroma.start_pleroma()
+
+    stream =
+      User
+      |> where(local: true)
+      |> Pleroma.Repo.stream()
+
+    Pleroma.Repo.transaction(fn ->
+      Enum.each(stream, &User.update_note_count/1)
+    end)
+
+    Mix.Pleroma.shell_info("Done")
+  end
+end
diff --git a/test/tasks/count_statuses_test.exs b/test/tasks/count_statuses_test.exs
new file mode 100644
index 000000000..6035da3c3
--- /dev/null
+++ b/test/tasks/count_statuses_test.exs
@@ -0,0 +1,39 @@
+# Pleroma: A lightweight social networking server
+# Copyright © 2017-2019 Pleroma Authors <https://pleroma.social/>
+# SPDX-License-Identifier: AGPL-3.0-only
+
+defmodule Mix.Tasks.Pleroma.CountStatusesTest do
+  use Pleroma.DataCase
+
+  alias Pleroma.User
+  alias Pleroma.Web.CommonAPI
+
+  import ExUnit.CaptureIO, only: [capture_io: 1]
+  import Pleroma.Factory
+
+  test "counts statuses" do
+    user = insert(:user)
+    {:ok, _} = CommonAPI.post(user, %{"status" => "test"})
+    {:ok, _} = CommonAPI.post(user, %{"status" => "test2"})
+
+    user2 = insert(:user)
+    {:ok, _} = CommonAPI.post(user2, %{"status" => "test3"})
+
+    user = refresh_record(user)
+    user2 = refresh_record(user2)
+
+    assert %{info: %{note_count: 2}} = user
+    assert %{info: %{note_count: 1}} = user2
+
+    {:ok, user} = User.update_info(user, &User.Info.set_note_count(&1, 0))
+    {:ok, user2} = User.update_info(user2, &User.Info.set_note_count(&1, 0))
+
+    assert %{info: %{note_count: 0}} = user
+    assert %{info: %{note_count: 0}} = user2
+
+    assert capture_io(fn -> Mix.Tasks.Pleroma.CountStatuses.run([]) end) == "Done\n"
+
+    assert %{info: %{note_count: 2}} = refresh_record(user)
+    assert %{info: %{note_count: 1}} = refresh_record(user2)
+  end
+end