iqdb: update API client to use new version of IQDB.

Replace the old IQDB API client with a new client for the new forked
version of IQDB at https://github.com/danbooru/iqdb.

Changes:

* The /iqdb_queries endpoint now returns `hash` and `signature` fields.
  The `signature` is the full decoded Haar signature, while the `hash`
  is a encoded version of the signature.
* The /iqdb_queries endpoint no longer returns `width` and `height`
  fields in the response (these were always 128x128).
* We no longer need the IQDBs frontend server, now we talk to the IQDB
  instance directly.
* We no longer send add/remove image commands to IQDB through AWS SQS,
  now we send them to IQDB directly. They are sent in a delayed job so
  that if IQDB is down, uploading images is still possible, the add
  image commands will just get queued up.
* Fix a bug where regenerating an image's thumbnails didn't regenerate
  IQDB, because IQDB silently ignored add image commands when the image
  already existed in the database.

.env

@@ -78,8 +78,6 @@
 # export DANBOORU_AWS_ACCESS_KEY_ID=
 # export DANBOORU_AWS_SECRET_ACCESS_KEY=
 # export DANBOORU_AWS_SQS_REGION=
-# export DANBOORU_IQDBS_AUTH_KEY=
-# export DANBOORU_IQDBS_SERVER=
 # export DANBOORU_CLOUDFLARE_KEY=
 # export DANBOORU_CLOUDFLARE_EMAIL=
 # export DANBOORU_CLOUDFLARE_ZONE=

Procfile

@@ -1,5 +1,2 @@
 unicorn: bin/rails server -p 3000
 jobs: bin/rake jobs:work
-recommender: bundle exec ruby script/mock_services/recommender.rb
-iqdbs: bundle exec ruby script/mock_services/iqdbs.rb
-reportbooru: bundle exec ruby script/mock_services/reportbooru.rb

README.md

@@ -88,7 +88,7 @@ The following features requires a Google API account:
 
 ### IQDB Service
 
