# Wrap methods to keep track of ActiveRecord::Relation method calls and query
# executions.

module Immunio
  module RelationHooks
    extend ActiveSupport::Concern

    # ActiveRecord will "spawn", or clone, relations under many circumstances.
    # For example, #where will first spawn the relation, then add the where
    # clause to the new relation, then return the new relation. This is so
    # you can do something like:
    #
    # prods = Product.where(cost: 10)
    #
    # # Return blue products costing $10
    # prods_blue = prods.where(color: 'blue').to_a
    #
    # # Return count of all products costing $10
    # prods_count = prods.count
    #
    # We must track the prods relation and the relations returned by the color
    # where call separately.
    def clone_with_immunio
      cloned = clone_without_immunio

      Request.time "plugin", "Immunio::RelationTracking" do
        QueryTracker.instance.spawn_relation self, cloned
      end

      cloned
    end

    included do
      alias_method :clone_without_immunio, :clone
      alias_method :clone, :clone_with_immunio
    end
  end

  module SpawnHooks
    extend ActiveSupport::Concern
    # Sometimes ActiveRecord creates a new relation with a new condition and
    # merges it into an existing relation. I'm not sure why, and I'm not going
    # to ask. Just copy the context data from the other relation into this one.
    def merge_with_immunio(other)
      return merge_without_immunio(other) unless other && !other.is_a?(Array)

      # Rails 4 added the ability to call merge with a proc, like:
      #
      # relation.merge(-> { where(:foo) })
      #
      # We don't need to do anything here. If the proc calls relation methods,
      # they will be called on the right relation and everything will be good.
      if !other.is_a?(ActiveRecord::Relation) && other.respond_to?(:to_proc)
        return merge_without_immunio(other)
      end

      # Rails 4 added the ability to merge in a hash of conditions, like:
      #
      # Developer.merge(where: 'projects IS NOT NULL')
      #
      # Use the internal HashMerger to build a real relation from the hash
      # before merging values into the spawned relation.
      if other.is_a?(Hash)
        # This shouldn't happen, but let's be safe.
        unless defined? ActiveRecord::Relation::HashMerger
          return merge_without_immunio(other)
        end

        other = ActiveRecord::Relation::HashMerger.new(self, other).other
      end

      spawned = merge_without_immunio(other)
      Request.time "plugin", "Immunio::RelationTracking" do
        QueryTracker.instance.merge_relations spawned, other
      end
      spawned
    end

    included do
      alias_method :merge_without_immunio, :merge
      alias_method :merge, :merge_with_immunio
    end
  end

  module QueryingHooks
    extend ActiveSupport::Concern

    class << self
      attr_accessor :methods_to_wrap
    end

    self.methods_to_wrap = []

    # Wrap ActiveRecord::Relation methods to add API calls to the context data
    # for a query. Some additional methods are wrapped because they execute
    # queries. In those cases, we must associate the connection to the relation
    # so we can grab the context info when we see the query being executed.
    def self.wrap_method(method)
      method_with_immunio = "#{method}_with_immunio".to_sym
      method_without_immunio = "#{method}_without_immunio".to_sym

      define_method method_with_immunio do |*args, &block|
        Request.time "plugin", "Immunio::RelationTracking" do
          # Construct the context data. The name method returns the name of the
          # class of self, which is the name of the model. This allows us to
          # differentiate between scopes from the same lines of code but on
          # different models.
          #
          # In Ruby 2 and up, we can ask for specific frames and the bits we
          # need from them. This doesn't impact performance much. But in earlier
          # versions of Ruby we would need to gather the entire stack trace and
          # parse the strings for the frames we need. As that is too much of a
          # performance hit, we don't include caller information in the
          # additional context data.
          if defined? caller_locations
            stack = caller_locations(4..5)

            # If the method was called on a relation directly, the real caller
            # is four frames up. But if the method was called on the model
            # class, like `User.where`, then there is an extra frame for
            # delegating the method call to the result of the `all` method. We
            # look for the telltale sign of the caller being in the querying.rb
            # file of ActiveRecord and with the same name as the method.
            if stack[0].path.end_with?('querying.rb') && stack[0].label == method.to_s
              frame = stack[1]
            else
              frame = stack[0]
            end

            caller_method = frame.label
            caller_line = frame.lineno

            data = "Relation for #{name}, method called: #{method}, caller: #{caller_method}:#{caller_line}"
          else
            data = "Relation for #{name}, method called: #{method}"
          end

          # In Rails 3, #where and #having clone the relation and return it with
          # the conditions added, but use the original relation to actually
          # build the conditions. We have to add a hack to push the last cloned
          # relation onto the connection's relation stack so parameters are
          # associated with the cloned relation and not the original relation.
          relation = if method == :build_where && Rails::VERSION::MAJOR == 3
            QueryTracker.instance.last_spawned_relation connection
          else
            self
          end

          # Push the current relation onto the stack of relations for the
          # connection. The top relation on the stack at query execution time is
          # used for contextual data.
          QueryTracker.instance.push_relation relation

          begin
            # Call original method
            ret = Request.pause "plugin", "Immunio::RelationTracking" do
              send method_without_immunio, *args, &block
            end

            # These two methods create clones of the original without actually
            # calling #clone. Copy relation data over in this special case.
            if method == :except || method == :only
              QueryTracker.instance.spawn_relation self, ret
            end

            # If ret is a relation, such as when #where is called, we need to
            # add data about the API call to the relation. Otherwise, the API
            # call is returning actual data from a query execution, like #count.
            # We don't need to add context data about calls that execute queries
            # because the stack trace at query execution time will include this
            # call.
            if ret.is_a? ActiveRecord::Relation
              QueryTracker.instance.add_relation_data ret, data
            end
          ensure
            QueryTracker.instance.pop_relation relation
          end

          ret
        end
      end
    end

    if defined? ActiveRecord::Querying
      # Public ActiveRecord::Relation API methods we want to add context for.
      # Examples: #where, #group, #order, #joins
      self.methods_to_wrap = ActiveRecord::Querying.public_instance_methods

      # Additional methods that execute queries. We must push the
      # ActiveRecord::Relation onto the stack of relations for the connection so
      # we can grab its contextual data when the query is executed.
      self.methods_to_wrap += [
        :insert,
        :update_record,
        :_update_record,
        :exec_queries,
        :find_with_associations,
        :limited_ids_for,
        :execute_simple_calculation,
        :execute_grouped_calculation
      ]

      # Additional methods that sanitize SQL fragments, which we wrap to gather
      # parameters. We wrap these methods to ensure the proper relation is at
      # the top of the connection relation stack so we save parameters to the
      # correct relation.
      self.methods_to_wrap += [
        :having!,
        :where!,
        :build_where
      ]
    end

    self.methods_to_wrap.each do |method|
      wrap_method method
    end

    included do
      methods = Immunio::QueryingHooks.methods_to_wrap
      methods.keep_if do |method|
        next unless method_defined?(method) || private_method_defined?(method)

        method_with_immunio = "#{method}_with_immunio".to_sym
        method_without_immunio = "#{method}_without_immunio".to_sym

        alias_method method_without_immunio, method
        alias_method method, method_with_immunio

        true
      end

      Immunio.logger.debug {"Wrapped ActiveRecord::Relation methods for context: #{methods.join ", "}"}
    end
  end

  # Some simple queries, specifically #find and #find_by with simple conditions
  # may be saved in a statement cache so the statement doesn't need to be built
  # from an AST every time. When a query is executed from the cache, the cached
  # query statement is executed without going through any of the wrapped
  # ActiveRecord::Relation methods. We wrap the execution to ensure we push the
  # original relation for the cached statement onto the connection relation
  # stack.
  #
  # Note: Because cached statements won't build up the query from an AST, we
  # won't have AST context data. That should be ok because only the simplest
  # queries are cacheable, queries that won't vary in structure even for the
  # same stack trace and relation API calls.
  module StatementCacheHooks
    extend ActiveSupport::Concern

    module ClassMethods
      # When a statement is cached, store the final relation used to generate
      # the statement in a special immunio instance variable.
      def create_with_immunio(*args, &block)
        # The relation is returned by the block passed into the create method.
        # This is more fragile than most of our method wrappings as we depend on
        # *how* the method works, but there's no way around it.
        relation = block.call ActiveRecord::StatementCache::Params.new

        ret = create_without_immunio(*args) { relation }
        ret.instance_variable_set :@__immunio_relation, relation
        ret
      end
    end

    # When a statement is executed, push the original relation onto the
    # connection relation stack.
    def execute_with_immunio(*args)
      Request.time "plugin", "Immunio::RelationTracking" do
        QueryTracker.instance.push_relation @__immunio_relation
      end

      begin
        execute_without_immunio(*args)
      ensure
        Request.time "plugin", "Immunio::RelationTracking" do
          QueryTracker.instance.pop_relation @__immunio_relation
        end
      end
    end

    included do
      # StatementCache is available since 4.0, but not used by ActiveRecord
      # itself until 4.2. The signature changed in 4.2, so we wrap the 4.2
      # version and don't bother with previous versions. If any apps use
      # StatementCache directly on Rails 4.0 or 4.1, we will not propagate
      # relational context data properly.
      if ActiveRecord::StatementCache.singleton_class.method_defined? :create
        singleton_class.send :alias_method, :create_without_immunio, :create
        singleton_class.send :alias_method, :create, :create_with_immunio

        alias_method :execute_without_immunio, :execute
        alias_method :execute, :execute_with_immunio
      end
    end
  end

  module HasManyThroughAssociationHooks
    extend ActiveSupport::Concern

  private
    # One off, non-ActiveRecord::Relation method that under one condition
    # executes a query in Rails 4.1 and up. Unfortunately, wrapping won't work
    # as the relation used to generate the query is a temporary relation that is
    # created in the method. The easiest way to deal with it, though very hacky,
    # is to copy the method from upstream Rails and push the temporary relation
    # onto the stack for the connection right before the query is executed.
    def delete_records_with_immunio(records, method, *args)
      scope = through_association.scope

      unless method == :destroy && !scope.klass.primary_key
        return delete_records_without_immunio(records, method, *args)
      end

      # From here on down, copied from upstream Rails 4.2. Verified to work on
      # Rails 4.1, too.
      ensure_not_nested

      scope.where! construct_join_attributes(*records)

      scope.each do |record|
        record.run_callbacks :destroy
      end

      arel = scope.arel

      stmt = Arel::DeleteManager.new arel.engine
      stmt.from scope.klass.arel_table
      stmt.wheres = arel.constraints

      Request.time "plugin", "Immunio::RelationTracking" do
        QueryTracker.instance.push_relation scope
      end

      begin
        count = scope.klass.connection.delete(stmt, 'SQL', scope.bind_values)
      ensure
        Request.time "plugin", "Immunio::RelationTracking" do
          QueryTracker.instance.pop_relation scope
        end
      end

      delete_through_records(records)

      if source_reflection.options[:counter_cache] && method != :destroy
        counter = source_reflection.counter_cache_column
        klass.decrement_counter counter, records.map(&:id)
      end

      if through_reflection.collection? && update_through_counter?(method)
        update_counter(-count, through_reflection)
      end

      update_counter(-count)
    end

    included do
      if Rails::VERSION::MAJOR >= 4 && Rails::VERSION::MINOR >= 1
        alias_method :delete_records_without_immunio, :delete_records
        alias_method :delete_records, :delete_records_with_immunio
      end
    end
  end
end

ActiveRecord::Relation.send :include, Immunio::RelationHooks if defined? ActiveRecord::Relation
ActiveRecord::Relation.send :include, Immunio::SpawnHooks if defined? ActiveRecord::SpawnMethods
ActiveRecord::Relation.send :include, Immunio::QueryingHooks if defined? ActiveRecord::Querying
ActiveRecord::StatementCache.send :include, Immunio::StatementCacheHooks if defined? ActiveRecord::StatementCache
ActiveRecord::Associations::HasManyThroughAssociation.send :include, Immunio::HasManyThroughAssociationHooks if defined? ActiveRecord::Associations::HasManyThroughAssociation