out_mongo: Support dynamic database name by built-in placeholders

Signed-off-by: Hiroshi Hatake <hatake@clear-code.com>

Author: cosmo0920

lib/fluent/plugin/out_mongo.rb

@@ -173,16 +173,17 @@ def multi_workers_ready?
 
     def write(chunk)
       collection_name = extract_placeholders(@collection, chunk.metadata)
-      operate(format_collection_name(collection_name), collect_records(chunk))
+      database_name = extract_placeholders(@database, chunk.metadata)
+      operate(database_name, format_collection_name(collection_name), collect_records(chunk))
     end
 
     private
 
-    def client
+    def client(database = @database)
       if @connection_string
         Mongo::Client.new(@connection_string)
       else
-        @client_options[:database] = @database
+        @client_options[:database] = database
         @client_options[:user] = @user if @user
         @client_options[:password] = @password if @password
         Mongo::Client.new(@nodes, @client_options)
@@ -227,7 +228,8 @@ def collection_exists?(name)
       end
     end
 
-    def get_collection(name, options)
+    def get_collection(database, name,  options)
+      @client = client(database) if database && @database != database
       return @client[name] if @collections[name]
 
       unless collection_exists?(name)
@@ -242,7 +244,7 @@ def forget_collection(name)
       @collections.delete(name)
     end
 
-    def operate(collection, records)
+    def operate(database, collection, records)
       begin
         if @replace_dot_in_key_with
           records.map! do |r|
@@ -255,7 +257,7 @@ def operate(collection, records)
           end
         end
 
-        get_collection(collection, @collection_options).insert_many(records)
+        get_collection(database, collection, @collection_options).insert_many(records)
       rescue Mongo::Error::BulkWriteError => e
         log.warn "#{records.size - e.result["n_inserted"]} documents are not inserted. Maybe these documents are invalid as a BSON."
         forget_collection(collection)

test/plugin/test_out_mongo.rb

@@ -36,9 +36,9 @@ def default_config
     ]
   end
 
-  def setup_mongod
+  def setup_mongod(database = database_name)
     options = {}
-    options[:database] = database_name
+    options[:database] = database
     @client = ::Mongo::Client.new(["localhost:#{port}"], options)
   end
 
@@ -211,6 +211,35 @@ def test_write_with_collection_placeholder
     end
   end
 
+  class WriteWithDatabasePlaceholder < self
+    def setup
+      @tag = 'custom'
+      setup_mongod(@tag)
+    end
+
+    def teardown
+      teardown_mongod
+    end
+
+    def test_write_with_database_placeholder
+      d = create_driver(%[
+        @type mongo
+        database ${tag}
+        collection #{collection_name}
+        include_time_key true
+      ])
+      d.run(default_tag: @tag) do
+        emit_documents(d)
+      end
+
+      actual_documents = get_documents
+      time = event_time("2011-01-02 13:14:15 UTC")
+      expected = [{'a' => 1, d.instance.inject_config.time_key => Time.at(time).localtime},
+                  {'a' => 2, d.instance.inject_config.time_key => Time.at(time).localtime}]
+      assert_equal(expected, actual_documents)
+    end
+  end
+
   def test_write_at_enable_tag
     d = create_driver(default_config + %[
       include_tag_key true