Merge pull request #2977 from fluent/record_accessor-set-support

record_accessor: Add set method for nested record

lib/fluent/plugin_helper/record_accessor.rb

@@ -41,9 +41,11 @@ def initialize(param)
             else
               mcall = method(:call_dig)
               mdelete = method(:delete_nest)
+              mset = method(:set_nest)
               singleton_class.module_eval do
                 define_method(:call, mcall)
                 define_method(:delete, mdelete)
+                define_method(:set, mset)
               end
             end
           end
@@ -75,6 +77,18 @@ def delete_nest(r)
           nil
         end
 
+        def set(r, v)
+          r[@keys] = v
+        end
+
+        # set_nest doesn't create intermediate object. If key doesn't exist, no effect.
+        # See also: https://bugs.ruby-lang.org/issues/11747
+        def set_nest(r, v)
+          r.dig(*@dig_keys)&.[]=(@last_key, v)
+        rescue
+          nil
+        end
+
         def self.parse_parameter(param)
           if param.start_with?('$.')
             parse_dot_notation(param)

test/plugin_helper/test_record_accessor.rb

@@ -194,4 +194,45 @@ class Dummy < Fluent::Plugin::TestBase
       end
     end
   end
+
+  sub_test_case 'Fluent::PluginHelper::RecordAccessor::Accessor#set' do
+    setup do
+      @d = Dummy.new
+    end
+
+    data('normal' => 'key1',
+         'space' => 'ke y2',
+         'dot key' => 'this.is.key3')
+    test 'set top key' do |param|
+      r = {'key1' => 'v1', 'ke y2' => 'v2', 'this.is.key3' => 'v3'}
+      accessor = @d.record_accessor_create(param)
+      accessor.set(r, "test")
+      assert_equal "test", r[param]
+    end
+
+    test "set top key using bracket style" do
+      r = {'key1' => 'v1', 'ke y2' => 'v2', 'this.is.key3' => 'v3'}
+      accessor = @d.record_accessor_create('$["this.is.key3"]')
+      accessor.set(r, "test")
+      assert_equal "test", r["this.is.key3"]
+    end
+
+    data('bracket' => "$['key1'][0]['ke y2']",
+         'bracket w/ double quotes' => '$["key1"][0]["ke y2"]')
+    test "set nested keys ['key1', 0, 'ke y2']" do |param|
+      r = {'key1' => [{'ke y2' => "value"}]}
+      accessor = @d.record_accessor_create(param)
+      accessor.set(r, "nested_message")
+      assert_equal "nested_message", r['key1'][0]['ke y2']
+    end
+
+    test "don't raise an error when unexpected record is coming" do
+      r = {'key1' => [{'key3' => "value"}]}
+      accessor = @d.record_accessor_create("$['key1']['key2']['key3']")
+      assert_nothing_raised do
+        accessor.set(r, "unknown field")
+      end
+      assert_equal({'key1' => [{'key3' => "value"}]}, r)
+    end
+  end
 end