From 0f36bbf8d31f473f2393d21c3d2058dd45f0ab68 Mon Sep 17 00:00:00 2001 From: evazion Date: Wed, 16 Jun 2021 05:14:15 -0500 Subject: [PATCH] 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 | 2 - Procfile | 5 +- README.md | 2 +- app/controllers/iqdb_queries_controller.rb | 3 +- app/controllers/mock_services_controller.rb | 2 +- app/jobs/iqdb_add_post_job.rb | 5 ++ app/jobs/iqdb_remove_post_job.rb | 5 ++ app/logical/iqdb_client.rb | 86 +++++++++++++++++++ app/logical/iqdb_proxy.rb | 77 ----------------- app/logical/upload_service.rb | 2 + app/logical/upload_service/replacer.rb | 2 +- app/models/post.rb | 42 +++------ app/views/uploads/_related_posts.html.erb | 2 +- config/danbooru_default_config.rb | 12 ++- config/routes.rb | 4 +- lib/tasks/iqdb.rake | 3 +- .../iqdb_queries_controller_test.rb | 9 +- .../mock_services_controller_test.rb | 2 +- .../post_regenerations_controller_test.rb | 4 +- test/functional/uploads_controller_test.rb | 1 - test/test_helpers/iqdb_test_helper.rb | 24 +----- test/unit/post_test.rb | 6 +- 22 files changed, 136 insertions(+), 164 deletions(-) create mode 100644 app/jobs/iqdb_add_post_job.rb create mode 100644 app/jobs/iqdb_remove_post_job.rb create mode 100644 app/logical/iqdb_client.rb delete mode 100644 app/logical/iqdb_proxy.rb diff --git a/.env b/.env index 2c7731b6b..e6c4f8a24 100644 --- a/.env +++ b/.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= diff --git a/Procfile b/Procfile index f9ade14dc..f9cfe0664 100644 --- a/Procfile +++ b/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 +jobs: bin/rake jobs:work \ No newline at end of file diff --git a/README.md b/README.md index ad4ac60da..51771ad05 100644 --- a/README.md +++ b/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 diff --git a/app/controllers/iqdb_queries_controller.rb b/app/controllers/iqdb_queries_controller.rb index 659c22950..408f390f4 100644 --- a/app/controllers/iqdb_queries_controller.rb +++ b/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 diff --git a/app/controllers/mock_services_controller.rb b/app/controllers/mock_services_controller.rb index 67d387fb6..69d6cad61 100644 --- a/app/controllers/mock_services_controller.rb +++ b/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 diff --git a/app/jobs/iqdb_add_post_job.rb b/app/jobs/iqdb_add_post_job.rb new file mode 100644 index 000000000..a99daa6f0 --- /dev/null +++ b/app/jobs/iqdb_add_post_job.rb @@ -0,0 +1,5 @@ +class IqdbAddPostJob < ApplicationJob + def perform(post) + IqdbClient.new.add_post(post) + end +end diff --git a/app/jobs/iqdb_remove_post_job.rb b/app/jobs/iqdb_remove_post_job.rb new file mode 100644 index 000000000..3d4cda697 --- /dev/null +++ b/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 diff --git a/app/logical/iqdb_client.rb b/app/logical/iqdb_client.rb new file mode 100644 index 000000000..e8ad04861 --- /dev/null +++ b/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 diff --git a/app/logical/iqdb_proxy.rb b/app/logical/iqdb_proxy.rb deleted file mode 100644 index 41dc5cd63..000000000 --- a/app/logical/iqdb_proxy.rb +++ /dev/null @@ -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 diff --git a/app/logical/upload_service.rb b/app/logical/upload_service.rb index 5a5d415b9..53f9bf8bd 100644 --- a/app/logical/upload_service.rb +++ b/app/logical/upload_service.rb @@ -85,6 +85,8 @@ class UploadService ) end + @post.update_iqdb + upload.update(status: "completed", post_id: @post.id) @post diff --git a/app/logical/upload_service/replacer.rb b/app/logical/upload_service/replacer.rb index 3d19a86ad..b995bf247 100644 --- a/app/logical/upload_service/replacer.rb +++ b/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) diff --git a/app/models/post.rb b/app/models/post.rb index a77f22ac0..c433bf799 100644 --- a/app/models/post.rb +++ b/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 - extend ActiveSupport::Concern - - module ClassMethods - def iqdb_sqs_service - SqsService.new(Danbooru.config.aws_sqs_iqdb_url) - end - - def iqdb_enabled? - Danbooru.config.aws_sqs_iqdb_url.present? - end - - def remove_iqdb(post_id) - if iqdb_enabled? - iqdb_sqs_service.send_message("remove\n#{post_id}") - end - end + concerning :IqdbMethods do + def update_iqdb + # performs IqdbClient.new.add_post(post) + IqdbAddPostJob.perform_later(self) 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) + def remove_iqdb + # performs IqdbClient.new.remove(id) + IqdbRemovePostJob.perform_later(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"] diff --git a/app/views/uploads/_related_posts.html.erb b/app/views/uploads/_related_posts.html.erb index 9821ba3ae..c283c2b7f 100644 --- a/app/views/uploads/_related_posts.html.erb +++ b/app/views/uploads/_related_posts.html.erb @@ -16,7 +16,7 @@ <% end %> -<% if Danbooru.config.iqdbs_server %> +<% if Danbooru.config.iqdb_url %> <% if params[:url] %>

