nativepluck.rb 4.29 KB
Newer Older
Al Chou's avatar
Al Chou committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14
# Copyright 2018 Ohad Dahan, Al Chou
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  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.

Ohad Dahan's avatar
Ohad Dahan committed
15 16 17
require "nativepluck/version"

module Nativepluck
18
  @override_pluck = @nativepluck_init = false
Ohad Dahan's avatar
Ohad Dahan committed
19

Ohad Dahan's avatar
Ohad Dahan committed
20 21 22
  class << self
    attr_accessor :nativepluck_type_map_for_results, :nativepluck_type_map_for_queries
    attr_accessor :original_type_map_for_results, :original_type_map_for_queries
23
    attr_accessor :nativepluck_init, :override_pluck
24 25
  end

Ohad Dahan's avatar
Ohad Dahan committed
26

Ohad Dahan's avatar
Ohad Dahan committed
27 28 29

  def self.init_nativepluck
    return if @nativepluck_init
Ohad Dahan's avatar
Ohad Dahan committed
30 31 32 33
    @nativepluck_type_map_for_results = PG::BasicTypeMapForResults.new ActiveRecord::Base.connection.raw_connection
    @nativepluck_type_map_for_queries = PG::BasicTypeMapForQueries.new ActiveRecord::Base.connection.raw_connection
    @original_type_map_for_results    = ActiveRecord::Base.connection.raw_connection.type_map_for_results
    @original_type_map_for_queries    = ActiveRecord::Base.connection.raw_connection.type_map_for_queries
Ohad Dahan's avatar
Ohad Dahan committed
34
    @nativepluck_init = true
Ohad Dahan's avatar
Ohad Dahan committed
35 36 37
  end

  def self.set_pg_native_casters
Ohad Dahan's avatar
Ohad Dahan committed
38
    init_nativepluck
Ohad Dahan's avatar
Ohad Dahan committed
39 40 41 42 43
    ActiveRecord::Base.connection.raw_connection.type_map_for_results = Nativepluck.nativepluck_type_map_for_results
    ActiveRecord::Base.connection.raw_connection.type_map_for_queries = Nativepluck.nativepluck_type_map_for_queries
  end

  def self.return_original_casters
Ohad Dahan's avatar
Ohad Dahan committed
44
    init_nativepluck
Ohad Dahan's avatar
Ohad Dahan committed
45 46
    ActiveRecord::Base.connection.raw_connection.type_map_for_results = Nativepluck.original_type_map_for_results
    ActiveRecord::Base.connection.raw_connection.type_map_for_queries = Nativepluck.original_type_map_for_queries
47 48
  end

49
  def self.nativepluck(input)
50
    out = []
Ohad Dahan's avatar
Ohad Dahan committed
51 52 53 54 55 56 57 58 59 60 61 62
    begin
      sql = input.respond_to?(:to_sql) ? input.to_sql : input
      Nativepluck.set_pg_native_casters
      results = ActiveRecord::Base.connection.raw_connection.async_exec(sql)
      results.nfields == 1 ? out = results.column_values(0) : out = results.values
    ensure
      results.try(:clear)
      Nativepluck.return_original_casters
    end
    return out
  end

63
  module InstanceMethods
64
    def nativepluck(*column_names)
65 66 67
      # Extracted (before modifications) from:
      # ruby-2.5.0/gems/activerecord-5.2.1/lib/active_record/relation/calculations.rb
      if loaded? && (column_names.map(&:to_s) - @klass.attribute_names - @klass.attribute_aliases.keys).empty?
68
        return records.nativepluck(*column_names)
69 70 71 72
      end

      if has_include?(column_names.first)
        relation = apply_join_dependency
73
        relation.nativepluck(*column_names)
74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89
      else
        enforce_raw_sql_whitelist(column_names)
        relation = spawn
        relation.select_values = column_names.map { |cn|
          @klass.has_attribute?(cn) || @klass.attribute_alias?(cn) ? arel_attribute(cn) : cn
        }
        Nativepluck.nativepluck(relation)
      end
    end
  end

  module ::ActiveRecord
    module Calculations
      include Nativepluck::InstanceMethods
    end
    module Querying
90 91 92 93 94 95
      delegate :nativepluck, to: :all
    end
  end

  def self.set_override_pluck(selection)
    raise ArgumentError.new("#{__method__}:: Input should be true/false and not #{selection.class}") unless (selection.is_a?(TrueClass) || selection.is_a?(FalseClass))
96

97 98
    if selection
      ::ActiveRecord::Calculations.class_eval do
99
        alias_method :orig_pluck, :pluck
100 101 102
        alias_method :pluck, :nativepluck
      end
      ::ActiveRecord::Querying.class_eval do
103
        alias_method :orig_pluck, :pluck
104 105 106
        alias_method :pluck, :nativepluck
      end
    else
107 108 109
      unless Nativepluck.override_pluck
        raise ArgumentError.new("#{__method__}:: Cannot turn off pluck overriding since it wasn't set yet")
      end
110 111 112 113 114 115
      ::ActiveRecord::Calculations.class_eval do
        alias_method :pluck, :orig_pluck
      end
      ::ActiveRecord::Querying.class_eval do
        alias_method :pluck, :orig_pluck
      end
116
    end
117
    Nativepluck.override_pluck = selection
118
  end
Ohad Dahan's avatar
Ohad Dahan committed
119
end