Rails Integration for DynamoDB Profiling

Inspired by the riak integration, I wanted to have better overviews for my non-Activerecord app about what queries are actually being issued:

# config/initializers/dynamodb_profiler.rb

# Instrument Aws::DynamoDB::Client calls so they appear in Rack::MiniProfiler's SQL Summary.
# Works transparently with aws-record because aws-record uses Aws::DynamoDB::Client internally.

begin
  require 'rack-mini-profiler'
rescue LoadError
  # Rack::MiniProfiler not present; no-op.
end

begin
  require 'aws-sdk-dynamodb'
rescue LoadError
  # AWS SDK not present; no-op. The patch only applies when Aws::DynamoDB::Client exists.
end

require 'json'

if defined?(::Rack::MiniProfiler) && defined?(::Aws::DynamoDB::Client)
  Rails.application.config.after_initialize do
    # Avoid double-prepend on reload
    module RackMiniProfilerDynamoDB
      module Instrumentation
        # @param params [Hash] Request parameters for DynamoDB get_item.
        # @param options [Hash] Additional AWS SDK options for the call.
        # @return [Seahorse::Client::Response] DynamoDB response object.
        def get_item(params = {}, options = {})
          profile_dynamodb('get_item', params) { super }
        end

        # @param params [Hash] Request parameters for DynamoDB put_item.
        # @param options [Hash] Additional AWS SDK options for the call.
        # @return [Seahorse::Client::Response] DynamoDB response object.
        def put_item(params = {}, options = {})
          profile_dynamodb('put_item', params) { super }
        end

        # @param params [Hash] Request parameters for DynamoDB update_item.
        # @param options [Hash] Additional AWS SDK options for the call.
        # @return [Seahorse::Client::Response] DynamoDB response object.
        def update_item(params = {}, options = {})
          profile_dynamodb('update_item', params) { super }
        end

        # @param params [Hash] Request parameters for DynamoDB delete_item.
        # @param options [Hash] Additional AWS SDK options for the call.
        # @return [Seahorse::Client::Response] DynamoDB response object.
        def delete_item(params = {}, options = {})
          profile_dynamodb('delete_item', params) { super }
        end

        # @param params [Hash] Request parameters for DynamoDB query.
        # @param options [Hash] Additional AWS SDK options for the call.
        # @return [Seahorse::Client::Response] DynamoDB response object.
        def query(params = {}, options = {})
          profile_dynamodb('query', params) { super }
        end

        # @param params [Hash] Request parameters for DynamoDB scan.
        # @param options [Hash] Additional AWS SDK options for the call.
        # @return [Seahorse::Client::Response] DynamoDB response object.
        def scan(params = {}, options = {})
          profile_dynamodb('scan', params) { super }
        end

        # @param params [Hash] Request parameters for DynamoDB batch_get_item.
        # @param options [Hash] Additional AWS SDK options for the call.
        # @return [Seahorse::Client::Response] DynamoDB response object.
        def batch_get_item(params = {}, options = {})
          profile_dynamodb('batch_get_item', params) { super }
        end

        # @param params [Hash] Request parameters for DynamoDB batch_write_item.
        # @param options [Hash] Additional AWS SDK options for the call.
        # @return [Seahorse::Client::Response] DynamoDB response object.
        def batch_write_item(params = {}, options = {})
          profile_dynamodb('batch_write_item', params) { super }
        end

        private

        # Profiles a DynamoDB operation so it shows up in Rack::MiniProfiler's SQL panel.
        #
        # - Gates off active MiniProfiler request (no overhead when disabled).
        # - Measures elapsed time using a monotonic clock.
        # - Records the "SQL" via Rack::MiniProfiler.record_sql (this is what populates SQL Summary).
        # - Also wraps in a `step` so entries appear in the timeline.
        #
        # @param operation [String] DynamoDB operation name (e.g., 'query', 'get_item').
        # @param params [Hash] The exact params hash sent to AWS SDK (symbols or strings).
        # @yieldreturn [Seahorse::Client::Response] The AWS SDK response.
        # @return [Seahorse::Client::Response] The AWS SDK response (unchanged).
        def profile_dynamodb(operation, params = {}, &blk)
          return yield unless mp_should_measure?

          base_label = build_ddb_label(operation, params)

          # Wrap in a step for the timeline view AND record SQL timing
          ::Rack::MiniProfiler.step("SQL: #{base_label}") do
            start_mono = Process.clock_gettime(Process::CLOCK_MONOTONIC)
            response = blk.call
            elapsed_ms = mp_elapsed_ms(start_mono)

            # Append response size to the recorded SQL line
            size = ddb_response_size_bytes(response)
            sql_text = "#{base_label} [#{size} bytes]"

            # Record as SQL timing (shows in SQL Summary)
            begin
              ::Rack::MiniProfiler.record_sql(sql_text, elapsed_ms, nil)
            rescue => e
              Rails.logger.debug { "DynamoDB profiling record_sql failed: #{e.class}: #{e.message}" }
            end

            response
          end
        rescue => e
          # Preserve original exceptions from the AWS SDK.
          raise e
        end

        # True if MiniProfiler is currently measuring this request.
        #
        # @return [Boolean]
        def mp_should_measure?
          c = ::Rack::MiniProfiler.current
          !!(c && c.measure && c.current_timer)
        end

        # Converts a monotonic start time to elapsed milliseconds.
        #
        # @param start_mono [Float] Monotonic clock reading captured before the operation.
        # @return [Float] Elapsed time in milliseconds with 0.1ms precision.
        def mp_elapsed_ms(start_mono)
          ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_mono) * 1000.0).round(1)
        end

        # Builds a concise, informative label for the operation.
        # Includes operation, table name (if provided), key info for item ops,
        # and Query/Scan expressions with placeholder substitution.
        #
        # @param operation [String]
        # @param params [Hash]
        # @return [String]
        def build_ddb_label(operation, params)
          # Support symbol and string keys (SDK accepts both).
          label = "#{operation}"

          case operation
          when 'query', 'scan'
            names = params[:expression_attribute_names] || params['ExpressionAttributeNames']
            values = params[:expression_attribute_values] || params['ExpressionAttributeValues']
            key_cond = params[:key_condition_expression] || params['KeyConditionExpression']
            filter = params[:filter_expression] || params['FilterExpression']

            parts = []
            parts << "#{substitute_expressions(key_cond, names, values)}" if key_cond
            parts << "#{substitute_expressions(filter, names, values)}" if filter
            label += " (#{parts.join('; ')})" if parts.any?

            idx = params[:index_name] || params['IndexName']
            label += " [idx: #{idx}]" if idx
          when 'get_item', 'update_item', 'delete_item'
            key = params[:key] || params['Key']
            if key && key.respond_to?(:map)
              key_info = key.map do |k, v|
                raw = extract_attr_value(v)
                "#{k}=#{truncate_value(raw)}"
              end.join(', ')
              label += " (#{key_info})" unless key_info.empty?
            end
          when 'batch_get_item'
            reqs = params[:request_items] || params['RequestItems']
            if reqs && reqs.respond_to?(:values)
              all_keys = []
              reqs.each do |_table, config|
                keys = Array((config[:keys] || config['Keys']))
                keys.each do |key|
                  # Extract the hash key value (usually the first key)
                  key_str = key.map { |k, v| "#{k}=#{truncate_value(extract_attr_value(v))}" }.join(',')
                  all_keys << key_str
                end
              end

              total = all_keys.size
              if total <= 10 && total.positive?
                # Show actual keys if 10 or fewer
                label += " [#{all_keys.join(', ')}]"
              elsif total.positive?
                # Show count if more than 10
                label += " [#{total} keys]"
              end
            end
          when 'batch_write_item'
            reqs = params[:request_items] || params['RequestItems']
            if reqs && reqs.respond_to?(:values)
              all_ops = []
              reqs.each do |_table, operations|
                Array(operations).each do |op|
                  if op[:put_request] || op['PutRequest']
                    item = (op[:put_request] || op['PutRequest'])[:item] || (op[:put_request] || op['PutRequest'])['Item']
                    if item && (key = item[:hash_key] || item['hash_key'])
                      all_ops << "put:#{truncate_value(extract_attr_value(key))}"
                    else
                      all_ops << 'put'
                    end
                  elsif op[:delete_request] || op['DeleteRequest']
                    key = (op[:delete_request] || op['DeleteRequest'])[:key] || (op[:delete_request] || op['DeleteRequest'])['Key']
                    if key && (hash_key = key[:hash_key] || key['hash_key'])
                      all_ops << "del:#{truncate_value(extract_attr_value(hash_key))}"
                    else
                      all_ops << 'del'
                    end
                  end
                end
              end

              total = all_ops.size
              if total <= 10 && total.positive?
                # Show actual operations if 10 or fewer
                label += " [#{all_ops.join(', ')}]"
              elsif total.positive?
                # Show count if more than 10
                label += " [#{total} operations]"
              end
            end
          end

          label
        rescue => e
          Rails.logger.debug { "DynamoDB profiling label build failed: #{e.class}: #{e.message}" }
          "DynamoDB.#{operation}"
        end

        # Replaces ExpressionAttributeNames and ExpressionAttributeValues in an expression
        # so the profiler shows meaningful text.
        #
        # @param expr [String, nil]
        # @param names [Hash, nil]
        # @param values [Hash, nil]
        # @return [String]
        def substitute_expressions(expr, names, values)
          return '' unless expr
          s = expr.to_s.dup

          if names.respond_to?(:each)
            names.each { |placeholder, real| s.gsub!(placeholder.to_s, real.to_s) }
          end

          if values.respond_to?(:each)
            values.each do |placeholder, typed|
              raw = extract_attr_value(typed)
              s.gsub!(placeholder.to_s, raw.inspect)
            end
          end

          s
        rescue => e
          Rails.logger.debug { "DynamoDB profiling expression substitution failed: #{e.class}: #{e.message}" }
          expr.to_s
        end

        # Extracts the raw value from a DynamoDB typed attribute (e.g., {"S"=>"abc"}) or returns the input.
        #
        # @param val [Object]
        # @return [Object]
        def extract_attr_value(val)
          if val.is_a?(Hash) && val.size == 1
            val.values.first
          else
            val
          end
        end

        # Computes approximate response size in bytes by JSON-encoding response.to_h.
        #
        # @param response [Object]
        # @return [Integer]
        def ddb_response_size_bytes(response)
          h = response.respond_to?(:to_h) ? response.to_h : {}
          JSON.generate(h).bytesize
        rescue => e
          Rails.logger.debug { "DynamoDB profiling size calc failed: #{e.class}: #{e.message}" }
          0
        end

        # Allocation-light truncation that avoids ActiveSupport.
        #
        # @param val [Object]
        # @param max [Integer]
        # @return [String]
        def truncate_value(val, max = 60)
          s = val.to_s
          return s if s.length <= max
          s[0, max - 3] + '...'
        end
      end
    end

    unless ::Aws::DynamoDB::Client.ancestors.include?(RackMiniProfilerDynamoDB::Instrumentation)
      ::Aws::DynamoDB::Client.prepend(RackMiniProfilerDynamoDB::Instrumentation)
    end

    Rails.logger.info 'Rack::MiniProfiler DynamoDB profiling enabled'
  end
end