Handle YAML configuration format on configuration file

lib/fluent/command/fluentd.rb

@@ -126,6 +126,14 @@
   opts[:without_source] = b
 }
 
+op.on('--guess-config-file-type VALU', 'guessing file type of fluentd configuration. yaml or yml') { |s|
+  if (s == 'yaml') || (s == 'yml')
+    opts[:guess_config_file_type] = s.to_sym
+  else
+    usage '--guess-config-file-type accepts yaml or yml'
+  end
+}
+
 op.on('--use-v1-config', "Use v1 configuration format (default)", TrueClass) {|b|
   opts[:use_v1_config] = b
 }

lib/fluent/config.rb

@@ -17,6 +17,7 @@
 require 'fluent/config/error'
 require 'fluent/config/element'
 require 'fluent/configurable'
+require 'fluent/config/yaml_parser'
 
 module Fluent
   module Config
@@ -25,7 +26,16 @@ module Config
     # @param additional_config [String] config which is added to last of config body
     # @param use_v1_config [Bool] config is formatted with v1 or not
     # @return [Fluent::Config]
-    def self.build(config_path:, encoding: 'utf-8', additional_config: nil, use_v1_config: true)
+    def self.build(config_path:, encoding: 'utf-8', additional_config: nil, use_v1_config: true, type: nil)
+      config_file_ext = File.extname(config_path)
+      if config_file_ext == '.yaml' || config_file_ext == '.yml'
+        type = :yaml
+      end
+
+      if type == :yaml || type == :yml
+        return Fluent::Config::YamlParser.parse(config_path)
+      end
+
       config_fname = File.basename(config_path)
       config_basedir = File.dirname(config_path)
       config_data = File.open(config_path, "r:#{encoding}:utf-8") do |f|
@@ -36,6 +46,7 @@ def self.build(config_path:, encoding: 'utf-8', additional_config: nil, use_v1_c
         end
         s
       end
+
       Fluent::Config.parse(config_data, config_fname, config_basedir, use_v1_config)
     end
 

lib/fluent/config/yaml_parser.rb

@@ -0,0 +1,56 @@
+#
+# Fluentd
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+
+require 'fluent/config/yaml_parser/loader'
+require 'fluent/config/yaml_parser/parser'
+require 'pathname'
+
+module Fluent
+  module Config
+    module YamlParser
+      def self.parse(path)
+        context = Kernel.binding
+
+        unless context.respond_to?(:use_nil)
+          context.define_singleton_method(:use_nil) do
+            raise Fluent::SetNil
+          end
+        end
+
+        unless context.respond_to?(:use_default)
+          context.define_singleton_method(:use_default) do
+            raise Fluent::SetDefault
+          end
+        end
+
+        unless context.respond_to?(:hostname)
+          context.define_singleton_method(:hostname) do
+            Socket.gethostname
+          end
+        end
+
+        unless context.respond_to?(:worker_id)
+          context.define_singleton_method(:worker_id) do
+            ENV['SERVERENGINE_WORKER_ID'] || ''
+          end
+        end
+
+        s = Fluent::Config::YamlParser::Loader.new(context).load(Pathname.new(path))
+        Fluent::Config::YamlParser::Parser.new(s).build.to_element
+      end
+    end
+  end
+end

lib/fluent/config/yaml_parser/fluent_value.rb

@@ -0,0 +1,47 @@
+#
+# Fluentd
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+
+module Fluent
+  module Config
+    module YamlParser
+      module FluentValue
+        JsonValue = Struct.new(:val) do
+          def to_s
+            val.to_json
+          end
+
+          def to_element
+            to_s
+          end
+        end
+
+        StringValue = Struct.new(:val, :context) do
+          def to_s
+            context.instance_eval("\"#{val}\"")
+          rescue Fluent::SetNil => _
+            ''
+          rescue Fluent::SetDefault => _
+            ':default'
+          end
+
+          def to_element
+            to_s
+          end
+        end
+      end
+    end
+  end
+end

lib/fluent/config/yaml_parser/loader.rb

@@ -0,0 +1,91 @@
+#
+# Fluentd
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+
+require 'psych'
+require 'json'
+require 'fluent/config/error'
+require 'fluent/config/yaml_parser/fluent_value'
+
+# Based on https://github.com/eagletmt/hako/blob/34cdde06fe8f3aeafd794be830180c3cedfbb4dc/lib/hako/yaml_loader.rb
+
+module Fluent
+  module Config
+    module YamlParser
+      class Loader
+        INCLUDE_TAG = 'tag:include'.freeze
+        FLUENT_JSON_TAG = 'tag:fluent/json'.freeze
+        FLUENT_STR_TAG = 'tag:fluent/s'.freeze
+        SHOVEL = '<<'.freeze
+
+        def initialize(context = Kernel.binding)
+          @context = context
+          @current_path = nil
+        end
+
+        # @param [String] path
+        # @return [Hash]
+        def load(path)
+          class_loader = Psych::ClassLoader.new
+          scanner = Psych::ScalarScanner.new(class_loader)
+
+          visitor = Visitor.new(scanner, class_loader)
+
+          visitor._register_domain(INCLUDE_TAG) do |_, val|
+            load(path.parent.join(val))
+          end
+
+          visitor._register_domain(FLUENT_JSON_TAG) do |_, val|
+            Fluent::Config::YamlParser::FluentValue::JsonValue.new(val)
+          end
+
+          visitor._register_domain(FLUENT_STR_TAG) do |_, val|
+            Fluent::Config::YamlParser::FluentValue::StringValue.new(val, @context)
+          end
+
+          path.open do |f|
+            visitor.accept(Psych.parse(f))
+          end
+        end
+
+        class Visitor < Psych::Visitors::ToRuby
+          def initialize(scanner, class_loader)
+            super(scanner, class_loader)
+          end
+
+          def _register_domain(name, &block)
+            @domain_types.merge!({ name => [name, block] })
+          end
+
+          def revive_hash(hash, o)
+            super(hash, o).tap do |r|
+              if r[SHOVEL].is_a?(Hash)
+                h2 = {}
+                r.each do |k, v|
+                  if k == SHOVEL
+                    h2.merge!(v)
+                  else
+                    h2[k] = v
+                  end
+                end
+                r.replace(h2)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end