postgresql_cursor 0.4.3 → 0.5.0

data/lib/postgresql_cursor/active_record/connection_adapters/postgresql_type_map.rb

@@ -0,0 +1,17 @@
+# lib/postgresql_cursor/active_record/connection_adapters/postgresql_type_map
+module PostgreSQLCursor
+  module ActiveRecord
+    module ConnectionAdapters
+      module PostgreSQLTypeMap
+        # Returns the private "type_map" needed for the cursor operation
+        def get_type_map # :nodoc:
+          if ::ActiveRecord::VERSION::MAJOR == 4 && ::ActiveRecord::VERSION::MINOR == 0
+            ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::OID::TYPE_MAP
+          else
+            type_map
+          end
+        end
+      end
+    end
+  end
+end

data/lib/postgresql_cursor/active_record/relation/cursor_iterators.rb

@@ -0,0 +1,64 @@
+# Defines extension to ActiveRecord/AREL to use this library
+module PostgreSQLCursor
+  module ActiveRecord
+    module Relation
+      module CursorIterators
+
+        # Public: Executes the query, returning each row as a hash
+        # to the given block.
+        #
+        # options     - Hash to control
+        #   fraction: 0.1..1.0    - The cursor_tuple_fraction (default 1.0)
+        #   block_size: 1..n      - The number of rows to fetch per db block fetch
+        #   while: value          - Exits loop when block does not return this value.
+        #   until: value          - Exits loop when block returns this value.
+        #
+        # Example:
+        #   Post.where(user_id:123).each_row { |hash| Post.process(hash) }
+        #   Post.each_row.map {|r| r["id"].to_i }
+        #
+        # Returns the number of rows yielded to the block
+        def each_row(options={}, &block)
+          options = {:connection => self.connection}.merge(options)
+          cursor  = PostgreSQLCursor::Cursor.new(to_sql, options)
+          return cursor.each_row(&block) if block_given?
+          cursor
+        end
+        alias :each_hash :each_row
+
+        # Public: Like each_row, but returns an instantiated model object to the block
+        #
+        # Paramaters: same as each_row
+        #
+        # Example:
+        #   Post.where(user_id:123).each_instance { |post| post.process }
+        #   Post.where(user_id:123).each_instance.map { |post| post.process }
+        #
+        # Returns the number of rows yielded to the block
+        def each_instance(options={}, &block)
+          options = {:connection => self.connection}.merge(options)
+          cursor = PostgreSQLCursor::Cursor.new(to_sql, options)
+          return cursor.each_instance(self, &block) if block_given?
+          cursor.iterate_type(self)
+        end
+
+        # Plucks the column names from the rows, and return them in an array
+        def pluck_rows(*cols)
+          options = cols.last.is_a?(Hash) ? cols.pop : {}
+          options[:connection] = self.connection
+          self.each_row(options).pluck(*cols)
+        end
+        alias :pluck_row :pluck_rows
+
+        # Plucks the column names from the instances, and return them in an array
+        def pluck_instances(*cols)
+          options = cols.last.is_a?(Hash) ? cols.pop : {}
+          options[:connection] = self.connection
+          self.each_instance(options).pluck(*cols)
+        end
+        alias :pluck_instance :pluck_instances
+
+      end
+    end
+  end
+end

data/lib/postgresql_cursor/active_record/sql_cursor.rb

