handle no @timestamp and fix generateId

review changes

Author: colinsurprenant

lib/logstash/outputs/mongodb.rb

@@ -46,19 +46,22 @@ class LogStash::Outputs::Mongodb < LogStash::Outputs::Base
   # Mutex used to synchronize access to 'documents'
   @@mutex = Mutex.new
 
-  public
   def register
+    if @bulk_size > 1000
+      raise LogStash::ConfigurationError, "Bulk size must be lower than '1000', currently '#{@bulk_size}'"
+    end
+
     Mongo::Logger.logger = @logger
     conn = Mongo::Client.new(@uri)
     @db = conn.use(@database)
 
-    if @bulk_size > 1000
-      raise LogStash::ConfigurationError, "Bulk size must be lower than '1000', currently '#{@bulk_size}'"
-    end
+    @closed = Concurrent::AtomicBoolean.new(false)
     @documents = {}
-    Thread.new do
-      loop do
-        sleep @bulk_interval
+
+    @bulk_thread = Thread.new(@bulk_interval) do |bulk_interval|
+      while @closed.false? do
+        sleep(bulk_interval)
+
         @@mutex.synchronize do
           @documents.each do |collection, values|
             if values.length > 0
@@ -69,23 +72,31 @@ def register
         end
       end
     end
-  end # def register
+  end
 
   def receive(event)
     begin
       # Our timestamp object now has a to_bson method, using it here
       # {}.merge(other) so we don't taint the event hash innards
       document = {}.merge(event.to_hash)
+
       if !@isodate
-        # not using timestamp.to_bson
-        document["@timestamp"] = event.timestamp.to_json
+        timestamp = event.timestamp
+        if timestamp
+          # not using timestamp.to_bson
+          document["@timestamp"] = timestamp.to_json
+        else
+          @logger.warn("Cannot set MongoDB document `@timestamp` field because it does not exist in the event", :event => event)
+        end
       end
+
       if @generateId
-        document["_id"] = BSON::ObjectId.new(nil, event.timestamp)
+        document["_id"] = BSON::ObjectId.new
       end
+
       if @bulk
+        collection = event.sprintf(@collection)
         @@mutex.synchronize do
-          collection = event.sprintf(@collection)
           if(!@documents[collection])
             @documents[collection] = []
           end
@@ -99,20 +110,25 @@ def receive(event)
       else
         @db[event.sprintf(@collection)].insert_one(document)
       end
-
     rescue => e
-      @logger.warn("Failed to send event to MongoDB", :event => event, :exception => e,
-                   :backtrace => e.backtrace)
       if e.message =~ /^E11000/
-          # On a duplicate key error, skip the insert.
-          # We could check if the duplicate key err is the _id key
-          # and generate a new primary key.
-          # If the duplicate key error is on another field, we have no way
-          # to fix the issue.
+        # On a duplicate key error, skip the insert.
+        # We could check if the duplicate key err is the _id key
+        # and generate a new primary key.
+        # If the duplicate key error is on another field, we have no way
+        # to fix the issue.
+        @logger.warn("Skipping insert because of a duplicate key error", :event => event, :exception => e)
       else
-        sleep @retry_delay
+        @logger.warn("Failed to send event to MongoDB, retrying in #{@retry_delay.to_s} seconds", :event => event, :exception => e)
+        sleep(@retry_delay)
         retry
       end
     end
-  end # def receive
-end # class LogStash::Outputs::Mongodb
+  end
+
+  def close
+    @closed.make_true
+    @bulk_thread.wakeup
+    @bulk_thread.join
+  end
+end

spec/outputs/mongodb_spec.rb

@@ -8,23 +8,21 @@
   let(:database)   { 'logstash' }
   let(:collection) { 'logs' }
 
-  let(:config) do
-    { "uri" => uri, "database" => database, "collection" => collection }
-  end
+  let(:config) {{
+    "uri" => uri,
+    "database" => database,
+    "collection" => collection
+  }}
 
-  it "should register" do
+  it "should register and close" do
     plugin = LogStash::Plugin.lookup("output", "mongodb").new(config)
     expect {plugin.register}.to_not raise_error
+    plugin.close
   end
 
-  describe "#send" do
-
+  describe "receive" do
     subject! { LogStash::Outputs::Mongodb.new(config) }
 
-    let(:properties) { { "message" => "This is a message!",
-                         "uuid" => SecureRandom.uuid,
-                         "number" => BigDecimal.new("4321.1234"),
-                         "utf8" => "żółć"} }
     let(:event)      { LogStash::Event.new(properties) }
     let(:connection) { double("connection") }
     let(:client)     { double("client") }
@@ -38,10 +36,51 @@
       subject.register
     end
 
-    it "should send the event to the database" do
-      expect(collection).to receive(:insert_one)
-      subject.receive(event)
+    after(:each) do
+      subject.close
+    end
+
+    describe "#send" do
+      let(:properties) {{
+        "message" => "This is a message!",
+        "uuid" => SecureRandom.uuid,
+        "number" => BigDecimal.new("4321.1234"),
+        "utf8" => "żółć"
+      }}
+
+      it "should send the event to the database" do
+        expect(collection).to receive(:insert_one)
+        subject.receive(event)
+      end
     end
-  end
 
+    describe "no event @timestamp" do
+      let(:properties) { { "message" => "foo" } }
+
+      it "should not contain a @timestamp field in the mongo document" do
+        expect(event).to receive(:timestamp).and_return(nil)
+        expect(event).to receive(:to_hash).and_return(properties)
+        expect(collection).to receive(:insert_one).with(properties)
+        subject.receive(event)
+      end
+    end
+
+    describe "generateId" do
+      let(:properties) { { "message" => "foo" } }
+      let(:config) {{
+          "uri" => uri,
+          "database" => database,
+          "collection" => collection,
+          "generateId" => true
+      }}
+
+      it "should contain a BSON::ObjectId as _id" do
+        expect(BSON::ObjectId).to receive(:new).and_return("BSON::ObjectId")
+        expect(event).to receive(:timestamp).and_return(nil)
+        expect(event).to receive(:to_hash).and_return(properties)
+        expect(collection).to receive(:insert_one).with(properties.merge("_id" => "BSON::ObjectId"))
+        subject.receive(event)
+      end
+    end
+  end
 end