From 031198cf146e4e52035c89caf693409028e0f30e Mon Sep 17 00:00:00 2001
From: Holger Just <hello@holgerjust.de>
Date: Thu, 9 Feb 2017 00:15:15 +0100
Subject: [PATCH] Add Hash#set to set a field from a block if it doesn't exist
 yet

This allows to define default values for certain fields which can be
inserted just before a Buffer is flushed. They won;t overwrite prior
user-provided fields.

Through the use of a block, expensive calculations for a field could be
avoided if the field is not going to be inserted, e.g. because it exists
already or is forbidden from being set.
---
 lib/rackstash/fields/hash.rb       | 27 ++++++++++++++++++++++++
 spec/rackstash/fields/hash_spec.rb | 33 ++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/lib/rackstash/fields/hash.rb b/lib/rackstash/fields/hash.rb
index f2c3120..a9fa271 100644
--- a/lib/rackstash/fields/hash.rb
+++ b/lib/rackstash/fields/hash.rb
@@ -183,6 +183,33 @@ module Rackstash
         @forbidden_keys.include?(key)
       end
 
+      # Set a `key` of `self` to the returned value of the passed block.
+      # If the key is forbidden from being set or already exists with a value
+      # other than `nil`, the block will not be called and the value will not be
+      # set.
+      #
+      # @param key [#to_s] the field name. When setting the field, this name
+      #   will be normalized as a frozen UTF-8 string.
+      #
+      # @yield [key] if the key doesn't exist yet, we call the block and use its
+      #    return value as the value to insert at `key`
+      # @yieldparam key [String] The normalized key where the value is being
+      #    inserted
+      # @yieldreturn [Object] the intended new value for `key` to be merged into
+      #   `self` at `key`.
+      #
+      # @return [Object, nil] The return value of the block or `nil` if no
+      #   insertion happened. Note that `nil` is also a valid value to insert
+      #   into the hash.
+      def set(key)
+        key = utf8_encode(key)
+
+        return if forbidden_key?(key)
+        return unless @raw[key] == nil
+
+        @raw[key] = normalize(yield(key))
+      end
+
       # @return [::Array] a new array populated with the values from this hash.
       # @see #keys
       def values
diff --git a/spec/rackstash/fields/hash_spec.rb b/spec/rackstash/fields/hash_spec.rb
index 4989faa..167d376 100644
--- a/spec/rackstash/fields/hash_spec.rb
+++ b/spec/rackstash/fields/hash_spec.rb
@@ -317,6 +317,39 @@ describe Rackstash::Fields::Hash do
     end
   end
 
+  describe '#set' do
+    it 'allows to set a normalized value' do
+      expect(hash).to receive(:normalize).with(:value).and_call_original
+
+      hash.set(:symbol) { :value }
+
+      expect(hash['symbol']).to eql 'value'
+    end
+
+    it 'ignores forbidden keys' do
+      forbidden_keys << 'forbidden'
+
+      expect { |b| hash.set(:forbidden, &b) }.not_to yield_control
+      expect { |b| hash.set('forbidden', &b) }.not_to yield_control
+
+      expect(hash['forbidden']).to be_nil
+    end
+
+    it 'ignores existing keys' do
+      hash['key'] = 'value'
+
+      expect { |b| hash.set(:key, &b) }.not_to yield_control
+      expect { |b| hash.set('key', &b) }.not_to yield_control
+
+      expect(hash['key']).to eql 'value'
+    end
+
+    it 'overwrites nil value' do
+      hash['nil'] = nil
+      expect { |b| hash.set('nil', &b) }.to yield_control
+    end
+  end
+
   describe '#values' do
     it 'returns an array of values' do
       hash['string'] = 'beep'