iqdb changes

Author: r888888888

README

@@ -36,3 +36,11 @@ correct.
 debug your Nginx configuration file.
 
 4) Check all log files.
+
+=== IQDB Integration
+
+In order to enable IQDB integration, you must compile and run the IQDB server somewhere (preferably on the local server). There are two Danbooru configuration settings that you must then set: iqdb_hostname_and_port and iqdb_file.
+
+You must then populate the initial database. There is a fix script called 028_iqdb_import.rb to do this for you.
+
+From then on, all new uploads will asynchronously create two tasks: one to update the IQDB database through the server interface (which updates the in-memory representation), and another to the command interface (which updates the disk representation). Expunging posts will also update the database.

app/logical/iqdb/command.rb

@@ -0,0 +1,29 @@
+module Iqdb
+  class Command
+    attr_reader :database
+
+    def initialize(database)
+      @database = database
+    end
+
+    def process(&block)
+      IO.popen("iqdb #{database}", "w", &block)
+    end
+
+    def add(post)
+      hex = post.id.to_s(16)
+      process do |io|
+        io.puts "add 0 #{hex} :#{post.preview_file_path}"
+        io.puts "quit"
+      end
+    end
+
+    def remove(post_id)
+      hex = post_id.to_s(16)
+      process do |io|
+        io.puts "remove 0 #{hex}"
+        io.puts "quit"
+      end
+    end
+  end
+end

app/logical/iqdb/server.rb

@@ -5,26 +5,7 @@ class Server
     FLAG_WIDTH_AS_SET = 0x08
     FLAG_DISCARD_COMMON_COEFFS = 0x16
 
-    attr_reader :hostname, :port
-
-    def self.import(database)
-      IO.popen("iqdb #{database}", "w") do |io|
-        Post.find_each do |post|
-          puts "Adding #{post.id}"
-          io.puts "#{post.id.to_s(16)} :#{post.preview_file_path}"
-        end
-      end
-    end
-
-    def self.add(database, image_id, filename)
-      image_id_hex = image_id.to_s(16)
-      `iqdb add #{database} #{image_id_hex} :#{filename}`
-    end
-
-    def self.remove(database, image_id)
-      image_id_hex = image_id.to_s(16)
-      `iqdb remove 0 #{image_id_hex} #{database}`
-    end
+    attr_reader :hostname, :port, :socket
 
     def initialize(hostname, port)
       @hostname = hostname
@@ -36,7 +17,7 @@ def open
     end
 
     def close
-      @socket.close
+      socket.close
     end
 
     def request
@@ -46,9 +27,23 @@ def request
       close
     end
 
-    def query(dbid, results, filename, flags = FLAG_DISCARD_COMMON_COEFFS)
+    def add(post)
+      request do
+        hex = post.id.to_s(16)
+        socket.puts "add 0 #{hex} :#{post.preview_file_path}"
+      end
+    end
+
+    def remove(post_id)
+      request do
+        hext = post_id.to_s(16)
+        socket.puts "remove 0 #{hex}"
+      end
+    end
+
+    def query(results, filename, flags = FLAG_DISCARD_COMMON_COEFFS)
       request do
-        @socket.puts "query #{dbid} #{flags} #{results} #{filename}"
+        socket.puts "query 0 #{flags} #{results} #{filename}"
         responses = Responses::Collection.new(@socket.read)
       end
     end

app/models/post.rb

@@ -9,8 +9,8 @@ class SearchError < Exception ; end
   after_save :create_version
   after_save :update_parent_on_save
   after_save :apply_post_metatags, :on => :create
-  # after_save :update_iqdb, :on => :create
-  # after_destroy :remove_iqdb
+  # after_save :update_iqdb_async, :on => :create
+  # after_destroy :remove_iqdb_async
   before_save :merge_old_changes
   before_save :normalize_tags
   before_save :update_tag_post_counts
@@ -1272,17 +1272,31 @@ def parse_pixiv_id
   end
 
   module IqdbMethods
-    def update_iqdb
+    extend ActiveSupport::Concern
+
+    module ClassMethods
+      def remove_iqdb(post_id)
+        Iqdb::Server.new(*Danbooru.config.iqdb_hostname_and_port).remove(post_id)
+        Iqdb::Command.new(Danbooru.config.iqdb_file).remove(post_id)
+      end
+    end
+
+    def update_iqdb_async
       Danbooru.config.all_server_hosts.each do |host|
-        Iqdb::Server.delay(:queue => host).add(Danbooru.config.iqdb_file, id, preview_file_path)
+        delay(:queue => host).update_iqdb
       end
     end
 
-    def remove_iqdb
+    def remove_iqdb_async
       Danbooru.config.all_server_hosts.each do |host|
-        Iqdb::Server.delay(:queue => host).remove(Danbooru.config.iqdb_file, id)
+        Post.delay(:queue => host).remove_iqdb(id)
       end
     end
+
+    def update_iqdb
+      Iqdb::Server.new(*Danbooru.config.iqdb_hostname_and_port).add(self)
+      Iqdb::Command.new(Danbooru.config.iqdb_file).add(self)
+    end
   end
 
   include FileMethods
@@ -1303,6 +1317,7 @@ def remove_iqdb
   include ApiMethods
   extend SearchMethods
   include PixivMethods
+  include IqdbMethods
 
   def visible?
     return false if !Danbooru.config.can_user_see_post?(CurrentUser.user, self)

script/fixes/029_iqdb_import.rb

@@ -2,4 +2,4 @@
 
 require File.expand_path(File.join(File.dirname(__FILE__), '..', '..', 'config', 'environment'))
 
-Iqdb::Server.import("/var/www/danbooru2/shared/iqdb.db")
+Iqdb::Server.import(Danbooru.config.iqdb_file)