[Elasticsearch] Implement Response

app/controllers/admin/reowner_controller.rb

@@ -17,7 +17,7 @@ module Admin
       end
 
       moved_post_ids = []
-      Post.tag_match("user:!#{@old_user.id} #{query}").limit(300).records.each do |p|
+      Post.tag_match("user:!#{@old_user.id} #{query}").limit(300).each do |p|
         moved_post_ids << p.id
         p.do_not_version_changes = true
         p.update({ uploader_id: @new_user.id })

app/controllers/admin/stuck_dnp_controller.rb

@@ -16,7 +16,7 @@ module Admin
 
       dnp_tags = %w[avoid_posting conditional_dnp]
       post_ids = []
-      Post.tag_match_system("#{query} ~avoid_posting ~conditional_dnp").limit(1000).records.each do |p|
+      Post.tag_match_system("#{query} ~avoid_posting ~conditional_dnp").limit(1000).each do |p|
         previous_tags = p.fetch_tags(*dnp_tags)
 
         p.do_not_version_changes = true

app/controllers/posts_controller.rb

@@ -89,7 +89,7 @@ class PostsController < ApplicationController
 
   def random
     tags = params[:tags] || ''
-    @post = Post.tag_match(tags + " order:random").limit(1).records[0]
+    @post = Post.tag_match(tags + " order:random").limit(1).first
     raise ActiveRecord::RecordNotFound if @post.nil?
     respond_with(@post) do |format|
       format.html { redirect_to post_path(@post, :tags => params[:tags]) }

app/logical/danbooru/paginator/document_store_extensions.rb

@@ -57,7 +57,7 @@ module Danbooru
       end
     end
 
-    module ElasticsearchExtensions
+    module DocumentStoreExtensions
       include BaseExtension
 
       def paginate(page, options)
@@ -71,7 +71,7 @@ module Danbooru
           current_page: current_page,
         }
 
-        PaginatedArray.new(records(includes: options[:includes]).to_a, new_opts)
+        PaginatedArray.new(records(includes: options[:includes]), new_opts)
       end
 
       def paginate_numbered

app/logical/document_store/model.rb

@@ -36,7 +36,7 @@ module DocumentStore
 
       def document_store_search(body)
         search = SearchRequest.new({ index: __elasticsearch__.index_name, body: body }, document_store_client)
-        Elasticsearch::Model::Response::Response.new(self, search)
+        Response.new(self, search)
       end
 
       def document_store_create_index!(delete_existing: false)

app/logical/document_store/response.rb

@@ -0,0 +1,33 @@
+module DocumentStore
+  class Response
+    include Danbooru::Paginator::DocumentStoreExtensions
+
+    delegate_missing_to :records
+    attr_reader :klass, :search
+
+    def initialize(klass, search)
+      @klass = klass
+      @search = search
+    end
+
+    def response
+      @response ||= @search.execute!
+    end
+
+    def hits
+      response["hits"]["hits"]
+    end
+
+    def relation
+      klass.where(id: hits.pluck("_id"))
+    end
+
+    def records(includes: nil)
+      @records ||= begin
+        sql_records = relation
+        sql_records = sql_records.includes(includes) if includes
+        sql_records.records.sort_by { |sql_record| hits.index { |hit| hit["_id"] == sql_record.id.to_s } }
+      end
+    end
+  end
+end

app/models/artist.rb

@@ -240,7 +240,7 @@ class Artist < ApplicationRecord
       Cache.fetch("artist-domains-#{id}", expires_in: 1.day) do
         re = /\.(png|jpeg|jpg|webm|mp4)$/m
         counted = Hash.new(0)
-        sources = Post.tag_match(name, resolve_aliases: false).limit(100).records.pluck(:source).each do |source_string|
+        sources = Post.tag_match(name, resolve_aliases: false).limit(100).pluck(:source).each do |source_string|
           sources = source_string.split("\n")
           # try to filter out direct file urls
           domains = sources.filter {|s| !re.match?(s) }.map do |x|

app/models/post.rb

@@ -1430,7 +1430,7 @@ class Post < ApplicationRecord
     end
 
     def sample(query, sample_size)
-      tag_match_system("#{query} order:random", free_tags_count: 1).limit(sample_size).records
+      tag_match_system("#{query} order:random", free_tags_count: 1).limit(sample_size).relation
     end
 
     # unflattens the tag_string into one tag per row.

app/models/takedown.rb

@@ -96,7 +96,7 @@ class Takedown < ApplicationRecord
     end
 
     def add_posts_by_tags!(tag_string)
-      new_ids = Post.tag_match_system("#{tag_string} -status:deleted").limit(1000).results.map(&:id)
+      new_ids = Post.tag_match_system("#{tag_string} -status:deleted").limit(1000).pluck(&:id)
       add_posts_by_ids!(new_ids.join(" "))
     end
 

app/presenters/user_presenter.rb

@@ -52,7 +52,7 @@ class UserPresenter
   end
 
   def uploads
-    Post.tag_match("user:#{user.name}").limit(6).records
+    Post.tag_match("user:#{user.name}").limit(6)
   end
 
   def has_uploads?

config/initializers/elasticsearch.rb

@@ -1,4 +1 @@
 Elasticsearch::Model.client = Elasticsearch::Client.new host: Danbooru.config.elasticsearch_host
-Rails.configuration.to_prepare do
-  Elasticsearch::Model::Response::Response.include(Danbooru::Paginator::ElasticsearchExtensions)
-end

test/unit/post_test.rb

@@ -2,7 +2,7 @@ require 'test_helper'
 
 class PostTest < ActiveSupport::TestCase
   def assert_tag_match(posts, query)
-    assert_equal(posts.map(&:id), Post.tag_match(query).records.pluck(:id))
+    assert_equal(posts.map(&:id), Post.tag_match(query).pluck(:id))
   end
 
   setup do