-IQDB integration is delegated to the [IQDBS service](https://github.com/r888888888/iqdbs).
+IQDB integration is delegated to the [IQDB service](https://github.com/danbooru/iqdb).
 
 ### Archive Service
 

app/controllers/iqdb_queries_controller.rb

@@ -5,7 +5,8 @@ class IqdbQueriesController < ApplicationController
     # XXX allow bare search params for backwards compatibility.
     search_params.merge!(params.slice(:url, :image_url, :file_url, :post_id, :limit, :similarity, :high_similarity).permit!)
 
-    @high_similarity_matches, @low_similarity_matches, @matches = IqdbProxy.new.search(search_params)
+    iqdb_params = search_params.to_h.symbolize_keys
+    @high_similarity_matches, @low_similarity_matches, @matches = IqdbClient.new.search(**iqdb_params)
     respond_with(@matches, template: "iqdb_queries/show")
   end
 

app/controllers/mock_services_controller.rb

@@ -31,7 +31,7 @@ class MockServicesController < ApplicationController
     render json: @data
   end
 
-  def iqdbs_similar
+  def iqdb_query
     @data = posts.map { |post| { post_id: post.id, score: rand(0..100)} }
     render json: @data
   end

app/jobs/iqdb_add_post_job.rb

@@ -0,0 +1,5 @@
+class IqdbAddPostJob < ApplicationJob
+  def perform(post)
+    IqdbClient.new.add_post(post)
+  end
+end

app/jobs/iqdb_remove_post_job.rb

@@ -0,0 +1,5 @@
+class IqdbRemovePostJob < ApplicationJob
+  def perform(post_id)
+    IqdbClient.new.remove(post_id)
+  end
+end

app/logical/iqdb_client.rb

@@ -0,0 +1,86 @@
+class IqdbClient
+  class Error < StandardError; end
+  attr_reader :iqdb_url, :http
+
+  def initialize(iqdb_url: Danbooru.config.iqdb_url.to_s, http: Danbooru::Http.new)
+    @iqdb_url = iqdb_url.chomp("/")
+    @http = http
+  end
+
+  concerning :QueryMethods do
+    def search(post_id: nil, file: nil, url: nil, image_url: nil, file_url: nil, similarity: 0.0, high_similarity: 65.0, limit: 20)
+      limit = limit.to_i.clamp(1, 1000)
+      similarity = similarity.to_f.clamp(0.0, 100.0)
+      high_similarity = high_similarity.to_f.clamp(0.0, 100.0)
+
+      if file.blank?
+        if url.present?
+          file = download(url, :preview_url)
+        elsif image_url.present?
+          file = download(image_url, :url)
+        elsif file_url.present?
+          file = download(file_url, :image_url)
+        elsif post_id.present?
+          file = Post.find(post_id).file(:preview)
+        else
+          return [[], [], []]
+        end
+      end
+
+      results = query(file, limit: limit)
+      results = results.select { |result| result["score"] >= similarity }.take(limit)
+      matches = decorate_posts(results)
+      high_similarity_matches, low_similarity_matches = matches.partition { |match| match["score"] >= high_similarity }
+
+      [high_similarity_matches, low_similarity_matches, matches]
+    ensure
+      file.try(:close)
+    end
+
+    def download(url, type)
+      strategy = Sources::Strategies.find(url)
+      download_url = strategy.send(type)
+      file = strategy.download_file!(download_url)
+      file
+    end
+
+    def decorate_posts(json)
+      post_ids = json.map { |match| match["post_id"] }
+      posts = Post.where(id: post_ids).group_by(&:id).transform_values(&:first)
+
+      json.map do |match|
+        post = posts.fetch(match["post_id"], nil)
+        match.with_indifferent_access.merge(post: post) if post
+      end.compact
+    end
+  end
+
+  def add_post(post)
+    return unless post.has_preview?
+    preview_file = post.file(:preview)
+    add(post.id, preview_file)
+  end
+
+  concerning :HttpMethods do
+    def query(file, limit: 20)
+      file = HTTP::FormData::File.new(file)
+      request(:post, "query", form: { file: file }, params: { limit: limit })
+    end
+
+    def add(post_id, file)
+      file = HTTP::FormData::File.new(file)
+      request(:post, "images/#{post_id}", form: { file: file })
+    end
+
+    def remove(post_id)
+      request(:delete, "images/#{post_id}")
+    end
+
+    def request(method, url, **options)
+      return [] if iqdb_url.blank? # do nothing if iqdb isn't configured
+      response = http.timeout(30).send(method, "#{iqdb_url}/#{url}", **options)
+      raise Error, "IQDB error: #{response.parse}" if response.status != 200
+      response.parse
+    end
+  end
+end

app/logical/iqdb_proxy.rb

@@ -1,77 +0,0 @@
-class IqdbProxy
-  class Error < StandardError; end
-  attr_reader :http, :iqdbs_server
-
-  def initialize(http: Danbooru::Http.new, iqdbs_server: Danbooru.config.iqdbs_server)
-    @iqdbs_server = iqdbs_server
-    @http = http
-  end
-
-  def enabled?
-    iqdbs_server.present?
-  end
-
-  def download(url, type)
-    strategy = Sources::Strategies.find(url)
-    download_url = strategy.send(type)
-    file = strategy.download_file!(download_url)
-    file
-  end
-
-  def search(params)
-    limit = params[:limit]&.to_i&.clamp(1, 1000) || 20
-    similarity = params[:similarity]&.to_f&.clamp(0.0, 100.0) || 0.0
-    high_similarity = params[:high_similarity]&.to_f&.clamp(0.0, 100.0) || 65.0
-
-    if params[:file].present?
-      file = params[:file]
-      results = query(file: file, limit: limit)
-    elsif params[:url].present?
-      file = download(params[:url], :preview_url)
-      results = query(file: file, limit: limit)
-    elsif params[:image_url].present?
-      file = download(params[:image_url], :url)
-      results = query(file: file, limit: limit)
-    elsif params[:file_url].present?
-      file = download(params[:file_url], :image_url)
-      results = query(file: file, limit: limit)
-    elsif params[:post_id].present?
-      url = Post.find(params[:post_id]).preview_file_url
-      results = query(url: url, limit: limit)
-    else
-      results = []
-    end
-
-    results = results.select { |result| result["score"] >= similarity }.take(limit)
-    matches = decorate_posts(results)
-    high_similarity_matches, low_similarity_matches = matches.partition { |match| match["score"] >= high_similarity }
-
-    [high_similarity_matches, low_similarity_matches, matches]
-  ensure
-    file.try(:close)
-  end
-
-  def query(file: nil, url: nil, limit: 20)
-    raise NotImplementedError, "the IQDBs service isn't configured" unless enabled?
-
-    file = HTTP::FormData::File.new(file) if file
-    form = { file: file, url: url, limit: limit }.compact
-    response = http.timeout(30).post("#{iqdbs_server}/similar", form: form)
-
-    raise Error, "IQDB error: #{response.status}" if response.status != 200
-    raise Error, "IQDB error: #{response.parse["error"]}" if response.parse.is_a?(Hash)
-    raise Error, "IQDB error: #{response.parse.first}" if response.parse.try(:first).is_a?(String)
-
-    response.parse
-  end
-
-  def decorate_posts(json)
-    post_ids = json.map { |match| match["post_id"] }
-    posts = Post.where(id: post_ids).group_by(&:id).transform_values(&:first)
-
-    json.map do |match|
-      post = posts.fetch(match["post_id"], nil)
-      match.with_indifferent_access.merge(post: post) if post
-    end.compact
-  end
-end

app/logical/upload_service.rb

@@ -85,6 +85,8 @@ class UploadService
       )
     end
 
+    @post.update_iqdb
+
     upload.update(status: "completed", post_id: @post.id)
 
     @post

app/logical/upload_service/replacer.rb

@@ -128,7 +128,7 @@ class UploadService
       replacement.save!
       post.save!
 
-      post.update_iqdb_async
+      post.update_iqdb
     end
 
     def purge_cached_urls(post)

app/models/post.rb

@@ -36,8 +36,6 @@ class Post < ApplicationRecord
   after_save :update_parent_on_save
   after_save :apply_post_metatags
   after_commit :delete_files, :on => :destroy
-  after_commit :remove_iqdb_async, :on => :destroy
-  after_commit :update_iqdb_async, :on => :create
 
   belongs_to :approver, class_name: "User", optional: true
   belongs_to :uploader, :class_name => "User", :counter_cache => "post_upload_count"
@@ -977,6 +975,8 @@ class Post < ApplicationRecord
           update_parent_on_destroy
         end
       end
+
+      remove_iqdb # this is non-transactional
     end
 
     def ban!
@@ -1319,16 +1319,15 @@ class Post < ApplicationRecord
 
     def regenerate!(category, user)
       if category == "iqdb"
-        update_iqdb_async
+        update_iqdb
 
         ModAction.log("<@#{user.name}> regenerated IQDB for post ##{id}", :post_regenerate_iqdb, user)
       else
         media_file = MediaFile.open(file, frame_data: pixiv_ugoira_frame_data)
         UploadService::Utils.process_resizes(self, nil, id, media_file: media_file)
 
-        # XXX This may be racy; the thumbnail may not be purged from Cloudflare by the time IQDB tries to download it.
         purge_cached_urls!
-        update_iqdb_async
+        update_iqdb
 
         ModAction.log("<@#{user.name}> regenerated image samples for post ##{id}", :post_regenerate, user)
       end
@@ -1340,33 +1339,15 @@ class Post < ApplicationRecord
     end
   end
 
-  module IqdbMethods
+  concerning :IqdbMethods do
-    extend ActiveSupport::Concern
+    def update_iqdb
-
+      # performs IqdbClient.new.add_post(post)
-    module ClassMethods
+      IqdbAddPostJob.perform_later(self)
-      def iqdb_sqs_service
-        SqsService.new(Danbooru.config.aws_sqs_iqdb_url)
     end
 
-      def iqdb_enabled?
+    def remove_iqdb
-        Danbooru.config.aws_sqs_iqdb_url.present?
+      # performs IqdbClient.new.remove(id)
-      end
+      IqdbRemovePostJob.perform_later(id)
-
-      def remove_iqdb(post_id)
-        if iqdb_enabled?
-          iqdb_sqs_service.send_message("remove\n#{post_id}")
-        end
-      end
-    end
-
-    def update_iqdb_async
-      if Post.iqdb_enabled? && has_preview?
-        Post.iqdb_sqs_service.send_message("update\n#{id}\n#{preview_file_url}")
-      end
-    end
-
-    def remove_iqdb_async
-      Post.remove_iqdb(id)
     end
   end
 
@@ -1460,7 +1441,6 @@ class Post < ApplicationRecord
   include ApiMethods
   extend SearchMethods
   include PixivMethods
-  include IqdbMethods
   include ValidationMethods
 
   has_bit_flags ["has_embedded_notes", "has_cropped"]

app/views/uploads/_related_posts.html.erb

@@ -16,7 +16,7 @@
   </div>
 <% end %>
 
-<% if Danbooru.config.iqdbs_server %>
+<% if Danbooru.config.iqdb_url %>
   <% if params[:url] %>
     <div class="input" id="iqdb-similar">
       <p><em>Loading similar...</em></p>

config/danbooru_default_config.rb

@@ -481,11 +481,12 @@ module Danbooru
     def reportbooru_key
     end
 
-    # The URL for the IQDBs server (https://github.com/evazion/iqdbs).
+    # The URL for the IQDB server (https://github.com/danbooru/iqdb). Optional.
-    # Optional. Used for dupe detection and reverse image searches.
+    # Used for dupe detection and reverse image searches. Set this to
-    # Set to http://localhost/mock/iqdbs to enable a fake iqdb server for
+    # http://localhost:3000/mock/iqdb to enable a fake iqdb server for
     # development purposes.
-    def iqdbs_server
+    def iqdb_url
+      # "http://localhost:3000/mock/iqdb"
     end
 
     def aws_credentials
@@ -501,9 +502,6 @@ module Danbooru
     def aws_sqs_region
     end
 
-    def aws_sqs_iqdb_url
-    end
-
     def aws_sqs_archives_url
     end
 

config/routes.rb

@@ -329,8 +329,8 @@ Rails.application.routes.draw do
   get "/mock/reportbooru/missed_searches" => "mock_services#reportbooru_missed_searches", as: "mock_reportbooru_missed_searches"
   get "/mock/reportbooru/post_searches/rank" => "mock_services#reportbooru_post_searches", as: "mock_reportbooru_post_searches"
   get "/mock/reportbooru/post_views/rank" => "mock_services#reportbooru_post_views", as: "mock_reportbooru_post_views"
-  get "/mock/iqdbs/similar" => "mock_services#iqdbs_similar", as: "mock_iqdbs_similar"
+  get "/mock/iqdb/query" => "mock_services#iqdb_query", as: "mock_iqdb_query"
-  post "/mock/iqdbs/similar" => "mock_services#iqdbs_similar"
+  post "/mock/iqdb/query" => "mock_services#iqdb_query"
 
   match "*other", to: "static#not_found", via: :all
 end

lib/tasks/iqdb.rake

@@ -2,8 +2,7 @@ namespace :iqdb do
   task reindex_posts: :environment do
     STDIN.each_line do |post_id|
       puts post_id
-      post.remove_iqdb_async
+      post.update_iqdb
-      post.update_iqdb_async
     end
   end
 end

test/functional/iqdb_queries_controller_test.rb

@@ -11,7 +11,7 @@ class IqdbQueriesControllerTest < ActionDispatch::IntegrationTest
       context "with a url parameter" do
         should "render a response" do
           @url = "https://google.com"
-          @matches = [{ "post_id" => @post.id, "width" => 128, "height" => 128, "score" => 95.0 }]
+          @matches = [{ post_id: @post.id, score: 95.0 }]
           mock_iqdb_matches(@matches)
 
           get_auth iqdb_queries_path, @user, as: :javascript, params: { url: @url }
@@ -22,10 +22,13 @@ class IqdbQueriesControllerTest < ActionDispatch::IntegrationTest
       end
 
       context "with a post_id parameter" do
-        should "redirect to iqdbs" do
+        should "render a response" do
-          @matches = [{ "post_id" => @post.id, "width" => 128, "height" => 128, "score" => 95.0 }]
+          @matches = [{ post_id: @post.id, score: 95.0 }]
           mock_iqdb_matches(@matches)
 
+          # Make the call to `@post.file(:preview)` work.
+          Post.any_instance.stubs(:file).returns(File.open("test/files/test.jpg"))
+
           get_auth iqdb_queries_path, @user, params: { post_id: @post.id }
 
           assert_response :success

test/functional/mock_services_controller_test.rb

@@ -15,7 +15,7 @@ class MockServicesControllerTest < ActionDispatch::IntegrationTest
           mock_reportbooru_missed_searches_path,
           mock_reportbooru_post_searches_path,
           mock_reportbooru_post_views_path,
-          mock_iqdbs_similar_path,
+          mock_iqdb_query_path,
         ]
 
         paths.each do |path|

