Merge pull request #1438 from fluent/local-storage-path-using-usage

fix to generate storage local path using usage
fluent · Jan 30, 2017 · 59ab3e6 · 59ab3e6
2 parents b710417 + 57f957a
commit 59ab3e6
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 2 deletions.
diff --git a/lib/fluent/plugin/storage_local.rb b/lib/fluent/plugin/storage_local.rb
@@ -64,7 +64,8 @@ def configure(conf)
             end
           end
         elsif root_dir = owner.plugin_root_dir
-          @path = File.join(root_dir, 'storage.json')
+          basename = (conf.arg && !conf.arg.empty?) ? "storage.#{conf.arg}.json" : "storage.json"
+          @path = File.join(root_dir, basename)
           @multi_workers_available = true
         else
           if @persistent

diff --git a/lib/fluent/plugin_helper/storage.rb b/lib/fluent/plugin_helper/storage.rb
@@ -33,6 +33,9 @@ def storage_create(usage: '', type: nil, conf: nil, default_type: nil)
         if conf && !conf.arg.empty?
           usage = conf.arg
         end
+        if !usage.empty? && usage !~ /^[a-zA-Z][-_.a-zA-Z0-9]*$/
+          raise Fluent::ConfigError, "Argument in <storage ARG> uses invalid characters: '#{usage}'"
+        end
 
         s = @_storages[usage]
         if s && s.running
@@ -59,7 +62,7 @@ def storage_create(usage: '', type: nil, conf: nil, default_type: nil)
                      conf = Hash[conf.map{|k,v| [k.to_s, v]}]
                      Fluent::Config::Element.new('storage', usage, conf, [])
                    when nil
-                     Fluent::Config::Element.new('storage', usage, {}, [])
+                     Fluent::Config::Element.new('storage', usage, {'@type' => type}, [])
                    else
                      raise ArgumentError, "BUG: conf must be a Element, Hash (or unspecified), but '#{conf.class}'"
                    end
@@ -98,6 +101,9 @@ def configure(conf)
         super
 
         @storage_configs.each do |section|
+          if !section.usage.empty? && section.usage !~ /^[a-zA-Z][-_.a-zA-Z0-9]*$/
+            raise Fluent::ConfigError, "Argument in <storage ARG> uses invalid characters: '#{section.usage}'"
+          end
           if @_storages[section.usage]
             raise Fluent::ConfigError, "duplicated storages configured: #{section.usage}"
           end

diff --git a/test/plugin/test_storage_local.rb b/test/plugin/test_storage_local.rb
@@ -186,6 +186,20 @@ class MyInput < Fluent::Plugin::Input
       assert_equal '2', @p.get('key1')
       assert_equal 4, @p.get('key2')
     end
+
+    test 'works with customized path by specified usage' do
+      root_dir = File.join(TMP_DIR, 'root')
+      expected_storage_path = File.join(root_dir, 'worker0', 'local_storage_test', 'storage.usage.json')
+      conf = config_element('ROOT', 'usage', {'@id' => 'local_storage_test'})
+      Fluent::SystemConfig.overwrite_system_config('root_dir' => root_dir) do
+        @d.configure(conf)
+      end
+      @d.start
+      @p = @d.storage_create(usage: 'usage', type: 'local')
+
+      assert_equal expected_storage_path, @p.path
+      assert @p.store.empty?
+    end
   end
 
   sub_test_case 'configured with root-dir and plugin id, and multi workers' do

diff --git a/test/plugin_helper/test_storage.rb b/test/plugin_helper/test_storage.rb
@@ -135,6 +135,17 @@ class Dummy2 < Fluent::Plugin::TestBase
     end
   end
 
+  test 'raises config error if config argument has invalid characters' do
+    d = Dummy.new
+    assert_raise Fluent::ConfigError.new("Argument in <storage ARG> uses invalid characters: 'yaa y'") do
+      d.configure(config_element('root', '', {}, [config_element('storage', 'yaa y', {'@type' => 'local'})]))
+    end
+    d.configure(config_element())
+    assert_raise Fluent::ConfigError.new("Argument in <storage ARG> uses invalid characters: 'a,b'") do
+      d.storage_create(usage: 'a,b', type: 'local')
+    end
+  end
+
   test 'can be configured without storage sections' do
     d = Dummy.new
     assert_nothing_raised do