out_solr.rb in fluent-plugin-output-solr-1.0.0

- old
+ new

@@ -1,87 +1,86 @@
 require 'securerandom'
 require 'rsolr'
 require 'zk'
 require 'rsolr/cloud'
+require 'fluent/plugin/output'
 
-module Fluent
-  class SolrOutput < BufferedOutput
+module Fluent::Plugin
+  class SolrOutput < Output
     Fluent::Plugin.register_output('solr', self)
 
-    DEFAULT_COLLECTION = 'collection1'
+    helpers :inject, :compat_parameters
 
-    DEFAULT_TAG_FIELD = 'tag'
-
-    DEFAULT_TIME_FIELD = 'time'
-    DEFAULT_TIME_FORMAT = '%FT%TZ'
-    DEFAULT_MILLISECOND = false
-
+    DEFAULT_COLLECTION = 'collection1'
     DEFAULT_IGNORE_UNDEFINED_FIELDS = false
-
+    DEFAULT_TAG_FIELD = 'tag'
+    DEFAULT_TIMESTAMP_FIELD = 'event_timestamp'
     DEFAULT_FLUSH_SIZE = 100
+    DEFAULT_BUFFER_TYPE = "memory"
     DEFAULT_COMMIT_WITH_FLUSH = true
 
     MODE_STANDALONE = 'Standalone'
     MODE_SOLRCLOUD = 'SolrCloud'
 
-    include Fluent::SetTagKeyMixin
     config_set_default :include_tag_key, false
-
-    include Fluent::SetTimeKeyMixin
     config_set_default :include_time_key, false
 
-    config_param :base_url, :string, :default => nil,
-                 :desc => 'The Solr base url (for example http://localhost:8983/solr).'
+    config_param :url, :string, :default => nil,
+                 :desc => 'The Solr server url (for example http://localhost:8983/solr/collection1).'
 
     config_param :zk_host, :string, :default => nil,
                  :desc => 'The ZooKeeper connection string that SolrCloud refers to (for example localhost:2181/solr).'
-
     config_param :collection, :string, :default => DEFAULT_COLLECTION,
-                 :desc => 'The Solr collection/core name (default collection1).'
+                 :desc => 'The SolrCloud collection name (default collection1).'
 
+    config_param :defined_fields, :array, :default => nil,
+                 :desc => 'The defined fields in the Solr schema.xml. If omitted, it will get fields via Solr Schema API.'
     config_param :ignore_undefined_fields, :bool, :default => DEFAULT_IGNORE_UNDEFINED_FIELDS,
                  :desc => 'Ignore undefined fields in the Solr schema.xml.'
 
+    config_param :unique_key_field, :string, :default => nil,
+                 :desc => 'A field name of unique key in the Solr schema.xml. If omitted, it will get unique key via Solr Schema API.'
     config_param :tag_field, :string, :default => DEFAULT_TAG_FIELD,
-                 :desc => 'A field name of fluentd tag in the Solr schema.xml (default tag).'
+                 :desc => 'A field name of fluentd tag in the Solr schema.xml (default event_timestamp).'
+    config_param :timestamp_field, :string, :default => DEFAULT_TIMESTAMP_FIELD,
+                 :desc => 'A field name of event timestamp in the Solr schema.xml (default event_timestamp).'
 
-    config_param :time_field, :string, :default => DEFAULT_TIME_FIELD,
-                 :desc => 'A field name of event timestamp in the Solr schema.xml (default time).'
-    config_param :time_format, :string, :default => DEFAULT_TIME_FORMAT,
-                 :desc => 'The format of the time field (default %d/%b/%Y:%H:%M:%S %z).'
-    config_param :millisecond, :bool, :default => DEFAULT_MILLISECOND,
-                 :desc => 'Output millisecond to Solr (default false).'
-
     config_param :flush_size, :integer, :default => DEFAULT_FLUSH_SIZE,
                  :desc => 'A number of events to queue up before writing to Solr (default 100).'
 
     config_param :commit_with_flush, :bool, :default => DEFAULT_COMMIT_WITH_FLUSH,
                  :desc => 'Send commit command to Solr with flush (default true).'
 
+    config_section :buffer do
+      config_set_default :@type, DEFAULT_BUFFER_TYPE
+      config_set_default :chunk_keys, ['tag']
+    end
+
     def initialize
       super
     end
 
     def configure(conf)
+      compat_parameters_convert(conf, :inject)
       super
     end
 
     def start
       super
 
       @mode = nil
-      if ! @base_url.nil? then
+      if ! @url.nil? then
         @mode = MODE_STANDALONE
       elsif ! @zk_host.nil?
         @mode = MODE_SOLRCLOUD
       end
 
       @solr = nil
       @zk = nil
 
       if @mode == MODE_STANDALONE then
-        @solr = RSolr.connect :url => @base_url.end_with?('/') ? @base_url + @collection : @base_url + '/' + @collection
+        @solr = RSolr.connect :url => @url
       elsif @mode == MODE_SOLRCLOUD then
         @zk = ZK.new(@zk_host)
         cloud_connection = RSolr::Cloud::Connection.new(@zk)
         @solr = RSolr::Client.new(cloud_connection, read_timeout: 60, open_timeout: 60)
       end
@@ -94,122 +93,115 @@
         @zk.close
       end
     end
 
     def format(tag, time, record)
-      [tag, time, record].to_msgpack
+      [time, record].to_msgpack
     end
 
+    def formatted_to_msgpack_binary
+      true
+    end
+
+    def multi_workers_ready?
+      true
+    end
+
     def write(chunk)
       documents = []
 
