From a9b7503aa7411c121b85a63f414ffdd0365e9928 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Fri, 6 Sep 2019 16:18:29 -0500
Subject: [PATCH] pools: allow searching pools by post id or post count.

All pools containing post #1:

    https://danbooru.donmai.us/pools?search[post_ids_include]=1

All pools containing either post #1 or #2:

    https://danbooru.donmai.us/pools?search[post_ids_include]=1,2
    https://danbooru.donmai.us/pools?search[post_ids_include]=1+2

Pools with 1-100 posts:

    https://danbooru.donmai.us/pools?search[post_id_count]=1..100

Pools with no posts (empty pools):

    https://danbooru.donmai.us/pools?search[post_id_count]=0
---
 app/models/application_record.rb | 36 ++++++++++++++++++++++++++++++--
 app/models/pool.rb               |  2 +-
 test/unit/pool_test.rb           | 19 +++++++++++++++++
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/app/models/application_record.rb b/app/models/application_record.rb
index 6d5ab6320..eed756081 100644
--- a/app/models/application_record.rb
+++ b/app/models/application_record.rb
@@ -35,6 +35,18 @@ class ApplicationRecord < ActiveRecord::Base
         where.not("#{qualified_column_for(attr)} ~ ?", "(?e)" + value)
       end
 
+      def where_array_includes(attr, values)
+        where("#{qualified_column_for(attr)} && ARRAY[?]", values)
+      end
+
+      def where_array_count(attr, value)
+        relation = all
+        qualified_column = "cardinality(#{qualified_column_for(attr)})"
+        parsed_range = Tag.parse_helper(value, :integer)
+
+        PostQueryBuilder.new(nil).add_range_relation(parsed_range, qualified_column, relation)
+      end
+
       def search_boolean_attribute(attribute, params)
         return all unless params[attribute]
 
@@ -81,7 +93,8 @@ class ApplicationRecord < ActiveRecord::Base
       end
 
       def search_attribute(name, params)
-        type = type_for_attribute(name).type || reflect_on_association(name)&.class_name
+        column = column_for_attribute(name)
+        type = column.type || reflect_on_association(name)&.class_name
 
         case type
         when "User"
@@ -93,7 +106,11 @@ class ApplicationRecord < ActiveRecord::Base
         when :boolean
           search_boolean_attribute(name, params)
         when :integer, :datetime
-          numeric_attribute_matches(name, params[name])
+          if column.array?
+            search_array_attribute(name, type, params)
+          else
+            numeric_attribute_matches(name, params[name])
+          end
         else
           raise NotImplementedError, "unhandled attribute type"
         end
@@ -149,6 +166,21 @@ class ApplicationRecord < ActiveRecord::Base
         relation
       end
 
+      def search_array_attribute(name, type, params)
+        relation = all
+
+        if params[:"#{name}_include"] && type == :integer
+          items = params[:"#{name}_include"].to_s.scan(/\d+/).map(&:to_i)
+          relation = relation.where_array_includes(name, items)
+        end
+
+        if params[:"#{name.to_s.singularize}_count"]
+          relation = relation.where_array_count(name, params[:"#{name.to_s.singularize}_count"])
+        end
+
+        relation
+      end
+
       def apply_default_order(params)
         if params[:order] == "custom"
           parse_ids = Tag.parse_helper(params[:id])
diff --git a/app/models/pool.rb b/app/models/pool.rb
index ce1251362..1517da54e 100644
--- a/app/models/pool.rb
+++ b/app/models/pool.rb
@@ -57,7 +57,7 @@ class Pool < ApplicationRecord
     def search(params)
       q = super
 
-      q = q.search_attributes(params, :creator, :is_active, :is_deleted, :name, :description)
+      q = q.search_attributes(params, :creator, :is_active, :is_deleted, :name, :description, :post_ids)
       q = q.text_attribute_matches(:description, params[:description_matches])
 
       if params[:name_matches].present?
diff --git a/test/unit/pool_test.rb b/test/unit/pool_test.rb
index 3e4ccb2c7..5da15deba 100644
--- a/test/unit/pool_test.rb
+++ b/test/unit/pool_test.rb
@@ -59,6 +59,25 @@ class PoolTest < ActiveSupport::TestCase
       assert_equal(@pool.id, Pool.find_by_name("test pool").id)
       assert_equal(@pool.id, Pool.search(name_matches: "test pool").first.id)
     end
+
+    should "find pools by post id" do
+      @pool1 = create(:pool, name: "pool1")
+      @pool2 = create(:pool, name: "pool2")
+      @post1 = create(:post, tag_string: "pool:pool1")
+      @post2 = create(:post, tag_string: "pool:pool2")
+
+      assert_equal([@pool1.id], Pool.search(post_ids_include: @post1.id).pluck(:id))
+      assert_equal([@pool2.id, @pool1.id], Pool.search(post_ids_include: "#{@post1.id} #{@post2.id}").pluck(:id))
+    end
+
+    should "find pools by post id count" do
+      @pool1 = create(:pool, name: "pool1")
+      @pool2 = create(:pool, name: "pool2")
+      @post1 = create(:post, tag_string: "pool:pool1")
+      @post2 = create(:post, tag_string: "pool:pool1")
+
+      assert_equal([@pool1.id], Pool.search(post_id_count: 2).pluck(:id))
+    end
   end
 
   context "Creating a pool" do