Loading similar...

diff --git a/config/danbooru_default_config.rb b/config/danbooru_default_config.rb index 5503f1f38..2d6e3fffc 100644 --- a/config/danbooru_default_config.rb +++ b/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). - # Optional. Used for dupe detection and reverse image searches. - # Set to http://localhost/mock/iqdbs to enable a fake iqdb server for + # The URL for the IQDB server (https://github.com/danbooru/iqdb). Optional. + # Used for dupe detection and reverse image searches. Set this to + # 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 diff --git a/config/routes.rb b/config/routes.rb index 753d3d991..bd74dcb88 100644 --- a/config/routes.rb +++ b/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" - post "/mock/iqdbs/similar" => "mock_services#iqdbs_similar" + get "/mock/iqdb/query" => "mock_services#iqdb_query", as: "mock_iqdb_query" + post "/mock/iqdb/query" => "mock_services#iqdb_query" match "*other", to: "static#not_found", via: :all end diff --git a/lib/tasks/iqdb.rake b/lib/tasks/iqdb.rake index 2e156e26d..c4549dcc0 100644 --- a/lib/tasks/iqdb.rake +++ b/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_async + post.update_iqdb end end end diff --git a/test/functional/iqdb_queries_controller_test.rb b/test/functional/iqdb_queries_controller_test.rb index e9af6770d..f10a70c07 100644 --- a/test/functional/iqdb_queries_controller_test.rb +++ b/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 - @matches = [{ "post_id" => @post.id, "width" => 128, "height" => 128, "score" => 95.0 }] + should "render a response" do + @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 diff --git a/test/functional/mock_services_controller_test.rb b/test/functional/mock_services_controller_test.rb index 5f33cc691..02612a42b 100644 --- a/test/functional/mock_services_controller_test.rb +++ b/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| diff --git a/test/functional/post_regenerations_controller_test.rb b/test/functional/post_regenerations_controller_test.rb index 6f54e47af..75a34670f 100644 --- a/test/functional/post_regenerations_controller_test.rb +++ b/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 diff --git a/test/functional/uploads_controller_test.rb b/test/functional/uploads_controller_test.rb index 77bc57388..59e5d1c6f 100644 --- a/test/functional/uploads_controller_test.rb +++ b/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 diff --git a/test/test_helpers/iqdb_test_helper.rb b/test/test_helpers/iqdb_test_helper.rb index 4b52eade7..69179553e 100644 --- a/test/test_helpers/iqdb_test_helper.rb +++ b/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 diff --git a/test/unit/post_test.rb b/test/unit/post_test.rb index 361e9c03f..4ad48f51a 100644 --- a/test/unit/post_test.rb +++ b/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