@@ -0,0 +1,92 @@
+module PostgreSQLCursor
+  module ActiveRecord
+    module SqlCursor
+      # Public: Executes the query, returning each row as a hash
+      # to the given block.
+      #
+      # options     - Hash to control
+      #   fraction: 0.1..1.0    - The cursor_tuple_fraction (default 1.0)
+      #   block_size: 1..n      - The number of rows to fetch per db block fetch
+      #   while: value          - Exits loop when block does not return this value.
+      #   until: value          - Exits loop when block returns this value.
+      #
+      # Example:
+      #   Post.each_row { |hash| Post.process(hash) }
+      #
+      # Returns the number of rows yielded to the block
+      def each_row(options={}, &block)
+        options = {:connection => self.connection}.merge(options)
+        all.each_row(options, &block)
+      end
+      alias :each_hash :each_row
+
+      # Public: Like each_row, but returns an instantiated model object to the block
+      #
+      # Paramaters: same as each_row
+      #
+      # Example:
+      #   Post.each_instance { |post| post.process }
+      #
+      # Returns the number of rows yielded to the block
+      def each_instance(options={}, &block)
+        options = {:connection => self.connection}.merge(options)
+        all.each_instance(options, &block)
+      end
+
+      # Public: Returns each row as a hash to the given block
+
+      # sql         - Full SQL statement, variables interpolated
+      # options     - Hash to control
+      #   fraction: 0.1..1.0    - The cursor_tuple_fraction (default 1.0)
+      #   block_size: 1..n      - The number of rows to fetch per db block fetch
+      #   while: value          - Exits loop when block does not return this value.
+      #   until: value          - Exits loop when block returns this value.
+      #
+      # Example:
+      #   Post.each_row_by_sql("select * from posts") { |hash| Post.process(hash) }
+      #   Post.each_row_by_sql("select * from posts").count
+      #
+      # Returns the number of rows yielded to the block
+      def each_row_by_sql(sql, options={}, &block)
+        options = {:connection => self.connection}.merge(options)
+        cursor  = PostgreSQLCursor::Cursor.new(sql, options)
+        return cursor.each_row(&block) if block_given?
+        cursor
+      end
+      alias :each_hash_by_sql :each_row_by_sql
+
+      # Public: Returns each row as a model instance to the given block
+      # As this instantiates a model object, it is slower than each_row_by_sql
+      #
+      # Paramaters: see each_row_by_sql
+      #
+      # Example:
+      #   Post.each_instance_by_sql("select * from posts") { |post| post.process }
+      #   Post.each_instance_by_sql("select * from posts").count
+      #
+      # Returns the number of rows yielded to the block
+      def each_instance_by_sql(sql, options={}, &block)
+        options = {:connection => self.connection}.merge(options)
+        cursor  = PostgreSQLCursor::Cursor.new(sql, options)
+        return cursor.each_instance(self, &block) if block_given?
+        cursor.iterate_type(self)
+      end
+
+      # Returns and array of the given column names. Use if you need cursors and don't expect
+      # this to comsume too much memory. Values are strings. Like ActiveRecord's pluck.
+      def pluck_rows(*cols)
+        options = cols.last.is_a?(Hash) ? cols.pop : {}
+        all.each_row(options).pluck(*cols)
+      end
+      alias :pluck_row :pluck_rows
+
+      # Returns and array of the given column names. Use if you need cursors and don't expect
+      # this to comsume too much memory. Values are instance types. Like ActiveRecord's pluck.
+      def pluck_instances(*cols)
+        options = cols.last.is_a?(Hash) ? cols.pop : {}
+        all.each_instance(options).pluck(*cols)
+      end
+      alias :pluck_instance :pluck_instances
+    end
+  end
+end

data/lib/postgresql_cursor/cursor.rb

