oneapm_rpm 1.1.0

data/lib/one_apm/agent/database.rb

@@ -0,0 +1,308 @@
+# encoding: utf-8
+
+require 'singleton'
+require 'one_apm/agent/database/obfuscation_helpers'
+require 'one_apm/agent/database/obfuscator'
+require 'one_apm/agent/database/postgres_explain_obfuscator'
+
+module OneApm
+  # columns for a mysql explain plan
+  MYSQL_EXPLAIN_COLUMNS = [
+                           "Id",
+                           "Select Type",
+                           "Table",
+                           "Type",
+                           "Possible Keys",
+                           "Key",
+                           "Key Length",
+                           "Ref",
+                           "Rows",
+                           "Extra"
+                          ].freeze
+
+  module Agent
+    module Database
+      MAX_QUERY_LENGTH = 16384
+
+      extend self
+
+      def capture_query(query)
+        Helper.correctly_encoded(truncate_query(query))
+      end
+
+      def truncate_query(query)
+        if query.length > (MAX_QUERY_LENGTH - 4)
+          query[0..MAX_QUERY_LENGTH - 4] + '...'
+        else
+          query
+        end
+      end
+
+      def obfuscate_sql(sql)
+        Obfuscator.instance.obfuscator.call(sql)
+      end
+
+      def set_sql_obfuscator(type, &block)
+        Obfuscator.instance.set_sql_obfuscator(type, &block)
+      end
+
+      def record_sql_method(config_section=:transaction_tracer)
+        case Agent.config["#{config_section}.record_sql".to_sym].to_s
+        when 'off'
+          :off
+        when 'none'
+          :off
+        when 'false'
+          :off
+        when 'raw'
+          :raw
+        else
+          :obfuscated
+        end
+      end
+
+      RECORD_FOR = [:raw, :obfuscated].freeze
+
+      def should_record_sql?(config_section=:transaction_tracer)
+        RECORD_FOR.include?(record_sql_method(config_section))
+      end
+
+      def should_collect_explain_plans?(config_section=:transaction_tracer)
+        should_record_sql?(config_section) &&
+          Agent.config["#{config_section}.explain_enabled".to_sym]
+      end
+
+      def get_connection(config, &connector)
+        ConnectionManager.instance.get_connection(config, &connector)
+      end
+
+      def close_connections
+        ConnectionManager.instance.close_connections
+      end
+
+      # This takes a connection config hash from ActiveRecord or Sequel and
+      # returns a string describing the associated database adapter
+      def adapter_from_config(config)
+        if config[:adapter]
+          return config[:adapter].to_s
+        elsif config[:uri] && config[:uri].to_s =~ /^jdbc:([^:]+):/
+          # This case is for Sequel with the jdbc-mysql, jdbc-postgres, or
+          # jdbc-sqlite3 gems.
+          return $1
+        end
+      end
+
+      # Perform this in the runtime environment of a managed
+      # application, to explain the sql statement executed within a
+      # segment of a transaction sample. Returns an array of
+      # explanations (which is an array rows consisting of an array of
+      # strings for each column returned by the the explain query)
+      # Note this happens only for statements whose execution time
+      # exceeds a threshold (e.g. 500ms) and only within the slowest
+      # transaction in a report period, selected for shipment to OneApm
+      def explain_sql(sql, connection_config, &explainer)
+        return nil unless sql && connection_config
+        statement = sql.split(";\n")[0] # only explain the first
+        explain_plan = explain_statement(statement, connection_config, &explainer)
+        return explain_plan || []
+      end
+
+      SUPPORTED_ADAPTERS_FOR_EXPLAIN = %w[postgres postgresql mysql2 mysql sqlite].freeze
+
+      def explain_statement(statement, config, &explainer)
+        return unless is_select?(statement)
+
+        if statement[-3,3] == '...'
+          OneApm::Agent.logger.debug('Unable to collect explain plan for truncated query.')
+          return
+        end
+
+        if parameterized?(statement)
+          OneApm::Agent.logger.debug('Unable to collect explain plan for parameterized query.')
+          return
+        end
+
+        adapter = adapter_from_config(config)
+        if !SUPPORTED_ADAPTERS_FOR_EXPLAIN.include?(adapter)
+          OneApm::Agent.logger.debug("Not collecting explain plan because an unknown connection adapter ('#{adapter}') was used.")
+          return
+        end
+
+        handle_exception_in_explain do
+          start = Time.now
+          plan = explainer.call(config, statement)
+          ::OneApm::Agent.record_metric("Supportability/Database/execute_explain_plan", Time.now - start)
+          return process_resultset(plan, adapter) if plan
+        end
+      end
+
+      def process_resultset(results, adapter)
+        case adapter.to_s
+        when 'postgres', 'postgresql'
+          process_explain_results_postgres(results)
+        when 'mysql2'
+          process_explain_results_mysql2(results)
+        when 'mysql'
+          process_explain_results_mysql(results)
+        when 'sqlite'
+          process_explain_results_sqlite(results)
+        end
+      end
+
+      QUERY_PLAN = 'QUERY PLAN'.freeze
+
+      def process_explain_results_postgres(results)
+        if results.is_a?(String)
+          query_plan_string = results
+        else
+          lines = []
+          results.each { |row| lines << row[QUERY_PLAN] }
+          query_plan_string = lines.join("\n")
+        end
+
+        unless record_sql_method == :raw
+          query_plan_string = OneApm::Agent::Database::PostgresExplainObfuscator.obfuscate(query_plan_string)
+        end
+        values = query_plan_string.split("\n").map { |line| [line] }
+
+        [[QUERY_PLAN], values]
+      end
+
+      # Sequel returns explain plans as just one big pre-formatted String
+      # In that case, we send a nil headers array, and the single string
+      # wrapped in an array for the values.
+      # Note that we don't use this method for Postgres explain plans, since
+      # they need to be passed through the explain plan obfuscator first.
+      def string_explain_plan_results(results)
+        [nil, [results]]
+      end
+
+      def process_explain_results_mysql(results)
+        return string_explain_plan_results(results) if results.is_a?(String)
+        headers = []
+        values  = []
+        if results.is_a?(Array)
+          # We're probably using the jdbc-mysql gem for JRuby, which will give
+          # us an array of hashes.
+          headers = results.first.keys
+          results.each do |row|
+            values << headers.map { |h| row[h] }
+          end
+        else
+          # We're probably using the native mysql driver gem, which will give us
+          # a Mysql::Result object that responds to each_hash
+          results.each_hash do |row|
+            headers = row.keys
+            values << headers.map { |h| row[h] }
+          end
+        end
+        [headers, values]
+      end
+
+      def process_explain_results_mysql2(results)
+        return string_explain_plan_results(results) if results.is_a?(String)
+        headers = results.fields
+        values  = []
+        results.each { |row| values << row }
+        [headers, values]
+      end
+
+      SQLITE_EXPLAIN_COLUMNS = %w[addr opcode p1 p2 p3 p4 p5 comment]
+
+      def process_explain_results_sqlite(results)
+        return string_explain_plan_results(results) if results.is_a?(String)
+        headers = SQLITE_EXPLAIN_COLUMNS
+        values  = []
+        results.each do |row|
+          values << headers.map { |h| row[h] }
+        end
+        [headers, values]
+      end
+
+      def handle_exception_in_explain
+        yield
+      rescue => e
+        begin
+          # guarantees no throw from explain_sql
+          ::OneApm::Agent.logger.error("Error getting query plan:", e)
+          nil
+        rescue
+          # double exception. throw up your hands
+          nil
+        end
+      end
+
+      KNOWN_OPERATIONS = [
+        'alter',
+        'select',
+        'update',
+        'delete',
+        'insert',
+        'create',
+        'show',
+        'set',
+        'exec',
+        'execute',
+        'call'
+      ]
+
+      SQL_COMMENT_REGEX = Regexp.new('/\*.*?\*/', Regexp::MULTILINE).freeze
+
+      def parse_operation_from_query(sql)
+        sql = sql.gsub(SQL_COMMENT_REGEX, '')
+        if sql =~ /(\w+)/
+          op = $1.downcase
+          return op if KNOWN_OPERATIONS.include?(op)
+        end
+      end
+
+      def is_select?(statement)
+        parse_operation_from_query(statement) == 'select'
+      end
+
+      def parameterized?(statement)
+        Obfuscator.instance.obfuscate_single_quote_literals(statement) =~ /\$\d+/
+      end
+
+      class ConnectionManager
+        include Singleton
+
+        # Returns a cached connection for a given ActiveRecord
+        # configuration - these are stored or reopened as needed, and if
+        # we cannot get one, we ignore it and move on without explaining
+        # the sql
+        def get_connection(config, &connector)
+          @connections ||= {}
+
+          connection = @connections[config]
+
+          return connection if connection
+
+          begin
+            @connections[config] = connector.call(config)
+          rescue => e
+            ::OneApm::Agent.logger.error("Caught exception trying to get connection to DB for explain.", e)
+            nil
+          end
+        end
+
+        # Closes all the connections in the internal connection cache
+        def close_connections
+          @connections ||= {}
+          @connections.values.each do |connection|
+            begin
+              connection.disconnect!
+            rescue
+            end
+          end
+
+          @connections = {}
+        end
+      end
+
+      class Statement < String
+        attr_accessor :adapter, :config, :explainer
+      end
+    end
+  end
+end

