# encoding: utf-8

require 'one_apm/agent/datastores/metric_helper'

module OneApm
  module Agent
    module Datastores

      # Add Datastore tracing to a method. This properly generates the metrics
      # for OneApm's Datastore features. It does not capture the actual
      # query content into Transaction Traces. Use wrap if you want to provide
      # that functionality.
      #
      # @param [Class] clazz the class to instrument
      #
      # @param [String, Symbol] method_name the name of instance method to
      # instrument
      #
      # @param [String] product name of your datastore for use in metric naming, e.g. "Redis"
      #
      # @param [optional,String] operation the name of operation if different
      # than the instrumented method name
      #
      # @api public
      #
      def self.trace(clazz, method_name, product, operation = method_name)
        clazz.class_eval do
          method_name_without_oneapm = "#{method_name}_without_oneapm"

          if OneApm::Helper.instance_methods_include?(clazz, method_name) &&
             !OneApm::Helper.instance_methods_include?(clazz, method_name_without_oneapm)

            visibility = OneApm::Helper.instance_method_visibility(clazz, method_name)

            alias_method method_name_without_oneapm, method_name

            define_method(method_name) do |*args, &blk|
              metrics = MetricHelper.metrics_for(product, operation)
              OneApm::Support::MethodTracer.trace_execution_scoped(metrics) do
                send(method_name_without_oneapm, *args, &blk)
              end
            end

            send visibility, method_name
            send visibility, method_name_without_oneapm
          end
        end
      end

      # Wrap a call to a datastore and record OneApm Datastore metrics. This
      # method can be used when a collection (i.e. table or model name) is
      # known at runtime to be included in the metric naming. It is intended
      # for situations that the simpler OneApm::Agent::Datastores.trace can't
      # properly handle.
      #
      # To use this, wrap the datastore operation in the block passed to wrap.
      #
      #   OneApm::Agent::Datastores.wrap("FauxDB", "find", "items") do
      #     FauxDB.find(query)
      #   end
      #
      # @param [String] product the datastore name for use in metric naming,
      # e.g. "FauxDB"
      #
      # @param [String,Symbol] operation the name of operation (e.g. "select"),
      # often named after the method that's being instrumented.
      #
      # @param [optional, String] collection the collection name for use in
      # statement-level metrics (i.e. table or model name)
      #
      # @param [Proc,#call] callback proc or other callable to invoke after
      # running the datastore block. Receives three arguments: result of the
      # yield, the most specific (scoped) metric name, and elapsed time of the
      # call. An example use is attaching SQL to Transaction Traces at the end
      # of a wrapped datastore call.
      #
      #   callback = Proc.new do |result, metrics, elapsed|
      #     OneApm::Agent::Datastores.notice_sql(query, metrics, elapsed)
      #   end
      #
      #   OneApm::Agent::Datastores.wrap("FauxDB", "find", "items", callback) do
      #     FauxDB.find(query)
      #   end
      #
      # **NOTE: THERE ARE SECURITY CONCERNS WHEN CAPTURING SQL!**
      # OneApm's Transaction Tracing and Slow SQL features will
      # attempt to apply obfuscation to the passed queries, but it is possible
      # for a query format to be unsupported and result in exposing user
      # information embedded within captured queries.
      #
      # @api public
      #
      def self.wrap(product, operation, collection = nil, callback = nil)
        return yield unless operation

        metrics = MetricHelper.metrics_for(product, operation, collection)
        scoped_metric = metrics.first
        OneApm::Support::MethodTracer.trace_execution_scoped(metrics) do
          t0 = Time.now
          begin
            result = yield
          ensure
            if callback
              elapsed_time = (Time.now - t0).to_f
              callback.call(result, scoped_metric, elapsed_time)
            end
          end
        end
      end

      # Wrapper for simplifying attaching SQL queries during a transaction.
      #
      # If you are recording non-SQL data, please use the notice_statement
      # method instead.
      #
      #   OneApm::Agent::Datastores.notice_sql(query, metrics, elapsed)
      #
      # @param [String] query the SQL text to be captured. Note that depending
      # on user settings, this string will be run through obfuscation, but
      # some dialects of SQL (or non-SQL queries) are not guaranteed to be
      # properly obfuscated by these routines!
      #
      # @param [String] scoped_metric The most specific metric relating to this
      # query. Typically the result of
      # OneApm::Agent::Datastores::MetricHelper#metrics_for
      #
      # @param [Float] elapsed the elapsed time during query execution
      #
      # **NOTE: THERE ARE SECURITY CONCERNS WHEN CAPTURING SQL!**
      # OneApm's Transaction Tracing and Slow SQL features will
      # attempt to apply obfuscation to the passed queries, but it is possible
      # for a query format to be unsupported and result in exposing user
      # information embedded within captured queries.
      #
      def self.notice_sql(query, scoped_metric, elapsed)
        agent = OneApm::Agent.instance
        agent.transaction_sampler.notice_sql(query, nil, elapsed)
        agent.sql_sampler.notice_sql(query, scoped_metric, nil, elapsed)
        nil
      end

      # Wrapper for simplifying attaching non-SQL data statements to a
      # transaction. For instance, Mongo or CQL queries, Memcached or Redis
      # keys would all be appropriate data to attach as statements.
      #
      # Data passed to this method is NOT obfuscated by OneApm, so please
      # ensure that user information is obfuscated if the agent setting
      # `transaction_tracer.record_sql` is set to `obfuscated`
      #
      #   OneApm::Agent::Datastores.notice_statement("key", elapsed)
      #
      # @param [String] statement text of the statement to capture.
      #
      # @param [Float] elapsed the elapsed time during query execution
      #
      # **NOTE: THERE ARE SECURITY CONCERNS WHEN CAPTURING STATEMENTS!**
      # This method will properly ignore statements when the user has turned
      # off capturing queries, but it is not able to obfuscate arbitrary data!
      # To prevent exposing user information embedded in captured queries,
      # please ensure all data passed to this method is safe to transmit to
      # OneApm.
      #
      def self.notice_statement(statement, elapsed)
        # Settings may change eventually, but for now we follow the same
        # capture rules as SQL for non-SQL statements.
        return unless OneApm::Agent::Database.should_record_sql?

        agent = OneApm::Agent.instance
        agent.transaction_sampler.notice_nosql_statement(statement, elapsed)
        nil
      end

    end
  end
end