From 1d4bbec6b3239bb83b500a6a90e6686cb682cfac Mon Sep 17 00:00:00 2001
From: lain <lain@soykaf.club>
Date: Wed, 16 May 2018 17:55:20 +0200
Subject: [PATCH] Fix User search.

Now uses a trigram based search. This is a lot faster and gives better
results. Closes #185.
---
 lib/mix/tasks/sample_psql.eex                 |  1 +
 lib/pleroma/user.ex                           | 21 ++++++++++-------
 .../20180516144508_add_trigram_extension.exs  | 15 ++++++++++++
 ...180516154905_create_user_trigram_index.exs |  7 ++++++
 .../mastodon_api_controller_test.exs          | 23 +++++++++++++++----
 5 files changed, 54 insertions(+), 13 deletions(-)
 create mode 100644 priv/repo/migrations/20180516144508_add_trigram_extension.exs
 create mode 100644 priv/repo/migrations/20180516154905_create_user_trigram_index.exs

diff --git a/lib/mix/tasks/sample_psql.eex b/lib/mix/tasks/sample_psql.eex
index 18e322efc..bc22f166c 100644
--- a/lib/mix/tasks/sample_psql.eex
+++ b/lib/mix/tasks/sample_psql.eex
@@ -6,3 +6,4 @@ ALTER DATABASE pleroma_dev OWNER TO pleroma;
 \c pleroma_dev;
 --Extensions made by ecto.migrate that need superuser access
 CREATE EXTENSION IF NOT EXISTS citext;
+CREATE EXTENSION IF NOT EXISTS pg_trgm;
diff --git a/lib/pleroma/user.ex b/lib/pleroma/user.ex
index 207674999..399a66787 100644
--- a/lib/pleroma/user.ex
+++ b/lib/pleroma/user.ex
@@ -21,6 +21,7 @@ defmodule Pleroma.User do
     field(:local, :boolean, default: true)
     field(:info, :map, default: %{})
     field(:follower_address, :string)
+    field(:search_distance, :float, virtual: true)
     has_many(:notifications, Notification)
 
     timestamps()
@@ -399,19 +400,23 @@ def search(query, resolve) do
       User.get_or_fetch_by_nickname(query)
     end
 
-    q =
+    inner =
       from(
         u in User,
-        where:
-          fragment(
-            "(to_tsvector('english', ?) || to_tsvector('english', ?)) @@ plainto_tsquery('english', ?)",
+        select_merge: %{
+          search_distance: fragment(
+            "? <-> (? || ?)",
+            ^query,
             u.nickname,
-            u.name,
-            ^query
-          ),
-        limit: 20
+            u.name
+          )}
       )
 
+    q = from(s in subquery(inner),
+      order_by: s.search_distance,
+      limit: 20
+    )
+
     Repo.all(q)
   end
 
diff --git a/priv/repo/migrations/20180516144508_add_trigram_extension.exs b/priv/repo/migrations/20180516144508_add_trigram_extension.exs
new file mode 100644
index 000000000..f2f0fca86
--- /dev/null
+++ b/priv/repo/migrations/20180516144508_add_trigram_extension.exs
@@ -0,0 +1,15 @@
+defmodule Pleroma.Repo.Migrations.AddTrigramExtension do
+  use Ecto.Migration
+  require Logger
+
+  def up do
+    Logger.warn("ATTENTION ATTENTION ATTENTION\n")
+    Logger.warn("This will try to create the pg_trgm extension on your database. If your database user does NOT have the necessary rights, you will have to do it manually and re-run the migrations.\nYou can probably do this by running the following:\n")
+    Logger.warn("sudo -u postgres psql pleroma_dev -c \"create extension if not exists pg_trgm\"\n")
+    execute("create extension if not exists pg_trgm")
+  end
+
+  def down do
+    execute("drop extension if exists pg_trgm")
+  end
+end
diff --git a/priv/repo/migrations/20180516154905_create_user_trigram_index.exs b/priv/repo/migrations/20180516154905_create_user_trigram_index.exs
new file mode 100644
index 000000000..abfa4b3cc
--- /dev/null
+++ b/priv/repo/migrations/20180516154905_create_user_trigram_index.exs
@@ -0,0 +1,7 @@
+defmodule Pleroma.Repo.Migrations.CreateUserTrigramIndex do
+  use Ecto.Migration
+
+  def change do
+    create index(:users, ["(nickname || name) gist_trgm_ops"], name: :users_trigram_index, using: :gist)
+  end
+end
diff --git a/test/web/mastodon_api/mastodon_api_controller_test.exs b/test/web/mastodon_api/mastodon_api_controller_test.exs
index 882c92682..8d79c96b1 100644
--- a/test/web/mastodon_api/mastodon_api_controller_test.exs
+++ b/test/web/mastodon_api/mastodon_api_controller_test.exs
@@ -609,16 +609,29 @@ test "unimplemented mutes, follow_requests, blocks, domain blocks" do
 
   test "account search", %{conn: conn} do
     user = insert(:user)
-    _user_two = insert(:user, %{nickname: "shp@shitposter.club"})
+    user_two = insert(:user, %{nickname: "shp@shitposter.club"})
     user_three = insert(:user, %{nickname: "shp@heldscal.la", name: "I love 2hu"})
 
-    conn =
+    results =
+      conn
+      |> assign(:user, user)
+      |> get("/api/v1/accounts/search", %{"q" => "shp"})
+      |> json_response(200)
+
+    result_ids = for result <- results, do: result["acct"]
+
+    assert user_two.nickname in result_ids
+    assert user_three.nickname in result_ids
+
+    results =
       conn
       |> assign(:user, user)
       |> get("/api/v1/accounts/search", %{"q" => "2hu"})
+      |> json_response(200)
 
-    assert [account] = json_response(conn, 200)
-    assert account["id"] == to_string(user_three.id)
+    result_ids = for result <- results, do: result["acct"]
+
+    assert user_three.nickname in result_ids
   end
 
   test "search", %{conn: conn} do
@@ -642,7 +655,7 @@ test "search", %{conn: conn} do
 
     assert results = json_response(conn, 200)
 
-    [account] = results["accounts"]
+    [account | _] = results["accounts"]
     assert account["id"] == to_string(user_three.id)
 
     assert results["hashtags"] == []