test/functional/post_regenerations_controller_test.rb

@@ -6,6 +6,7 @@ class PostRegenerationsControllerTest < ActionDispatch::IntegrationTest
       @mod = create(:moderator_user, name: "yukari", created_at: 1.month.ago)
       @upload = assert_successful_upload("test/files/test.jpg", user: @mod)
       @post = @upload.post
+      perform_enqueued_jobs # add post to iqdb
     end
 
     context "create action" do
@@ -23,9 +24,6 @@ class PostRegenerationsControllerTest < ActionDispatch::IntegrationTest
 
       context "for an IQDB regeneration" do
         should "regenerate IQDB" do
-          mock_iqdb_service!
-          Post.iqdb_sqs_service.expects(:send_message).with("update\n#{@post.id}\n#{@post.preview_file_url}")
-
           post_auth post_regenerations_path, @mod, params: { post_id: @post.id, category: "iqdb" }
           perform_enqueued_jobs
         end

test/functional/uploads_controller_test.rb

@@ -4,7 +4,6 @@ class UploadsControllerTest < ActionDispatch::IntegrationTest
   context "The uploads controller" do
     setup do
       @user = create(:contributor_user, name: "marisa")
-      mock_iqdb_service!
     end
 
     context "image proxy action" do

test/test_helpers/iqdb_test_helper.rb

