nativepluck.rb 2.15 KB
Newer Older
Ohad Dahan's avatar
Ohad Dahan committed
1 2 3
require "nativepluck/version"

module Nativepluck
Ohad Dahan's avatar
Ohad Dahan committed
4 5 6 7 8 9
  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
  end


10 11
  def self.included(klass)
    klass.extend(ClassMethods)
Ohad Dahan's avatar
Ohad Dahan committed
12 13 14 15 16 17 18 19 20 21 22 23 24 25
    @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
  end

  def self.set_pg_native_casters
    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
    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
26 27 28 29 30
  end

  module ClassMethods
    def nativepluck(*columns, **opts)
      raise ArgumentError.new('No columns to pluck were provided') if columns.size == 0
31 32
      begin
        Nativepluck.set_pg_native_casters
33 34 35 36 37 38 39 40 41 42
        sql = "
        SELECT #{columns.join(',')}
        FROM #{self.table_name}
        #{"GROUP BY #{opts[:group].zip.map { |i| i.join(' ') }.join(',')}" if opts[:group]}
        #{"ORDER BY #{opts[:order].zip.map { |i| i.join(' ') }.join(',')}" if opts[:order]}
        #{"LIMIT #{opts[:limit]}"   if opts[:limit]}
        #{"OFFSET #{opts[:offset]}" if opts[:offset]}
        "
        results = ActiveRecord::Base.connection.raw_connection.async_exec(sql)
        puts "#{__method__} sql = #{sql}" if opts[:verbose]
43 44 45 46
        columns.size == 1 ? out = results.column_values(0) : out = results.values
      ensure
        Nativepluck.return_original_casters
      end
47
      return out
48 49
    end
  end
Ohad Dahan's avatar
Ohad Dahan committed
50
end