diff --git a/lib/rackstash/logger.rb b/lib/rackstash/logger.rb index ac9636f..4a51425 100644 --- a/lib/rackstash/logger.rb +++ b/lib/rackstash/logger.rb @@ -57,13 +57,24 @@ module Rackstash # By default we use {PROGNAME}. # @param formatter [#call] the log formatter for each individual buffered # line. See {#formatter} for details. - def initialize(*flows, level: DEBUG, progname: PROGNAME, formatter: Formatter.new) + # @yieldparam flow [Rackstash::Flow] if the given block accepts an argument, + # we yield the last {Flow} as a parameter. Without an expected argument, + # the block is directly executed in the context of the last {Flow}. + def initialize(*flows, level: DEBUG, progname: PROGNAME, formatter: Formatter.new, &block) @buffer_stack = Concurrent::ThreadLocalVar.new @sink = Rackstash::Sink.new(*flows) self.level = level self.progname = progname self.formatter = formatter + + if block_given? && (flow = @sink.flows.last) + if block.arity == 0 + flow.instance_eval(&block) + else + yield flow + end + end end # Add a message to the current buffer without any further formatting. If diff --git a/spec/rackstash/logger_spec.rb b/spec/rackstash/logger_spec.rb index d008434..a523451 100644 --- a/spec/rackstash/logger_spec.rb +++ b/spec/rackstash/logger_spec.rb @@ -35,6 +35,35 @@ describe Rackstash::Logger do logger = described_class.new('output.log', formatter: -> {}) expect(logger.formatter).to be_a Proc end + + it 'yields the last flow to a parameterized block' do + block_called = 0 + block_self = nil + block_args = nil + + described_class.new(StringIO.new, StringIO.new) do |*args| + block_called += 1 + block_self = self + block_args = args + end + + expect(block_called).to eql 1 + expect(block_self).to equal self + expect(block_args).to match [instance_of(Rackstash::Flow)] + end + + it 'instance_evals the parameter-less block in the last flow' do + block_called = 0 + block_self = nil + + described_class.new(StringIO.new, StringIO.new) do + block_called += 1 + block_self = self + end + + expect(block_called).to eql 1 + expect(block_self).to be_instance_of(Rackstash::Flow) + end end describe 'subscript accessors' do