data/lib/one_apm/agent/database/active_record_helper.rb

@@ -0,0 +1,80 @@
+# encoding: utf-8
+module OneApm
+  module Agent
+    module Instrumentation
+      module ActiveRecordHelper
+        module_function
+
+        def metric_for_name(name)
+          return unless name && name.respond_to?(:split)
+          parts = name.split(' ')
+          if parts.size == 2
+            model = parts.first
+            operation = parts.last.downcase
+            case operation
+            when 'load', 'count', 'exists'
+              op_name = 'find'
+            when 'indexes', 'columns'
+              op_name = nil # fall back to DirectSQL
+            when 'destroy', 'find', 'save', 'create'
+              op_name = operation
+            when 'update'
+              op_name = 'save'
+            else
+              if model == 'Join'
+                op_name = operation
+              end
+            end
+            "ActiveRecord/#{model}/#{op_name}" if op_name
+          end
+        end
+
+        def metric_for_sql(sql) #THREAD_LOCAL_ACCESS
+          txn = OneApm::Transaction.tl_current
+          metric = txn && txn.database_metric_name
+          if metric.nil?
+            operation = OneApm::Agent::Database.parse_operation_from_query(sql)
+            if operation
+              # Could not determine the model/operation so use a fallback metric
+              metric = "Database/SQL/#{operation}"
+            else
+              metric = "Database/SQL/other"
+            end
+          end
+          metric
+        end
+
+        # Given a metric name such as "ActiveRecord/model/action" this
+        # returns an array of rollup metrics:
+        # [ "Datastore/all", "ActiveRecord/all", "ActiveRecord/action" ]
+        # If the metric name is in the form of "ActiveRecord/action"
+        # this returns merely: [ "Datastore/all", "ActiveRecord/all" ]
+        def rollup_metrics_for(metric)
+          metrics = ["Datastore/all"]
+
+          # If we're outside of a web transaction, don't record any rollup
+          # database metrics. This is to prevent metrics from background tasks
+          # from polluting the metrics used to drive overview graphs.
+          if OneApm::Transaction.recording_web_transaction?
+            metrics << "ActiveRecord/all"
+          else
+            metrics << "Datastore/allOther"
+          end
+          metrics << "ActiveRecord/#{$1}" if metric =~ /ActiveRecord\/[\w|\:]+\/(\w+)/
+
+          metrics
+        end
+
+        # Given a database adapter name and a database server host
+        # this returns a metric name in the form:
+        # "RemoteService/sql/adapter/host"
+        # Host defaults to "localhost".
+        def remote_service_metric(adapter, host)
+          host ||= 'localhost'
+          type = adapter.to_s.sub(/\d*/, '')
+          "RemoteService/sql/#{type}/#{host}"
+        end
+      end
+    end
+  end
+end

