traject 1.1.0 → 2.0.0.rc.1

data/lib/traject/command_line.rb

@@ -12,7 +12,7 @@ module Traject
   #
   # A CommandLine object has a single persistent Indexer object it uses
   class CommandLine
-    # orig_argv is origina one passed in, remaining_argv is after destructive
+    # orig_argv is original one passed in, remaining_argv is after destructive
     # processing by slop, still has file args in it etc.
     attr_accessor :orig_argv, :remaining_argv
     attr_accessor :slop, :options

data/lib/traject/csv_writer.rb

@@ -0,0 +1,34 @@
+require 'traject/delimited_writer'
+require 'csv'
+
+# A CSV-writer, for folks who like that sort of thing.
+# Use DelimitedWriter for non-CSV lines (e.g., tab-delimited)
+#
+#
+
+class Traject::CSVWriter < Traject::DelimitedWriter
+
+  def initialize(*args)
+    super
+    self.delimiter = nil # Let CSV take care of it
+  end
+
+  def _write(data)
+    @output_file << data
+  end
+
+  # Turn the output file into a CSV writer
+  def open_output_file
+    of = super
+    CSV.new(of)
+  end
+
+  # Let CSV take care of the comma escaping
+  def escape(x)
+    x = x.to_s
+    x.gsub! internal_delimiter, @eidelim
+    x
+  end
+
+
+end

data/lib/traject/delimited_writer.rb

@@ -0,0 +1,110 @@
+require 'traject/line_writer'
+
+# A simple line writer that uses configuration to determine
+# how to produce a tab-delimited file
+#
+# Appropos settings:
+#
+# * output_file  -- the file to write to
+# * output_stream -- the stream to write to, if defined and output_file is not
+# * delimited_writer.delimiter  -- What to separate fields with; default is tab
+# * delimited_writer.internal_delimiter -- Delimiter _within_ a field, for multiple
+#   values. Default is pipe ( | )
+# * delimited_writer.fields -- comma-separated list of the fields to output
+# * delimited_writer.header (true/false) -- boolean that determines if we should output a header row. Default is true
+# * delimited_writer.escape -- If a value actually contains the delimited or internal_delimiter, what to do?
+#   If unset, will follow the procedure below. If set, will turn it into the character(s) given
+#
+#
+# If `delimited_writer.escape` is not set, the writer will automatically
+# escape delimiters/internal_delimiters in the following way:
+#  * If the delimiter is a tab, replace tabs in values with a single space
+#  * If the delimiter is anything else, prefix it with a backslash
+
+class Traject::DelimitedWriter < Traject::LineWriter
+
+  attr_reader :delimiter,  :internal_delimiter, :edelim, :eidelim
+  attr_accessor :header
+
+  def initialize(settings)
+    super
+
+    # fields to output
+
+    begin
+      @fields = settings['delimited_writer.fields'].split(",")
+    rescue NoMethodError => e
+    end
+
+    if e or @fields.empty?
+      raise ArgumentError.new("#{self.class.name} must have a comma-delimited list of field names to output set in setting 'delimited_writer.fields'")
+    end
+
+    self.delimiter = settings['delimited_writer.delimiter'] || "\t"
+    self.internal_delimiter = settings['delimited_writer.internal_delimiter'] || '|'
+    self.header = settings['delimited_writer.header'].to_s != 'false'
+
+    # Output the header if need be
+    write_header if @header
+  end
+
+  def escaped_delimiter(d)
+    return nil if d.nil?
+    d == "\t" ? ' ' : '\\' + d
+  end
+
+  def delimiter=(d)
+    @delimiter = d
+    @edelim = escaped_delimiter(d)
+    self
+  end
+
+  def internal_delimiter=(d)
+    @internal_delimiter = d
+    @eidelim =  escaped_delimiter(d)
+  end
+
+
+
+
+  def write_header
+    _write(@fields)
+  end
+
+  def _write(data)
+    output_file.puts(data.join(delimiter))
+  end
+
+  # Get the output values out of the context
+  def raw_output_values(context)
+    context.output_hash.values_at(*@fields)
+  end
+
+  # Escape the delimiters in whatever way has been defined
+  def escape(x)
+    x = x.to_s
+    x.gsub! @delimiter, @edelim if @delimiter
+    x.gsub! @internal_delimiter, @eidelim
+    x
+  end
+
+
+  # Derive actual output field values from the raw values
+  def output_values(raw)
+    raw.map do |x|
+      if x.is_a? Array
+        x.map!{|s| escape(s)}
+        x.join(@internal_delimiter)
+      else
+        escape(x)
+      end
+    end
+  end
+
+  # Spit out the escaped values joined by the delimiter
+  def serialize(context)
+    output_values(raw_output_values(context))
+  end
+
+
+end

