lib/ur/middleware.rb in ur-0.1.0 vs lib/ur/middleware.rb in ur-0.1.1

@@ -1,5 +1,7 @@
+# frozen_string_literal: true
+
 require 'ur' unless Object.const_defined?(:Ur)
 
 class Ur
   module Middleware
     def initialize(app, options = {})
@@ -7,10 +9,19 @@
       @options = options
     end
     attr_reader :app
     attr_reader :options
 
+    def begin_request(ur)
+      ur.metadata.begin!
+    end
+
+    def finish_request(ur)
+      ur.logger_tags(@options[:logger])
+      ur.metadata.finish!
+    end
+
     def invoke_callback(name, *a, &b)
       if @options[name]
         @options[name].call(*a, &b)
       end
     end
@@ -19,23 +30,25 @@
   class FaradayMiddleware < ::Faraday::Middleware
     include Ur::Middleware
     def call(request_env)
       ur = Ur.from_faraday_request(request_env)
       invoke_callback(:before_request, ur)
-      ur.logger = options[:logger] if options[:logger]
+      begin_request(ur)
       ur.faraday_on_complete(@app, request_env) do |response_env|
+        finish_request(ur)
         invoke_callback(:after_response, ur)
       end
     end
   end
 
   class RackMiddleware
     include Ur::Middleware
     def call(env)
       ur = Ur.from_rack_request(env)
       invoke_callback(:before_request, ur)
-      ur.logger = options[:logger] if options[:logger]
+      begin_request(ur)
       ur.with_rack_response(@app, env) do
+        finish_request(ur)
         invoke_callback(:after_response, ur)
       end
     end
   end
 end