data/lib/one_apm/agent/database/obfuscation_helpers.rb

@@ -0,0 +1,76 @@
+# encoding: utf-8
+
+module OneApm
+  module Agent
+    module Database
+      module ObfuscationHelpers
+        # Note that the following two regexes are applied to a reversed version
+        # of the query. This is why the backslash escape sequences (\' and \")
+        # appear reversed within them.
+        #
+        # Note that some database adapters (notably, PostgreSQL with
+        # standard_conforming_strings on and MySQL with NO_BACKSLASH_ESCAPES on)
+        # do not apply special treatment to backslashes within quoted string
+        # literals. We don't have an easy way of determining whether the
+        # database connection from which a query was captured was operating in
+        # one of these modes, but the obfuscation is done in such a way that it
+        # should not matter.
+        #
+        # Reversing the query string before obfuscation allows us to get around
+        # the fact that a \' appearing within a string may or may not terminate
+        # the string, because we know that a string cannot *start* with a \'.
+        REVERSE_SINGLE_QUOTES_REGEX = /'(?:''|'\\|[^'])*'/
+        REVERSE_ANY_QUOTES_REGEX    = /'(?:''|'\\|[^'])*'|"(?:""|"\\|[^"])*"/
+
+        NUMERICS_REGEX = /\b\d+\b/
+
+        # We take a conservative, overly-aggressive approach to obfuscating
+        # comments, and drop everything from the query after encountering any
+        # character sequence that could be a comment initiator. We do this after
+        # removal of string literals to avoid accidentally over-obfuscating when
+        # a string literal contains a comment initiator.
+        SQL_COMMENT_REGEX = Regexp.new('(?:/\*|--|#).*', Regexp::MULTILINE).freeze
+
+        # We use these to check whether the query contains any quote characters
+        # after obfuscation. If so, that's a good indication that the original
+        # query was malformed, and so our obfuscation can't reliabily find
+        # literals. In such a case, we'll replace the entire query with a
+        # placeholder.
+        LITERAL_SINGLE_QUOTE = "'".freeze
+        LITERAL_DOUBLE_QUOTE = '"'.freeze
+
+        PLACEHOLDER = '?'.freeze
+
+        def obfuscate_single_quote_literals(sql)
+          obfuscated = sql.reverse
+          obfuscated.gsub!(REVERSE_SINGLE_QUOTES_REGEX, PLACEHOLDER)
+          obfuscated.reverse!
+          obfuscated
+        end
+
+        def obfuscate_quoted_literals(sql)
+          obfuscated = sql.reverse
+          obfuscated.gsub!(REVERSE_ANY_QUOTES_REGEX, PLACEHOLDER)
+          obfuscated.reverse!
+          obfuscated
+        end
+
+        def obfuscate_numeric_literals(sql)
+          sql.gsub(NUMERICS_REGEX, PLACEHOLDER)
+        end
+
+        def remove_comments(sql)
+          sql.gsub(SQL_COMMENT_REGEX, PLACEHOLDER)
+        end
+
+        def contains_single_quotes?(str)
+          str.include?(LITERAL_SINGLE_QUOTE)
+        end
+
+        def contains_quotes?(str)
+          str.include?(LITERAL_SINGLE_QUOTE) || str.include?(LITERAL_DOUBLE_QUOTE)
+        end
+      end
+    end
+  end
+end