Use bucket based on placeholder

Signed-off-by: rmontenegroo <[email protected]>

Author: rmontenegroo

.gitignore

@@ -11,3 +11,5 @@ vendor
 .ruby-version
 
 test/tmp/
+
+docker/

VERSION

@@ -1 +1 @@
-1.4.0
+1.6.0

lib/fluent/plugin/out_s3.rb

@@ -87,6 +87,8 @@ def initialize
     config_param :aws_iam_retries, :integer, default: nil, deprecated: "Use 'instance_profile_credentials' instead"
     desc "S3 bucket name"
     config_param :s3_bucket, :string
+    desc "Set bucket name fallback if fails fetching from placeholders"
+    config_param :s3_bucket_fallback, :string, :default => nil
     desc "S3 region name"
     config_param :s3_region, :string, default: ENV["AWS_REGION"] || "us-east-1"
     desc "Use 's3_region' instead"
@@ -249,11 +251,6 @@ def start
 
       s3_client = Aws::S3::Client.new(options)
       @s3 = Aws::S3::Resource.new(client: s3_client)
-      @bucket = @s3.bucket(@s3_bucket)
-
-      check_apikeys if @check_apikey_on_start
-      ensure_bucket if @check_bucket
-      ensure_bucket_lifecycle
 
       super
     end
@@ -264,6 +261,7 @@ def format(tag, time, record)
     end
 
     def write(chunk)
+
       i = 0
       metadata = chunk.metadata
       previous_path = nil
@@ -273,6 +271,40 @@ def write(chunk)
                      @time_slice_with_tz.call(metadata.timekey)
                    end
 
+      bucket_name = nil
+
+      if @s3_bucket.match(/\$\{.*\}/)
+
+        bucket_name = extract_placeholders(@s3_bucket, chunk)
+
+        if bucket_name.match(/\$\{.*\}/)
+
+            log.warn "Trying to use @s3_bucket_fallback as a fallback bucket name"
+            
+            if @s3_bucket_fallback
+              
+              bucket_name = @s3_bucket_fallback
+            
+            else
+
+              raise "It was possible to extract_placeholder from @s3_bucket and there is no @s3_bucket_fallback set"
+
+            end
+
+        end
+         
+      else
+
+        bucket_name = @s3_bucket
+
+      end
+
+      bucket = @s3.bucket(bucket_name)
+
+      check_apikeys(bucket) if @check_apikey_on_start
+      ensure_bucket(bucket) if @check_bucket
+      ensure_bucket_lifecycle(bucket)
+
       if @check_object
         begin
           @values_for_s3_object_chunk[chunk.unique_id] ||= {
@@ -304,7 +336,7 @@ def write(chunk)
 
           i += 1
           previous_path = s3path
-        end while @bucket.object(s3path).exists?
+        end while bucket.object(s3path).exists?
       else
         if @localtime
           hms_slicer = Time.now.strftime("%H%M%S")
@@ -362,18 +394,19 @@ def write(chunk)
             put_options[:metadata][k] = extract_placeholders(v, chunk).gsub(%r(%{[^}]+}), {"%{index}" => sprintf(@index_format, i - 1)})
           end
         end
-        @bucket.object(s3path).put(put_options)
+        bucket.object(s3path).put(put_options)
 
         @values_for_s3_object_chunk.delete(chunk.unique_id)
 
         if @warn_for_delay
           if Time.at(chunk.metadata.timekey) < Time.now - @warn_for_delay
-            log.warn "out_s3: delayed events were put to s3://#{@s3_bucket}/#{s3path}"
+            log.warn "out_s3: delayed events were put to s3://#{bucket.name}/#{s3path}"
           end
         end
       ensure
         tmp.close(true) rescue nil
       end
+
     end
 
     private
@@ -399,40 +432,40 @@ def timekey_to_timeformat(timekey)
       end
     end
 
-    def ensure_bucket
-      if !@bucket.exists?
+    def ensure_bucket(bucket)
+      if !bucket.exists?
         if @auto_create_bucket
-          log.info "Creating bucket #{@s3_bucket} on #{@s3_endpoint}"
-          @s3.create_bucket(bucket: @s3_bucket)
+          log.info "Creating bucket #{bucket.name} on #{@s3_endpoint}"
+          @s3.create_bucket(bucket: bucket.name)          
         else
-          raise "The specified bucket does not exist: bucket = #{@s3_bucket}"
+          raise "The specified bucket does not exist: bucket = #{bucket.name}"
         end
       end
     end
 
-    def ensure_bucket_lifecycle
+    def ensure_bucket_lifecycle(bucket)
       unless @bucket_lifecycle_rules.empty?
-        old_rules = get_bucket_lifecycle_rules
+        old_rules = get_bucket_lifecycle_rules(bucket)
         new_rules = @bucket_lifecycle_rules.sort_by { |rule| rule.id }.map do |rule|
           { id: rule.id, expiration: { days: rule.expiration_days }, prefix: rule.prefix, status: "Enabled" }
         end
 
         unless old_rules == new_rules
-          log.info "Configuring bucket lifecycle rules for #{@s3_bucket} on #{@s3_endpoint}"
-          @bucket.lifecycle_configuration.put({ lifecycle_configuration: { rules: new_rules } })
+          log.info "Configuring bucket lifecycle rules for #{bucket.name} on #{@s3_endpoint}"
+          bucket.lifecycle_configuration.put({ lifecycle_configuration: { rules: new_rules } })
         end
       end
     end
 
-    def get_bucket_lifecycle_rules
+    def get_bucket_lifecycle_rules(bucket)
       begin
-        @bucket.lifecycle_configuration.rules.sort_by { |rule| rule[:id] }.map do |rule|
+        bucket.lifecycle_configuration.rules.sort_by { |rule| rule[:id] }.map do |rule|
           { id: rule[:id], expiration: { days: rule[:expiration][:days] }, prefix: rule[:prefix], status: rule[:status] }
         end
       rescue Aws::S3::Errors::NoSuchLifecycleConfiguration
         []
       end
-    end
+     end
 
     def process_s3_object_key_format
       %W(%{uuid} %{uuid:random} %{uuid:hostname} %{uuid:timestamp}).each { |ph|
@@ -461,8 +494,8 @@ def check_s3_path_safety(conf)
       end
     end
 
-    def check_apikeys
-      @bucket.objects(prefix: @path, :max_keys => 1).first
+    def check_apikeys(bucket)
+      bucket.objects(prefix: @path, :max_keys => 1).first
     rescue Aws::S3::Errors::NoSuchBucket
       # ignore NoSuchBucket Error because ensure_bucket checks it.
     rescue => e