data/lib/traject/indexer.rb

@@ -6,13 +6,16 @@ require 'traject/thread_pool'
 
 require 'traject/indexer/settings'
 require 'traject/marc_reader'
-require 'traject/marc4j_reader'
 require 'traject/json_writer'
-require 'traject/solrj_writer'
+require 'traject/solr_json_writer'
 
 require 'traject/macros/marc21'
 require 'traject/macros/basic'
 
+if defined? JRUBY_VERSION
+  require 'traject/marc4j_reader'
+end
+
 # This class does indexing for traject: Getting input records from a Reader
 # class, mapping the input records to an output hash, and then sending the output
 # hash off somewhere (usually Solr) with a Writer class.
@@ -53,8 +56,9 @@ require 'traject/macros/basic'
 #   2) Responds to the usual ruby #each, returning a source record from each #each.
 #      (Including Enumerable is prob a good idea too)
 #
-#  The default reader is the Traject::Marc4JReader, who's behavior is
-#  further customized by several settings in the Settings hash.
+#  The default reader is the Traject::MarcReader, who's behavior is
+#  further customized by several settings in the Settings hash. Jruby users
+#  with specialized needs may want to look at the gem traject-marc4j_reader.
 #
 #  Alternate readers can be set directly with the #reader_class= method, or
 #  with the "reader_class_name" Setting, a String name of a class
@@ -72,14 +76,22 @@ require 'traject/macros/basic'
 #  4) Optionally implements a #skipped_record_count method, returning int count of records
 #     that were skipped due to errors (and presumably logged)
 #
-#  The default writer is the SolrJWriter, using Java SolrJ to
-#  write to a Solr.  A few other built-in writers are available,
-#  but it's anticipated more will be created as plugins or local
-#  code for special purposes.
+#  Traject packages one solr writer: traject/solr_json_writer, which sends
+#  in json format and works under both ruby and  jruby, but only with solr version
+#  >= 3.2. To index to an older solr installation, you'll need to use jruby and
+#  install the gem traject-solrj_writer, which uses the solrj .jar underneath.
 #
 #  You can set alternate writers by setting a Class object directly
 #  with the #writer_class method, or by the 'writer_class_name' Setting,
-#  with a String name of class meeting the Writer contract.
+#  with a String name of class meeting the Writer contract. There are several
+#  that ship with traject itself:
+#
+#  * traject/json_writer (Traject::JsonWriter) -- write newline-delimied json files.
+#  * traject/yaml_writer (Traject::YamlWriter) -- write pretty yaml file; very human-readable
+#  * traject/debug_writer (Traject::DebugWriter) -- write a tab-delimited file where
+#    each line consists of the id, field, and value(s).
+#  * traject/delimited_writer and traject/csv_writer -- write character-delimited files
+#    (default is tab-delimited) or comma-separated-value files.
 #
 class Traject::Indexer
 
@@ -310,7 +322,13 @@ class Traject::Indexer
     reader = self.reader!(io_stream)
     writer = self.writer!
 
