From 013d0f7d92d79292c0b1f183b154aa28a9e1b23f Mon Sep 17 00:00:00 2001
From: Holger Just <hello@holgerjust.de>
Date: Thu, 17 Aug 2017 00:22:49 +0200
Subject: [PATCH] Auto-flush the buffer when adding fields using Logger#add

With this interface, the user adds fields the same way they would add
messages. The implicit assumption here is that the buffer is handled the
same way. Thus, if the current buffer is non-buffering, we will
automatically flush it to the sink and clear it, just the same way as we
would have done it for a message.
---
 lib/rackstash/buffer.rb       | 23 ++++++++++++
 lib/rackstash/logger.rb       |  3 +-
 spec/rackstash/buffer_spec.rb | 66 +++++++++++++++++++++++++++++++++++
 spec/rackstash/logger_spec.rb |  7 ++--
 4 files changed, 92 insertions(+), 7 deletions(-)

diff --git a/lib/rackstash/buffer.rb b/lib/rackstash/buffer.rb
index f64847a..d22b9a4 100644
--- a/lib/rackstash/buffer.rb
+++ b/lib/rackstash/buffer.rb
@@ -118,6 +118,29 @@ module Rackstash
       exception
     end
 
+    # Deep-merge fields to the buffer. This will mark the current buffer as
+    # {pending?} and will result in the eventual flush of the logged data.
+    #
+    # The buffer's timestamp will be initialized with the current time if it
+    # wasn't set earlier already.
+    #
+    # If the buffer is not {#buffering?}, it will be {#flush}ed and {#clear}ed
+    # after each added message. All fields, tags, and messages added before as
+    # well as the fields added with this method call will be flushed.
+    #
+    # @param hash (see Fields::Hash#deep_merge!)
+    # @raise (see Fields::Hash#deep_merge!)
+    # @return [Rackstash::Fields::Hash, ::Hash, Proc] the given `hash` value
+    #
+    # @see Fields::Hash#deep_merge!
+    def add_fields(hash)
+      timestamp
+      fields.deep_merge!(hash, force: true)
+      auto_flush
+
+      hash
+    end
+
     # Add a new message to the buffer. This will mark the current buffer as
     # {pending?} and will result in the eventual flush of the logged data.
     #
diff --git a/lib/rackstash/logger.rb b/lib/rackstash/logger.rb
index 937dd1d..b59fdea 100644
--- a/lib/rackstash/logger.rb
+++ b/lib/rackstash/logger.rb
@@ -322,8 +322,7 @@ module Rackstash
 
       case msg
       when Hash, Rackstash::Fields::Hash
-        buffer.fields.deep_merge!(msg)
-        msg
+        buffer.add_fields(msg)
       else
         time = Time.now.utc.freeze
         buffer.add_message Message.new(
diff --git a/spec/rackstash/buffer_spec.rb b/spec/rackstash/buffer_spec.rb
index 893dc2b..c8ffc67 100644
--- a/spec/rackstash/buffer_spec.rb
+++ b/spec/rackstash/buffer_spec.rb
@@ -77,6 +77,72 @@ describe Rackstash::Buffer do
     end
   end
 
+  describe '#add_fields' do
+    it 'deep-merges fields' do
+      buffer.add_fields(foo: :bar, number: 123)
+
+      expect(buffer.fields['foo']).to eql 'bar'
+      expect(buffer.fields['number']).to eql 123
+    end
+
+    it 'overwrites fields' do
+      buffer.fields['foo'] = 'initial'
+      buffer.add_fields(foo: 'overwritten')
+
+      expect(buffer.fields['foo']).to eql 'overwritten'
+    end
+
+    it 'raises ArgumentError when trying to set a forbidden key' do
+      expect { buffer.add_fields(message: 'oh no!') }.to raise_error ArgumentError
+    end
+
+    it 'sets the timestamp' do
+      expect(buffer).to receive(:timestamp)
+      buffer.add_fields(key: 'value')
+    end
+
+    context 'when buffering?' do
+      before do
+        buffer_options[:buffering] = true
+      end
+
+      it 'does not call #flush' do
+        expect(buffer).not_to receive(:flush)
+        buffer.add_fields(key: 'value')
+      end
+
+      it 'does not call #clear' do
+        expect(buffer).not_to receive(:clear)
+        buffer.add_fields(key: 'value')
+        expect(buffer.fields['key']).to eql 'value'
+      end
+    end
+
+    context 'when not buffering?' do
+      before do
+        buffer_options[:buffering] = false
+      end
+
+      it 'calls #flush' do
+        expect(buffer).to receive(:flush)
+        buffer.add_fields(key: 'value')
+      end
+
+      it 'calls #clear' do
+        allow(buffer).to receive(:flush)
+        expect(buffer).to receive(:clear).and_call_original
+        buffer.add_fields(key: 'value')
+        expect(buffer.fields).to be_empty
+        expect(buffer.pending?).to be false
+      end
+    end
+
+    it 'returns the given value' do
+      fields = {key: 'value'}
+      expect(buffer.add_fields(fields)).to equal fields
+    end
+  end
+
   describe '#add_message' do
     it 'adds a message to the buffer' do
       msg = double(message: 'Hello World', time: Time.now)
diff --git a/spec/rackstash/logger_spec.rb b/spec/rackstash/logger_spec.rb
index a9a3ecd..d008434 100644
--- a/spec/rackstash/logger_spec.rb
+++ b/spec/rackstash/logger_spec.rb
@@ -414,13 +414,10 @@ describe Rackstash::Logger do
     end
 
     it 'merges fields if message is a Hash' do
-      fields = instance_double('Rackstash::Fields::Hash')
-      expect(buffer).to receive(:fields).and_return(fields)
-      expect(fields).to receive(:deep_merge!).with(foo: 'bar')
-
+      expect(buffer).to receive(:add_fields).with(foo: 'bar')
       expect(buffer).not_to receive(:add_message)
 
-      expect(logger.add(0, foo: 'bar')).to eql(foo: 'bar')
+      logger.add(0, foo: 'bar')
     end
 
     it 'can use debug shortcut' do