@@ -1,28 +1,6 @@
 module IqdbTestHelper
-  def mock_iqdb_service!
-    mock_sqs_service = Class.new do
-      def initialize
-        @commands = []
-      end
-
-      def commands
-        @commands
-      end
-
-      def send_message(msg)
-        @commands << msg.split(/\n/).first
-      end
-    end
-
-    service = mock_sqs_service.new
-    Post.stubs(:iqdb_sqs_service).returns(service)
-    Post.stubs(:iqdb_enabled?).returns(true)
-
-    Danbooru.config.stubs(:iqdbs_server).returns("http://localhost:3004")
-  end
-
   def mock_iqdb_matches(matches)
-    Danbooru.config.stubs(:iqdbs_server).returns("http://localhost:3004")
+    Danbooru.config.stubs(:iqdb_url).returns("http://localhost:5588")
     response = HTTP::Response.new(status: 200, body: matches.to_json, headers: { "Content-Type": "application/json" }, version: "1.1")
     HTTP::Client.any_instance.stubs(:post).returns(response)
   end

test/unit/post_test.rb

@@ -35,6 +35,7 @@ class PostTest < ActiveSupport::TestCase
         @post = @upload.post
         Favorite.add(post: @post, user: @user)
         create(:favorite_group, post_ids: [@post.id])
+        perform_enqueued_jobs # perform IqdbAddPostJob
       end
 
       should "delete the files" do
@@ -85,10 +86,9 @@ class PostTest < ActiveSupport::TestCase
       end
 
       should "remove the post from iqdb" do
-        mock_iqdb_service!
-        Post.iqdb_sqs_service.expects(:send_message).with("remove\n#{@post.id}")
-
         @post.expunge!
+        perform_enqueued_jobs
+        assert_performed_jobs(1, only: IqdbRemovePostJob)
       end
 
       context "that is status locked" do