@@ -0,0 +1,199 @@
+################################################################################
+# PostgreSQLCursor: library class provides postgresql cursor for large result
+# set processing. Requires ActiveRecord, but can be adapted to other DBI/ORM libraries.
+# If you don't use AR, this assumes #connection and #instantiate methods are available.
+#
+# options     - Hash to control operation and loop breaks
+#   connection: instance  - ActiveRecord connection to use
+#   fraction: 0.1..1.0    - The cursor_tuple_fraction (default 1.0)
+#   block_size: 1..n      - The number of rows to fetch per db block fetch
+#   while: value          - Exits loop when block does not return this value.
+#   until: value          - Exits loop when block returns this value.
+#
+# Exmaples:
+#   PostgreSQLCursor::Cursor.new("select ...").each { |hash| ... }
+#   ActiveRecordModel.where(...).each_row { |hash| ... }
+#   ActiveRecordModel.each_row_by_sql("select ...") { |hash| ... }
+#   ActiveRecordModel.each_instance_by_sql("select ...") { |model| ... }
+#
+module PostgreSQLCursor
+  class Cursor
+    include Enumerable
+    attr_reader :sql, :options, :connection, :count, :result
+    @@cursor_seq = 0
+
+    # Public: Start a new PostgreSQL cursor query
+    # sql     - The SQL statement with interpolated values
+    # options - hash of processing controls
+    #   while: value    - Exits loop when block does not return this value.
+    #   until: value    - Exits loop when block returns this value.
+    #   fraction: 0.1..1.0    - The cursor_tuple_fraction (default 1.0)
+    #   block_size: 1..n      - The number of rows to fetch per db block fetch
+    #                           Defaults to 1000
+    #
+    # Examples
+    #
+    #   PostgreSQLCursor::Cursor.new("select ....")
+    #
+    # Returns the cursor object when called with new.
+    def initialize(sql, options={})
+      @sql        = sql
+      @options    = options
+      @connection = @options.fetch(:connection) { ::ActiveRecord::Base.connection }
+      @count      = 0
+      @iterate    = options[:instances] ? :each_instance : :each_row
+    end
+
+    # Specify the type to instantiate, or reset to return a Hash
+    def iterate_type(type=nil)
+      if type.nil? || type == Hash
+        @iterate = :each_row
+      else
+        @iterate = :each_instance
+        @type    = type
+      end
+      self
+    end
+
+    # Public: Yields each row of the result set to the passed block
+    #
+    # Yields the row to the block. The row is a hash with symbolized keys.
+    #   {colname: value, ....}
+    #
+    # Returns the count of rows processed
+    def each(&block)
+      if @iterate == :each_row
+        self.each_row(&block)
+      else
+        self.each_instance(@type, &block)
+      end
+    end
+
+    def each_row(&block)
+      self.each_tuple do |row|
+        row = row.symbolize_keys if @options[:symbolize_keys]
+        block.call(row)
+      end
+    end
+
+    def each_instance(klass=nil, &block)
+      klass ||= @type
+      self.each_tuple do |row|
+        if ::ActiveRecord::VERSION::MAJOR < 4
+          model = klass.send(:instantiate,row)
+        else
+          @column_types ||= column_types
+          model = klass.send(:instantiate, row, @column_types)
+        end
+        block.call(model)
+      end
+    end
+
+    # Returns an array of columns plucked from the result rows.
+    # Experimental function, as this could still use too much memory
+    # and negate the purpose of this libarary.
+    # Should this return a lazy enumerator instead?
+    def pluck(*cols)
+      options = cols.last.is_a?(Hash) ? cols.pop : {}
+      @options.merge!(options)
+      @options[:symbolize_keys] = true
+      self.iterate_type(options[:class]) if options[:class]
+      cols    = cols.map {|c| c.to_sym }
+      result  = []
+
+      self.each() do |row|
+        row = row.symbolize_keys if row.is_a?(Hash)
+        result << cols.map { |c| row[c] }
+      end
+
+      result.flatten! if cols.size == 1
+      result
+    end
+
+    def each_tuple(&block) #:nodoc:
+      has_do_until  = @options.has_key?(:until)
+      has_do_while  = @options.has_key?(:while)
+      @count        = 0
+      @column_types = nil
+      @connection.transaction do
+        begin
+          open
+          while (row = fetch) do
+            break if row.size==0
+            @count += 1
+            rc = block.call(row)
+            break if has_do_until && rc == @options[:until]
+            break if has_do_while && rc != @options[:while]
+          end
+        rescue Exception => e
+          raise e
+        ensure
+          close
+        end
+      end
+      @count
+    end
+
+    def cast_types(row)
+      row
+    end
+
+    def column_types
+      return nil if ::ActiveRecord::VERSION::MAJOR < 4
+      return @column_types if @column_types
+
+      types = {}
+      fields = @result.fields
+      fields.each_with_index do |fname, i|
+        ftype = @result.ftype i
+        fmod  = @result.fmod i
+        types[fname] = @connection.get_type_map.fetch(ftype, fmod) { |oid, mod|
+          warn "unknown OID: #{fname}(#{oid}) (#{sql})"
+          OID::Identity.new
+        }
+      end
+
+      @column_types = types
+    end
+
+    # Public: Opens (actually, "declares") the cursor. Call this before fetching
+    def open
+      set_cursor_tuple_fraction
+      @cursor = @@cursor_seq += 1
+      @result = @connection.execute("declare cursor_#{@cursor} cursor for #{@sql}")
+      @block = []
+    end
+
+    # Public: Returns the next row from the cursor, or empty hash if end of results
+    #
+    # Returns a row as a hash of {'colname'=>value,...}
+    def fetch
+      fetch_block if @block.size==0
+      @block.shift
+    end
+
+    # Private: Fetches the next block of rows into @block
+    def fetch_block(block_size=nil)
+      block_size ||= @block_size ||= @options.fetch(:block_size) { 1000 }
+      @result = @connection.execute("fetch #{block_size} from cursor_#{@cursor}")
+      @block = @result.collect {|row| row } # Make our own
+    end
+
+    # Public: Closes the cursor
+    def close
+      @connection.execute("close cursor_#{@cursor}")
+    end
+
+    # Private: Sets the PostgreSQL cursor_tuple_fraction value = 1.0 to assume all rows will be fetched
+    # This is a value between 0.1 and 1.0 (PostgreSQL defaults to 0.1, this library defaults to 1.0)
+    # used to determine the expected fraction (percent) of result rows returned the the caller.
+    # This value determines the access path by the query planner.
+    def set_cursor_tuple_fraction(frac=1.0)
+      @cursor_tuple_fraction ||= @options.fetch(:fraction) { 1.0 }
+      return @cursor_tuple_fraction if frac == @cursor_tuple_fraction
+      @cursor_tuple_fraction = frac
+      @result = @connection.execute("set cursor_tuple_fraction to  #{frac}")
+      frac
+    end
+  end
+end