-    thread_pool = Traject::ThreadPool.new(settings["processing_thread_pool"].to_i)
+
+    processing_threads = settings["processing_thread_pool"].to_i
+    if processing_threads > 0 and !(defined? JRuby)
+      processing_threads = 0
+      logger.warn "Processing threads set to 0 because we're not running under JRuby"
+    end
+    thread_pool = Traject::ThreadPool.new(processing_threads)
 
     logger.info "   Indexer with reader: #{reader.class.name} and writer: #{writer.class.name}"
 
@@ -326,7 +344,7 @@ class Traject::Indexer
       thread_pool.raise_collected_exception!
 
       if settings["debug_ascii_progress"].to_s == "true"
-        $stderr.write "." if count % settings["solrj_writer.batch_size"] == 0
+        $stderr.write "." if count % settings["solr_writer.batch_size"].to_i == 0
       end
 
       if log_batch_size && (count % log_batch_size == 0)

data/lib/traject/indexer/settings.rb

@@ -1,4 +1,5 @@
 require 'hashie'
+require 'concurrent'
 
 class Traject::Indexer
 
@@ -22,9 +23,6 @@ class Traject::Indexer
     include Hashie::Extensions::MergeInitializer # can init with hash
     include Hashie::Extensions::IndifferentAccess
 
-    # Hashie bug Issue #100 https://github.com/intridea/hashie/pull/100
-    alias_method :store, :indifferent_writer
-
     def initialize(*args)
       super
       self.default_proc = lambda do |hash, key|
@@ -59,19 +57,37 @@ class Traject::Indexer
     def fill_in_defaults!
       self.reverse_merge!(self.class.defaults)
     end
+    
+    
+    def self.mri_defaults
+      {
+        "reader_class_name"         => "Traject::MarcReader",
+        "writer_class_name"         => "Traject::SolrJsonWriter",
+        "marc_source.type"          => "binary",            
+        "solrj_writer.batch_size"   => 200,
+        "solrj_writer.thread_pool"  => 1,
+        "processing_thread_pool"    => self.default_processing_thread_pool,
+        "log.batch_size.severity"   => "info"
+      }
+    end
 
