1 files changed, 31 insertions, 10 deletions

diff --git a/lib/pleroma/bookmark.ex b/lib/pleroma/bookmark.ex
index 187749e..1a2a63b 100755..100644
--- a/lib/pleroma/bookmark.ex
+++ b/lib/pleroma/bookmark.ex
@@ -10,6 +10,7 @@ defmodule Pleroma.Bookmark do
 
   alias Pleroma.Activity
   alias Pleroma.Bookmark
+  alias Pleroma.BookmarkFolder
   alias Pleroma.Repo
   alias Pleroma.User
 
@@ -18,33 +19,46 @@ defmodule Pleroma.Bookmark do
   schema "bookmarks" do
     belongs_to(:user, User, type: FlakeId.Ecto.CompatType)
     belongs_to(:activity, Activity, type: FlakeId.Ecto.CompatType)
+    belongs_to(:folder, BookmarkFolder, type: FlakeId.Ecto.CompatType)
 
     timestamps()
   end
 
-  @spec create(FlakeId.Ecto.CompatType.t(), FlakeId.Ecto.CompatType.t()) ::
-          {:ok, Bookmark.t()} | {:error, Changeset.t()}
-  def create(user_id, activity_id) do
+  @spec create(Ecto.UUID.t(), Ecto.UUID.t()) ::
+          {:ok, Bookmark.t()} | {:error, Ecto.Changeset.t()}
+  def create(user_id, activity_id, folder_id \\ nil) do
     attrs = %{
       user_id: user_id,
-      activity_id: activity_id
+      activity_id: activity_id,
+      folder_id: folder_id
     }
 
     %Bookmark{}
-    |> cast(attrs, [:user_id, :activity_id])
+    |> cast(attrs, [:user_id, :activity_id, :folder_id])
     |> validate_required([:user_id, :activity_id])
     |> unique_constraint(:activity_id, name: :bookmarks_user_id_activity_id_index)
-    |> Repo.insert()
+    |> Repo.insert(
+      on_conflict: [set: [folder_id: folder_id]],
+      conflict_target: [:user_id, :activity_id]
+    )
   end
 
-  @spec for_user_query(FlakeId.Ecto.CompatType.t()) :: Ecto.Query.t()
-  def for_user_query(user_id) do
+  @spec for_user_query(Ecto.UUID.t()) :: Ecto.Query.t()
+  def for_user_query(user_id, folder_id \\ nil) do
     Bookmark
     |> where(user_id: ^user_id)
+    |> maybe_filter_by_folder(folder_id)
     |> join(:inner, [b], activity in assoc(b, :activity))
     |> preload([b, a], activity: a)
   end
 
+  defp maybe_filter_by_folder(query, nil), do: query
+
+  defp maybe_filter_by_folder(query, folder_id) do
+    query
+    |> where(folder_id: ^folder_id)
+  end
+
   def get(user_id, activity_id) do
     Bookmark
     |> where(user_id: ^user_id)
@@ -52,8 +66,8 @@ defmodule Pleroma.Bookmark do
     |> Repo.one()
   end
 
-  @spec destroy(FlakeId.Ecto.CompatType.t(), FlakeId.Ecto.CompatType.t()) ::
-          {:ok, Bookmark.t()} | {:error, Changeset.t()}
+  @spec destroy(Ecto.UUID.t(), Ecto.UUID.t()) ::
+          {:ok, Bookmark.t()} | {:error, Ecto.Changeset.t()}
   def destroy(user_id, activity_id) do
     from(b in Bookmark,
       where: b.user_id == ^user_id,
@@ -62,4 +76,11 @@ defmodule Pleroma.Bookmark do
     |> Repo.one()
     |> Repo.delete()
   end
+
+  def set_folder(bookmark, folder_id) do
+    bookmark
+    |> cast(%{folder_id: folder_id}, [:folder_id])
+    |> validate_required([:folder_id])
+    |> Repo.update()
+  end
 end