added update unread_count for notifications

lib/pleroma/marker.ex

@@ -11,6 +11,7 @@ defmodule Pleroma.Marker do
   alias Ecto.Multi
   alias Pleroma.Repo
   alias Pleroma.User
+  alias __MODULE__
 
   @timelines ["notifications"]
 
@@ -46,6 +47,41 @@ def upsert(%User{} = user, attrs) do
     |> Repo.transaction()
   end
 
+  @spec multi_set_unread_count(Multi.t(), User.t(), String.t()) :: Multi.t()
+  def multi_set_unread_count(multi, %User{} = user, "notifications") do
+    multi
+    |> Multi.run(:counters, fn _repo, _changes ->
+      query =
+        from(q in Pleroma.Notification,
+          where: q.user_id == ^user.id,
+          select: %{
+            timeline: "notifications",
+            user_id: ^user.id,
+            unread_count: fragment("SUM( CASE WHEN seen = false THEN 1 ELSE 0 END ) as unread_count")
+          }
+        )
+
+      {:ok, Repo.one(query)}
+    end)
+    |> Multi.insert(
+      :marker,
+      fn %{counters: attrs} ->
+        Marker
+        |> struct(attrs)
+        |> Ecto.Changeset.change()
+      end,
+      returning: true,
+      on_conflict: {:replace, [:last_read_id, :unread_count]},
+      conflict_target: [:user_id, :timeline]
+    )
+  end
+
+  def set_unread_count(%User{} = user, timeline) do
+    Multi.new()
+    |> multi_set_unread_count(user, timeline)
+    |> Repo.transaction()
+  end
+
   defp get_marker(user, timeline) do
     case Repo.find_resource(get_query(user, timeline)) do
       {:ok, marker} -> %__MODULE__{marker | user: user}

lib/pleroma/notification.ex

@@ -5,7 +5,9 @@
 defmodule Pleroma.Notification do
   use Ecto.Schema
 
+  alias Ecto.Multi
   alias Pleroma.Activity
+  alias Pleroma.Marker
   alias Pleroma.Notification
   alias Pleroma.Object
   alias Pleroma.Pagination
@@ -151,25 +153,23 @@ def for_user_since(user, date) do
     |> Repo.all()
   end
 
-  def set_read_up_to(%{id: user_id} = _user, id) do
+  def set_read_up_to(%{id: user_id} = user, id) do
     query =
       from(
         n in Notification,
         where: n.user_id == ^user_id,
         where: n.id <= ^id,
         where: n.seen == false,
-        update: [
-          set: [
-            seen: true,
-            updated_at: ^NaiveDateTime.utc_now()
-          ]
-        ],
         # Ideally we would preload object and activities here
         # but Ecto does not support preloads in update_all
         select: n.id
       )
 
-    {_, notification_ids} = Repo.update_all(query, [])
+    {:ok, %{ids: {_, notification_ids}}} =
+      Multi.new()
+      |> Multi.update_all(:ids, query, set: [seen: true, updated_at: NaiveDateTime.utc_now()])
+      |> Marker.multi_set_unread_count(user, "notifications")
+      |> Repo.transaction()
 
     Notification
     |> where([n], n.id in ^notification_ids)
@@ -186,11 +186,18 @@ def set_read_up_to(%{id: user_id} = _user, id) do
     |> Repo.all()
   end
 
+  @spec read_one(User.t(), String.t()) ::
+          {:ok, Notification.t()} | {:error, Ecto.Changeset.t()} | nil
   def read_one(%User{} = user, notification_id) do
     with {:ok, %Notification{} = notification} <- get(user, notification_id) do
-      notification
-      |> changeset(%{seen: true})
-      |> Repo.update()
+      Multi.new()
+      |> Multi.update(:update, changeset(notification, %{seen: true}))
+      |> Marker.multi_set_unread_count(user, "notifications")
+      |> Repo.transaction()
+      |> case do
+        {:ok, %{update: notification}} -> {:ok, notification}
+        {:error, :update, changeset, _} -> {:error, changeset}
+      end
     end
   end
 
@@ -243,8 +250,11 @@ def create_notifications(%Activity{data: %{"to" => _, "type" => "Create"}} = act
     object = Object.normalize(activity)
 
     unless object && object.data["type"] == "Answer" do
-      users = get_notified_from_activity(activity)
-      notifications = Enum.map(users, fn user -> create_notification(activity, user) end)
+      notifications =
+        activity
+        |> get_notified_from_activity()
+        |> Enum.map(&create_notification(activity, &1))
+
       {:ok, notifications}
     else
       {:ok, []}
@@ -253,8 +263,11 @@ def create_notifications(%Activity{data: %{"to" => _, "type" => "Create"}} = act
 
   def create_notifications(%Activity{data: %{"to" => _, "type" => type}} = activity)
       when type in ["Like", "Announce", "Follow"] do
-    users = get_notified_from_activity(activity)
-    notifications = Enum.map(users, fn user -> create_notification(activity, user) end)
+    notifications =
+      activity
+      |> get_notified_from_activity
+      |> Enum.map(&create_notification(activity, &1))
+
     {:ok, notifications}
   end
 
@@ -263,8 +276,11 @@ def create_notifications(_), do: {:ok, []}
   # TODO move to sql, too.
   def create_notification(%Activity{} = activity, %User{} = user) do
     unless skip?(activity, user) do
-      notification = %Notification{user_id: user.id, activity: activity}
-      {:ok, notification} = Repo.insert(notification)
+      {:ok, %{notification: notification}} =
+        Multi.new()
+        |> Multi.insert(:notification, %Notification{user_id: user.id, activity: activity})
+        |> Marker.multi_set_unread_count(user, "notifications")
+        |> Repo.transaction()
 
       ["user", "user:notification"]
       |> Streamer.stream(notification)

test/notification_test.exs

@@ -310,6 +310,13 @@ test "it sets all notifications as read up to a specified notification ID" do
       assert n1.seen == true
       assert n2.seen == true
       assert n3.seen == false
+
+      assert %Pleroma.Marker{unread_count: 1} =
+               Pleroma.Repo.get_by(
+                 Pleroma.Marker,
+                 user_id: other_user.id,
+                 timeline: "notifications"
+               )
     end
   end