From 35be52eb9febea9bfd16094f77cc14526d76632b Mon Sep 17 00:00:00 2001
From: itepechi <72330683+itepechi@users.noreply.github.com>
Date: Fri, 3 Nov 2023 05:42:14 +0900
Subject: [PATCH] Support for user search via PGroonga

---
 lib/mix/tasks/pleroma/database.ex             |  8 ++++
 lib/pleroma/user/search.ex                    | 41 +++++++++++++++++--
 ...31025184658_create_pgroonga_user_index.exs | 16 ++++++++
 3 files changed, 62 insertions(+), 3 deletions(-)
 create mode 100644 priv/repo/optional_migrations/pgroonga/20231025184658_create_pgroonga_user_index.exs

diff --git a/lib/mix/tasks/pleroma/database.ex b/lib/mix/tasks/pleroma/database.ex
index 2b5999db7..0f847b945 100644
--- a/lib/mix/tasks/pleroma/database.ex
+++ b/lib/mix/tasks/pleroma/database.ex
@@ -370,6 +370,7 @@ defmodule Mix.Tasks.Pleroma.Database do
     if pgroonga_enabled do
       shell_info("Recreate index with options #{pg_options}")
       Ecto.Adapters.SQL.query!(Pleroma.Repo, "DROP INDEX IF EXISTS object_content_pgroonga;")
+      Ecto.Adapters.SQL.query!(Pleroma.Repo, "DROP INDEX IF EXISTS nickname_name_pgroonga;")
 
       Ecto.Adapters.SQL.query!(
         Pleroma.Repo,
@@ -378,6 +379,13 @@ defmodule Mix.Tasks.Pleroma.Database do
         timeout: :infinity
       )
 
+      Ecto.Adapters.SQL.query!(
+        Pleroma.Repo,
+        "CREATE INDEX nickname_name_pgroonga ON users USING pgroonga ((ARRAY[nickname::text, name::text])) WITH (#{pg_options});",
+        [],
+        timeout: :infinity
+      )
+
       shell_info('Done.')
     else
       shell_info('PGroonga is not enabled.')
diff --git a/lib/pleroma/user/search.ex b/lib/pleroma/user/search.ex
index 08375bdb3..c0cbc8b57 100644
--- a/lib/pleroma/user/search.ex
+++ b/lib/pleroma/user/search.ex
@@ -12,6 +12,13 @@ defmodule Pleroma.User.Search do
   @limit 20
 
   def search(query_string, opts \\ []) do
+    query_type =
+      if Pleroma.Config.get([:database, :pgroonga_enabled]) do
+        :pgroonga
+      else
+        :gin
+      end
+
     resolve = Keyword.get(opts, :resolve, false)
     following = Keyword.get(opts, :following, false)
     result_limit = Keyword.get(opts, :limit, @limit)
@@ -32,7 +39,7 @@ defmodule Pleroma.User.Search do
 
     results =
       query_string
-      |> search_query(for_user, following, top_user_ids)
+      |> search_query(query_type, for_user, following, top_user_ids)
       |> Pagination.fetch_paginated(%{"offset" => offset, "limit" => result_limit}, :offset)
 
     results
@@ -85,14 +92,14 @@ defmodule Pleroma.User.Search do
     end
   end
 
-  defp search_query(query_string, for_user, following, top_user_ids) do
+  defp search_query(query_string, query_type, for_user, following, top_user_ids) do
     for_user
     |> base_query(following)
     |> filter_blocked_user(for_user)
     |> filter_invisible_users()
     |> filter_internal_users()
     |> filter_blocked_domains(for_user)
-    |> fts_search(query_string)
+    |> query_with(query_type, query_string)
     |> select_top_users(top_user_ids)
     |> trigram_rank(query_string)
     |> boost_search_rank(for_user, top_user_ids)
@@ -108,6 +115,14 @@ defmodule Pleroma.User.Search do
     )
   end
 
+  defp query_with(query, :gin, query_string) do
+    fts_search(query, query_string)
+  end
+
+  defp query_with(query, :pgroonga, query_string) do
+    pgroonga_search(query, query_string)
+  end
+
   defp fts_search(query, query_string) do
     query_string = to_tsquery(query_string)
 
@@ -160,6 +175,26 @@ defmodule Pleroma.User.Search do
     )
   end
 
+  defp pgroonga_search(query, query_string) do
+    from(
+      u in query,
+      where:
+        fragment(
+          """
+          ARRAY[?::text, ?::text]
+          &@~ (
+            ?,
+            ARRAY[2, 1],
+            'nickname_name_pgroonga'
+          )::pgroonga_full_text_search_condition
+          """,
+          u.nickname,
+          u.name,
+          ^query_string
+        )
+    )
+  end
+
   defp base_query(%User{} = user, true), do: User.get_friends_query(user)
   defp base_query(_user, _following), do: User
 
diff --git a/priv/repo/optional_migrations/pgroonga/20231025184658_create_pgroonga_user_index.exs b/priv/repo/optional_migrations/pgroonga/20231025184658_create_pgroonga_user_index.exs
new file mode 100644
index 000000000..7cf3fcca5
--- /dev/null
+++ b/priv/repo/optional_migrations/pgroonga/20231025184658_create_pgroonga_user_index.exs
@@ -0,0 +1,16 @@
+defmodule Pleroma.Repo.Migrations.CreatePgroongaUserIndex do
+  use Ecto.Migration
+
+  def up do
+    create_if_not_exists(
+      index(:users, ["(ARRAY[nickname::text, name::text])"],
+        using: :pgroonga,
+        name: :nickname_name_pgroonga
+      )
+    )
+  end
+
+  def down do
+    execute("DROP INDEX IF EXISTS nickname_name_pgroonga")
+  end
+end