-      # Get fields from Solr
-      fields = get_fields
+      @fields = @defined_fields.nil? ? get_fields : @defined_fields
+      @unique_key = @unique_key_field.nil? ? get_unique_key : @unique_key_field
+      tag = chunk.metadata.tag
+      chunk.msgpack_each do |time, record|
+        record = inject_values_to_record(tag, time, record)
 
-      # Get unique key field from Solr
-      unique_key = get_unique_key
-
-      chunk.msgpack_each do |tag, time, record|
-        # Set unique key and value
-        unless record.has_key?(unique_key) then
-          record.merge!({unique_key => SecureRandom.uuid})
+        unless record.has_key?(@unique_key) then
+          record.merge!({@unique_key => SecureRandom.uuid})
         end
 
-        # Set Fluentd tag to Solr tag field
         unless record.has_key?(@tag_field) then
           record.merge!({@tag_field => tag})
         end
 
-        # Set time
-        tmp_time = Time.at(time).utc
-        if record.has_key?(@time_field) then
-          # Parsing the time field in the record by the specified format.
+        if record.has_key?(@timestamp_field) then
           begin
-            tmp_time = Time.strptime(record[@time_field], @time_format).utc
-          rescue Exception => e
-            log.warn "An error occurred in parsing the time field: #{e.message}"
+            event_timestamp_dt = DateTime.strptime(record[@timestamp_field], "%d/%b/%Y:%H:%M:%S %z").to_s
+            record.merge!({@timestamp_field => Time.parse(event_timestamp_dt.to_s).utc.strftime('%FT%TZ')})
+          rescue
+            record.merge!({@timestamp_field => Time.at(time).utc.strftime('%FT%TZ')})
           end
-        end
-        if @millisecond then
-          record.merge!({@time_field => '%s.%03dZ' % [tmp_time.strftime('%FT%T'), tmp_time.usec / 1000.0]})
         else
-          record.merge!({@time_field => tmp_time.strftime('%FT%TZ')})
+          record.merge!({@timestamp_field => Time.at(time).utc.strftime('%FT%TZ')})
         end
 
-        # Ignore undefined fields
         if @ignore_undefined_fields then
           record.each_key do |key|
-            unless fields.include?(key) then
+            unless @fields.include?(key) then
               record.delete(key)
             end
           end
         end
 
-        # Add record to documents
         documents << record
 
-        # Update when flash size is reached
         if documents.count >= @flush_size
           update documents
           documents.clear
         end
       end
 
-      # Update remaining documents
       update documents unless documents.empty?
     end
 
     def update(documents)
-      begin
-        if @mode == MODE_STANDALONE then
-          @solr.add documents, :params => {:commit => @commit_with_flush}
-        elsif @mode == MODE_SOLRCLOUD then
-          @solr.add documents, collection: @collection, :params => {:commit => @commit_with_flush}
-        end
-        log.debug "Sent #{documents.count} document(s) to Solr"
-      rescue Exception
-        log.warn "An error occurred while sending #{documents.count} document(s) to Solr"
+      if @mode == MODE_STANDALONE then
+        @solr.add documents, :params => {:commit => @commit_with_flush}
+        log.debug "Added %d document(s) to Solr" % documents.count
+      elsif @mode == MODE_SOLRCLOUD then
+        @solr.add documents, collection: @collection, :params => {:commit => @commit_with_flush}
+        log.debug "Update: Added %d document(s) to Solr" % documents.count
       end
+      rescue Exception => e
+        log.warn "Update: An error occurred while indexing: #{e.message}"
     end
 
     def get_unique_key
-      unique_key = 'id'
+      response = nil
 
-      begin
-        response = nil
-        if @mode == MODE_STANDALONE then
-          response = @solr.get 'schema/uniquekey'
-        elsif @mode == MODE_SOLRCLOUD then
-          response = @solr.get 'schema/uniquekey', collection: @collection
-        end
-        unique_key = response['uniqueKey']
-        log.debug "Unique key: #{unique_key}"
-      rescue Exception
-        log.warn 'An error occurred while getting unique key'
+      if @mode == MODE_STANDALONE then
+        response = @solr.get 'schema/uniquekey'
+      elsif @mode == MODE_SOLRCLOUD then
+        response = @solr.get 'schema/uniquekey', collection: @collection
       end
 
+      unique_key = response['uniqueKey']
+      log.debug "Unique key: #{unique_key}"
+
       return unique_key
+
+      rescue Exception => e
+        log.warn "Unique key: #{e.message}"
     end
 
     def get_fields
-      fields = []
+      response = nil
 
-      begin
-        response = nil
+      if @mode == MODE_STANDALONE then
+        response = @solr.get 'schema/fields'
+      elsif @mode == MODE_SOLRCLOUD then
+        response = @solr.get 'schema/fields', collection: @collection
+      end
 
-        if @mode == MODE_STANDALONE then
-          response = @solr.get 'schema/fields'
-        elsif @mode == MODE_SOLRCLOUD then
-          response = @solr.get 'schema/fields', collection: @collection
-        end
-        response['fields'].each do |field|
-          fields.push(field['name'])
-        end
-        log.debug "Fields: #{fields}"
-      rescue Exception
-        log.warn 'An error occurred while getting fields'
+      fields = []
+      response['fields'].each do |field|
+        fields.push(field['name'])
       end
+      log.debug "Fields: #{fields}"
 
       return fields
+
+      rescue Exception => e
+        log.warn "Fields: #{e.message}"
     end
   end
 end