-    def self.defaults
-      @@defaults ||= {
-      "reader_class_name"         => "Traject::MarcReader",
-      "writer_class_name"         => "Traject::SolrJWriter",
-      "marc_source.type"          => "binary",            
-      "marc4j_reader.permissive"  => true,
-      "solrj_writer.batch_size"   => 200,
-      "solrj_writer.thread_pool"  => 1,
-      "processing_thread_pool"    => 3,
-      "log.batch_size.severity"   => "info"
+    def self.jruby_defaults
+      {
+        'reader_class_name' => "Traject::Marc4JReader",
+        'marc4j_reader.permissive' => true
       }
     end
+    
+        
+    def self.defaults
+      return @@defaults if defined? @@defaults
+      default_settings = self.mri_defaults
+      if defined? JRUBY_VERSION
+        default_settings.merge! self.jruby_defaults
+      end
+      
+      @@defaults = default_settings
+    end
 
     def inspect
       # Keep any key ending in password out of the inspect
@@ -80,5 +96,15 @@ class Traject::Indexer
         hash
       end.inspect
     end
+
+    protected
+    def self.default_processing_thread_pool
+      if ["jruby", "rbx"].include? ENV["RUBY_ENGINE"]
+        [1, Concurrent.processor_count - 1].max
+      else
+        1
+      end
+    end
+
   end
 end

data/lib/traject/line_writer.rb

@@ -16,14 +16,18 @@ require 'thread'
 # method. For instance, see JsonWriter.
 class Traject::LineWriter
   attr_reader :settings
-  attr_reader :write_mutex
+  attr_reader :write_mutex, :output_file
 
   def initialize(argSettings)
     @settings     = argSettings
     @write_mutex  = Mutex.new
 
     # trigger lazy loading now for thread-safety
-    output_file
+    @output_file = open_output_file
+  end
+
+  def _write(data)
+    output_file.puts(data)
   end
 
 
@@ -34,13 +38,13 @@ class Traject::LineWriter
   def put(context)
     serialized = serialize(context)
     write_mutex.synchronize do
-      output_file.puts(serialized)
+      _write(serialized)
     end
   end
 
-  def output_file
+  def open_output_file
     unless defined? @output_file
-      @output_file =
+      of =
         if settings["output_file"]
           File.open(settings["output_file"], 'w:UTF-8')
         elsif settings["output_stream"]
@@ -49,7 +53,7 @@ class Traject::LineWriter
           $stdout
         end
     end
-    return @output_file
+    return of
   end
 
   def close

data/lib/traject/marc_reader.rb

@@ -5,7 +5,8 @@ require 'traject/ndj_reader'
 # can read MARC ISO 2709 ('binary'), MARC-XML, and Marc-in-json (newline-delimited-json).
 #
 # Marc4JReader is an alternative to this class, powered by Marc4J. You may be interested
-# in comparing for performance, under your particular use case. 
+# in comparing for performance, under your particular use case. To use it, you'll need
+# the gem traject-marc4j_reader.
 #
 # By default assumes binary MARC encoding, please set marc_source.type setting
 # for XML or json. If binary, please set marc_source.encoding with char encoding. 

data/lib/traject/solr_json_writer.rb

@@ -0,0 +1,277 @@
+require 'yell'
+
+require 'traject'
+require 'traject/util'
+require 'traject/qualified_const_get'
+require 'traject/thread_pool'
+
+require 'json'
+require 'httpclient'
+
+require 'uri'
+require 'thread'     # for Mutex/Queue
+require 'concurrent' # for atomic_fixnum
+
+# Write to Solr using the JSON interface; only works for Solr >= 3.2
+#
+# This should work under both MRI and JRuby, with JRuby getting much
+# better performance due to the threading model.
+#
+# Relevant settings
+#
+# * solr.url (optional if solr.update_url is set) The URL to the solr core to index into
+#
+# * solr.update_url: The actual update url. If unset, we'll first see if
+#   "#{solr.url}/update/json" exists, and if not use "#{solr.url}/update"
+#
+# * solr_writer.batch_size: How big a batch to send to solr. Default is 100.
+#   My tests indicate that this setting doesn't change overall index speed by a ton.
+#
+# * solr_writer.thread_pool: How many threads to use for the writer. Default is 1.
+#   Likely useful even under MRI since thread will be waiting on Solr for some time. 
+#
+# * solr_writer.max_skipped: How many records skipped due to errors before we 
+#   bail out with a fatal error? Set to -1 for unlimited skips. Default 0, 
+#   raise and abort on a single record that could not be added to Solr. 
+#
+# * solr_writer.commit_on_close: Set to true (or "true") if you want to commit at the
+#   end of the indexing run. (Old "solrj_writer.commit_on_close" supported for backwards
+#   compat only.)
+#
+# * solr_writer.commit_timeout: If commit_on_close, how long to wait for Solr before
+#   giving up as a timeout. Default 10 minutes. Solr can be slow. 
+#
+# * solr_json_writer.http_client Mainly intended for testing, set your own HTTPClient
+#   or mock object to be used for HTTP. 
+
+
+class Traject::SolrJsonWriter
+  include Traject::QualifiedConstGet
+
+  DEFAULT_MAX_SKIPPED = 0
+  DEFAULT_BATCH_SIZE  = 100
+
+  # The passed-in settings
+  attr_reader :settings, :thread_pool_size
+
+  # A queue to hold documents before sending to solr
+  attr_reader :batched_queue
+
+  def initialize(argSettings)
+    @settings = Traject::Indexer::Settings.new(argSettings)
+
+    # Set max errors
+    @max_skipped = (@settings['solr_writer.max_skipped'] || DEFAULT_MAX_SKIPPED).to_i
+    if @max_skipped < 0
+      @max_skipped = nil
+    end
+
+    @http_client = @settings["solr_json_writer.http_client"] || HTTPClient.new
+
+    @batch_size = (settings["solr_writer.batch_size"] || DEFAULT_BATCH_SIZE).to_i
+    @batch_size = 1 if @batch_size < 1
+
+    # Store error count in an AtomicInteger, so multi threads can increment
+    # it safely, if we're threaded.
+    @skipped_record_incrementer = Concurrent::AtomicFixnum.new(0)
+
+
+    # How many threads to use for the writer?
+    # if our thread pool settings are 0, it'll just create a null threadpool that
+    # executes in calling context.
+    @thread_pool_size = (@settings["solr_writer.thread_pool"] || 1).to_i
+
+    @batched_queue         = Queue.new
+    @thread_pool = Traject::ThreadPool.new(@thread_pool_size)
+
+    # old setting solrj_writer supported for backwards compat, as we make
+    # this the new default writer. 
+    @commit_on_close = (settings["solr_writer.commit_on_close"] || settings["solrj_writer.commit_on_close"]).to_s == "true"
+
+    # Figure out where to send updates
+    @solr_update_url = self.determine_solr_update_url
+
+    logger.info("   #{self.class.name} writing to '#{@solr_update_url}' in batches of #{@batch_size} with #{@thread_pool_size} bg threads")
+  end
+
+
+  # Add a single context to the queue, ready to be sent to solr
+  def put(context)
+    @thread_pool.raise_collected_exception!
+
+    @batched_queue << context
+    if @batched_queue.size >= @batch_size
+      batch = Traject::Util.drain_queue(@batched_queue)
+      @thread_pool.maybe_in_thread_pool(batch) {|batch_arg| send_batch(batch_arg) }
+    end
+  end
+
+  # Send the given batch of contexts. If something goes wrong, send
+  # them one at a time.
+  # @param [Array<Traject::Indexer::Context>] an array of contexts
+  def send_batch(batch)
+    return if batch.empty?
+    json_package = JSON.generate(batch.map { |c| c.output_hash })
+    begin
+      resp = @http_client.post @solr_update_url, json_package, "Content-type" => "application/json"
+    rescue StandardError => exception
+    end
+
+    if exception || resp.status != 200
+      error_message = exception ? 
+        Traject::Util.exception_to_log_message(exception) : 
+        "Solr response: #{resp.status}: #{resp.body}"
+
+      logger.error "Error in Solr batch add. Will retry documents individually at performance penalty: #{error_message}"
+      
+      batch.each do |c|
+        send_single(c)
+      end
+    end
+  end
+
+
+  # Send a single context to Solr, logging an error if need be
+  # @param [Traject::Indexer::Context] c The context whose document you want to send
+  def send_single(c)
+    json_package = JSON.generate([c.output_hash])
+    begin
+      resp = @http_client.post @solr_update_url, json_package, "Content-type" => "application/json"
+      # Catch Timeouts and network errors as skipped records, but otherwise
+      # allow unexpected errors to propagate up. 
+    rescue HTTPClient::TimeoutError, SocketError, Errno::ECONNREFUSED => exception
+    end
+
+    if exception || resp.status != 200
+      if exception
+        msg = Traject::Util.exception_to_log_message(e)
+      else
+        msg = "Solr error response: #{resp.status}: #{resp.body}"
+      end
+      logger.error "Could not add record #{record_id_from_context c} at source file position #{c.position}: #{msg}"
+      logger.debug(c.source_record.to_s)
+
+      @skipped_record_incrementer.increment
+      if @max_skipped and skipped_record_count > @max_skipped
+        raise RuntimeError.new("#{self.class.name}: Exceeded maximum number of skipped records (#{@max_skipped}): aborting")
+      end
+
+    end
+
+  end
+
+
+  # Get the logger from the settings, or default to an effectively null logger
+  def logger
+    settings["logger"] ||= Yell.new(STDERR, :level => "gt.fatal") # null logger
+  end
+
+  # Returns MARC 001, then a slash, then output_hash["id"] -- if both
+  # are present. Otherwise may return just one, or even an empty string. 
+  def record_id_from_context(context)
+    marc_id = context.source_record && context.source_record['001'] && context.source_record['001'].value
+    output_id = context.output_hash["id"]
+
+    return [marc_id, output_id].compact.join("/")
+  end
+
+
+  # On close, we need to (a) raise any exceptions we might have, (b) send off
+  # the last (possibly empty) batch, and (c) commit if instructed to do so
+  # via the solr_writer.commit_on_close setting.
+  def close
+    @thread_pool.raise_collected_exception!
+
+    # Finish off whatever's left. Do it in the thread pool for
+    # consistency, and to ensure expected order of operations, so
+    # it goes to the end of the queue behind any other work.
+    batch = Traject::Util.drain_queue(@batched_queue)
+    if batch.length > 0
+      @thread_pool.maybe_in_thread_pool { send_batch(batch) }
+    end
+
+    # Wait for shutdown, and time it.
+    logger.debug "#{self.class.name}: Shutting down thread pool, waiting if needed..."
+    elapsed = @thread_pool.shutdown_and_wait
+    if elapsed > 60
+      logger.warn "Waited #{elapsed} seconds for all threads, you may want to increase solr_writer.thread_pool (currently #{@settings["solr_writer.thread_pool"]})"
+    end
+    logger.debug "#{self.class.name}: Thread pool shutdown complete"
+    logger.warn "#{self.class.name}: #{skipped_record_count} skipped records" if skipped_record_count > 0
+
+    # check again now that we've waited, there could still be some
+    # that didn't show up before.
+    @thread_pool.raise_collected_exception!
+
+    # Commit if we're supposed to
+    if @commit_on_close
+      commit
+    end
+  end
+
+
+  # Send a commit
+  def commit
+    logger.info "#{self.class.name} sending commit to solr at url #{@solr_update_url}..."
+
+    original_timeout = @http_client.receive_timeout
+
+    @http_client.receive_timeout = (settings["commit_timeout"] || (10 * 60)).to_i
+
+    resp = @http_client.get(@solr_update_url, {"commit" => 'true'})
+    unless resp.status == 200
+      raise RuntimeError.new("Could not commit to Solr: #{resp.status} #{resp.body}")
+    end
+
+    @http_client.receive_timeout = original_timeout
+  end
+
+
+  # Return count of encountered skipped records. Most accurate to call
+  # it after #close, in which case it should include full count, even
+  # under async thread_pool.
+  def skipped_record_count
+    @skipped_record_incrementer.value
+  end
+
+
+  # Relatively complex logic to determine if we have a valid URL and what it is
+  def determine_solr_update_url
+    if settings['solr.update_url']
+      check_solr_update_url(settings['solr.update_url'])
+    else
+      derive_solr_update_url_from_solr_url(settings['solr.url'])
+    end
+  end
+
+
+  # If we've got a solr.update_url, make sure it's ok
+  def check_solr_update_url(url)
+    unless url =~ /^#{URI::regexp}$/
+      raise ArgumentError.new("#{self.class.name} setting `solr.update_url` doesn't look like a URL: `#{url}`")
+    end
+    url
+  end
+
+  def derive_solr_update_url_from_solr_url(url)
+    # Nil? Then we bail
+    if url.nil?
+      raise ArgumentError.new("#{self.class.name}: Neither solr.update_url nor solr.url set; need at least one")
+    end
+
+    # Not a URL? Bail
+    unless url =~ /^#{URI::regexp}$/
+      raise ArgumentError.new("#{self.class.name} setting `solr.url` doesn't look like a URL: `#{url}`")
+    end
+
+    # First, try the /update/json handler
+    candidate = [url.chomp('/'), 'update', 'json'].join('/')
+    resp      = @http_client.get(candidate)
+    if resp.status == 404
+      candidate = [url.chomp('/'), 'update'].join('/')
+    end
+    candidate
+  end
+
+
+end