middleware.rb in google-cloud-logging-0.23.1

- old
+ new

@@ -11,15 +11,22 @@
 # 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 "stackdriver/core"
+
 module Google
   module Cloud
     module Logging
       class Middleware
         ##
+        # A default value for the log_name_map argument. Directs health check
+        # logs to a separate log name so they don't spam the main log.
+        DEFAULT_LOG_NAME_MAP = { "/_ah/health" => "ruby_health_check_log" }
+
+        ##
         # The Google::Cloud::Logging::Logger instance
         attr_reader :logger
 
         ##
         # Create a new AppEngine logging Middleware.
@@ -27,17 +34,22 @@
         # @param [Rack Application] app Rack application
         # @param [Google::Cloud::Logging::Logger] logger A logger to be used by
         #   this middleware. The middleware will be interacting with the logger
         #   to track Stackdriver request trace ID. It also properly sets
         #   env["rack.logger"] to this assigned logger for accessing.
+        # @param [Hash] log_name_map A map from URI path to log name override.
+        #   The path may be a string or regex. If a request path matches an
+        #   entry in this map, the log name is set accordingly, otherwise the
+        #   logger's default log_name is used.
         #
         # @return [Google::Cloud::Logging::Middleware] A new
         #   Google::Cloud::Logging::Middleware instance
         #
-        def initialize app, logger: nil
+        def initialize app, logger: nil, log_name_map: DEFAULT_LOG_NAME_MAP
           @app = app
           @logger = logger
+          @log_name_map = log_name_map
         end
 
         ##
         # Rack middleware entry point. In most Rack based frameworks, a request
         # is served by one thread. So entry point, we associate the GCP request
@@ -49,32 +61,51 @@
         #
         # @return [Rack::Response] The response from downstream Rack app
         #
         def call env
           env["rack.logger"] = logger
-          trace_id = extract_trace_id(env)
-          logger.add_trace_id trace_id
-
+          trace_id = get_trace_id env
+          log_name = get_log_name env
+          logger.add_request_info trace_id: trace_id, log_name: log_name
           begin
             @app.call env
           ensure
-            logger.delete_trace_id
+            logger.delete_request_info
           end
         end
 
         ##
-        # Extract the trace_id from HTTP request header
-        # HTTP_X_CLOUD_TRACE_CONTEXT.
+        # Determine the trace ID for this request.
         #
-        # @return [String] The trace_id or nil if trace_id is empty.
-        def extract_trace_id env
-          trace_context = env["HTTP_X_CLOUD_TRACE_CONTEXT"].to_s
-          return nil if trace_context.empty?
-          trace_context.sub(%r{/.*}, "")
+        # @private
+        # @param [Hash] env The Rack environment.
+        # @return [String] The trace ID.
+        #
+        def get_trace_id env
+          trace_context = Stackdriver::Core::TraceContext.parse_rack_env env
+          trace_context.trace_id
         end
 
         ##
+        # Determine the log name override for this request, if any.
+        #
+        # @private
+        # @param [Hash] env The Rack environment.
+        # @return [String, nil] The log name override, or `nil` if there is
+        #     no override.
+        #
+        def get_log_name env
+          return nil unless @log_name_map
+          path = "#{env['SCRIPT_NAME']}#{env['PATH_INFO']}"
+          path = "/#{path}" unless path.start_with? "/"
+          @log_name_map.each do |k, v|
+            return v if k === path
+          end
+          nil
+        end
+
+        ##
         # Construct a monitored resource based on the given type and label if
         # both are provided. Otherwise, construct a default monitored resource
         # based on the current environment.
         #
         # @example If both type and labels are provided, it returns resource:
@@ -134,28 +165,32 @@
         ##
         # @private Extract information from current environment and construct
         # the correct monitoring resource types and labels.
         #
         # @example If running from GAE, returns default resource:
-        #   rc = Google::Cloud::Logging::Middleware.build_monitored_resource
+        #   rc = Google::Cloud::Logging::Middleware.send \
+        #          :default_monitored_resource
         #   rc.type   #=> "gae_app"
         #   rc.labels # { module_id: [GAE module name],
         #             #   version_id: [GAE module version] }
         #
         # @example If running from GKE, returns default resource:
-        #   rc = Google::Cloud::Logging::Middleware.build_monitored_resource
+        #   rc = Google::Cloud::Logging::Middleware.send \
+        #          :default_monitored_resource
         #   rc.type   #=> "container"
         #   rc.labels # { cluster_name: [GKE cluster name],
         #             #   namespace_id: [GKE namespace_id] }
         #
         # @example If running from GCE, return default resource:
-        #   rc = Google::Cloud::Logging::Middleware.build_monitored_resource
+        #   rc = Google::Cloud::Logging::Middleware.send \
+        #          :default_monitored_resource
         #   rc.type   #=> "gce_instance"
         #   rc.labels # { instance_id: [GCE VM instance id],
         #             #   zone: [GCE vm group zone] }
         #
         # @example Otherwise default to generic "global" type:
-        #   rc = Google::Cloud::Logging::Middleware.build_monitored_resource
+        #   rc = Google::Cloud::Logging::Middleware.send \
+        #          :default_monitored_resource
         #   rc.type   #=> "global"
         #   rc.labels #=> {}
         #
         # Reference https://cloud.google.com/logging/docs/api/ref_v2beta1/rest/v2beta1/MonitoredResource
         # for a full list of monitoring resources