diff --git a/app/logical/iqdb/download.rb b/app/logical/iqdb/download.rb deleted file mode 100644 index 200f17572..000000000 --- a/app/logical/iqdb/download.rb +++ /dev/null @@ -1,47 +0,0 @@ -module Iqdb - class Download - class Error < StandardError; end - - def self.enabled? - Danbooru.config.iqdbs_server.present? && Danbooru.config.iqdbs_auth_key.present? - end - - def self.get_referer(url) - headers = {} - datums = {} - - strategy = Sources::Strategies.find(url) - - [strategy.image_url, strategy.headers["Referer"]] - end - - def self.find_similar(source) - raise NotImplementedError, "the IQDBs service isn't configured. Similarity searches are not available." unless enabled? - - url, ref = get_referer(source) - params = { - "key" => Danbooru.config.iqdbs_auth_key, - "url" => url, - "ref" => ref - } - uri = URI.parse("#{Danbooru.config.iqdbs_server}/similar") - uri.query = URI.encode_www_form(params) - - resp = HTTParty.get(uri, Danbooru.config.httparty_options) - raise "HTTP error code: #{resp.code} #{resp.message}" unless resp.success? - - json = JSON.parse(resp.body) - raise "IQDB error: #{json["error"]}" unless json.is_a?(Array) - - post_ids = json.map { |match| match["post_id"] } - posts = Post.find(post_ids) - - json.map do |match| - post = posts.find { |post| post.id == match["post_id"] } - match.with_indifferent_access.merge({ post: post }) - end - rescue => e - raise Error, { message: e.message, iqdb_response: json } - end - end -end diff --git a/app/logical/iqdb_proxy.rb b/app/logical/iqdb_proxy.rb index 524bc7242..7f487ab28 100644 --- a/app/logical/iqdb_proxy.rb +++ b/app/logical/iqdb_proxy.rb @@ -1,19 +1,24 @@ class IqdbProxy - def self.search(params) - raise NotImplementedError unless Danbooru.config.iqdbs_server.present? + class Error < StandardError; end - limit = params[:limit].presence || 20 - limit = limit.to_i.clamp(1, 100) + def self.enabled? + Danbooru.config.iqdbs_server.present? + end + + def self.search(params) + raise NotImplementedError, "the IQDBs service isn't configured" unless enabled? + + limit = params[:limit]&.to_i&.clamp(1, 1000) || 20 similarity = params[:similarity].to_f.clamp(0.0, 100.0) if params[:file].present? - results = query_file(params[:file], limit) + results = query(file: params[:file], limit: limit) elsif params[:url].present? url = Sources::Strategies.find(params[:url]).image_url - results = query_url(url, limit) + results = query(url: url, limit: limit) elsif params[:post_id].present? url = Post.find(params[:post_id]).preview_file_url - results = query_url(url, limit) + results = query(url: url, limit: limit) else results = [] end @@ -22,26 +27,19 @@ class IqdbProxy decorate_posts(results) end - def self.query_url(url, limit) - query = { url: url, limit: limit } - response = HTTParty.get("#{Danbooru.config.iqdbs_server}/similar", query: query, **Danbooru.config.httparty_options) - response.parsed_response - end - - def self.query_file(file, limit) - body = { file: file, limit: limit } - response = HTTParty.post("#{Danbooru.config.iqdbs_server}/similar", body: body, **Danbooru.config.httparty_options) + def self.query(params) + response = HTTParty.post("#{Danbooru.config.iqdbs_server}/similar", body: params, **Danbooru.config.httparty_options) + raise Error, "HTTP error: #{response.code} #{response.message}" unless response.success? response.parsed_response end def self.decorate_posts(json) - json.map do |x| - begin - x["post"] = Post.find(x["post_id"]) - x - rescue ActiveRecord::RecordNotFound - nil - end + 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