data/lib/postgresql_cursor/version.rb

@@ -0,0 +1,3 @@
+module PostgresqlCursor
+  VERSION = "0.5.0"
+end

data/postgresql_cursor.gemspec

@@ -1,48 +1,27 @@
-# Generated by jeweler
-# DO NOT EDIT THIS FILE DIRECTLY
-# Instead, edit Jeweler::Tasks in Rakefile, and run 'rake gemspec'
-# -*- encoding: utf-8 -*-
-# stub: postgresql_cursor 0.4.3 ruby lib
+# coding: utf-8
+lib = File.expand_path('../lib', __FILE__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 'postgresql_cursor/version'
 
-Gem::Specification.new do |s|
-  s.name = "postgresql_cursor"
-  s.version = "0.4.3"
+Gem::Specification.new do |spec|
+  spec.name          = "postgresql_cursor"
+  spec.version       = PostgresqlCursor::VERSION
+  spec.authors       = ["Allen Fair"]
+  spec.email         = ["allen.fair@gmail.com"]
+  spec.summary       = "ActiveRecord PostgreSQL Adapter extension for using a cursor to return a large result set"
+  spec.description   = "PostgreSQL Cursor is an extension to the ActiveRecord PostgreSQLAdapter for very large result sets. It provides a cursor open/fetch/close interface to access data without loading all rows into memory, and instead loads the result rows in \"chunks\" (default of 1_000 rows), buffers them, and returns the rows one at a time."
+  spec.homepage      = "http://github.com/afair/postgresql_cursor"
+  spec.license       = "MIT"
 
-  s.required_rubygems_version = Gem::Requirement.new(">= 0") if s.respond_to? :required_rubygems_version=
-  s.require_paths = ["lib"]
-  s.authors = ["Allen Fair"]
-  s.date = "2014-06-06"
-  s.description = "PostgreSQL Cursor is an extension to the ActiveRecord PostgreSQLAdapter for very large result sets. It provides a cursor open/fetch/close interface to access data without loading all rows into memory, and instead loads the result rows in \"chunks\" (default of 10_000 rows), buffers them, and returns the rows one at a time."
-  s.email = "allen.fair@gmail.com"
-  s.extra_rdoc_files = [
-    "LICENSE",
-    "README.rdoc"
-  ]
-  s.files = [
-    ".document",
-    "LICENSE",
-    "README.rdoc",
-    "Rakefile",
-    "VERSION",
-    "lib/postgresql_cursor.rb",
-    "postgresql_cursor.gemspec",
-    "test/helper.rb",
-    "test/test_postgresql_cursor.rb"
-  ]
-  s.homepage = "http://github.com/afair/postgresql_cursor"
-  s.rubygems_version = "2.2.1"
-  s.summary = "ActiveRecord PostgreSQL Adapter extension for using a cursor to return a large result set"
+  spec.files         = `git ls-files -z`.split("\x0")
+  spec.executables   = spec.files.grep(%r{^bin/}) { |f| File.basename(f) }
+  spec.test_files    = spec.files.grep(%r{^(test|spec|features)/})
+  spec.require_paths = ["lib"]
 
-  if s.respond_to? :specification_version then
-    s.specification_version = 4
+ #spec.add_dependency "pg" # Remove this for jruby, which should specify 'activerecord-jdbcpostgresql-adapter'
+  spec.add_dependency "activerecord", ">= 3.2.0"
 
-    if Gem::Version.new(Gem::VERSION) >= Gem::Version.new('1.2.0') then
-      s.add_runtime_dependency(%q<activerecord>, [">= 0"])
-    else
-      s.add_dependency(%q<activerecord>, [">= 0"])
-    end
-  else
-    s.add_dependency(%q<activerecord>, [">= 0"])
-  end
+  spec.add_development_dependency "pg"
+  spec.add_development_dependency "rake"
+  spec.add_development_dependency "minitest"
 end
-