Cleaned up code

lib/postgresql_cursor/active_record/relation/cursor_iterators.rb

@@ -20,11 +20,9 @@ module CursorIterators
         # Returns the number of rows yielded to the block
         def each_row(options={}, &block)
           options = {:connection => self.connection}.merge(options)
-          if block_given?
-            PostgreSQLCursor::Cursor.new(to_sql, options).each(&block)
-          else
-            PostgreSQLCursor::Cursor.new(to_sql, options)
-          end
+          cursor  = PostgreSQLCursor::Cursor.new(to_sql, options)
+          return cursor.each_row(&block) if block_given?
+          cursor
         end
         alias :each_hash :each_row
 
@@ -39,39 +37,27 @@ def each_row(options={}, &block)
         # Returns the number of rows yielded to the block
         def each_instance(options={}, &block)
           options = {:connection => self.connection}.merge(options)
-          options[:symbolize_keys] = false # Must be strings to initiate
-
-          if block_given?
-            PostgreSQLCursor::Cursor.new(to_sql, options).each do |row, column_types|
-              model = ::ActiveRecord::VERSION::MAJOR < 4 ?  instantiate(row) : instantiate(row, column_types)
-              yield model
-            end
-          else
-            PostgreSQLCursor::Cursor.new(to_sql, options).instance_iterator(self)
-          end
+          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)
-          pluck_method(:each_row, *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)
-          pluck_method(:each_instance, *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
 
-        def pluck_method(method, *cols)
-          options = cols.last.is_a?(Hash) ? cols.pop : {}
-          cols = cols.map {|c| c.to_sym }
-          result = []
-          self.send(method, options) 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
       end
     end
   end

lib/postgresql_cursor/active_record/sql_cursor.rb

@@ -49,11 +49,9 @@ def each_instance(options={}, &block)
       # Returns the number of rows yielded to the block
       def each_row_by_sql(sql, options={}, &block)
         options = {:connection => self.connection}.merge(options)
-        if block_given?
-          PostgreSQLCursor::Cursor.new(sql, options).each(&block)
-        else
-          PostgreSQLCursor::Cursor.new(sql, options)
-        end
+        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
 
@@ -69,27 +67,24 @@ def each_row_by_sql(sql, options={}, &block)
       # Returns the number of rows yielded to the block
       def each_instance_by_sql(sql, options={}, &block)
         options = {:connection => self.connection}.merge(options)
-        if block_given?
-          PostgreSQLCursor::Cursor.new(sql, options).each do |row, column_types|
-            model = ::ActiveRecord::VERSION::MAJOR < 4 ?  instantiate(row) : instantiate(row, column_types)
-            yield model
-          end
-        else
-          PostgreSQLCursor::Cursor.new(sql, options).instance_iterator(self)
-        end
+        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)
-        all.pluck_row(*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)
-        all.pluck_instance(*cols)
+        options = cols.last.is_a?(Hash) ? cols.pop : {}
+        all.each_instance(options).pluck(*cols)
       end
       alias :pluck_instance :pluck_instances
     end

lib/postgresql_cursor/cursor.rb

@@ -41,23 +41,76 @@ def initialize(sql, options={})
       @options    = options
       @connection = @options.fetch(:connection) { ::ActiveRecord::Base.connection }
       @count      = 0
-      @iterate    = :rows
+      @iterate    = options[:instances] ? :each_instance : :each_row
     end
 
-    def instance_iterator(type)
-      @iterate = :instances
-      @type    = type
+    # 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
@@ -68,18 +121,7 @@ def each(&block)
           while (row = fetch) do
             break if row.size==0
             @count += 1
-            if @iterate == :instances
-              model = if ::ActiveRecord::VERSION::MAJOR < 4
-                @type.send(:instantiate,row)
-              else
-                @type.send(:instantiate,row, column_types)
-              end
-              rc = yield(model)
-            else
-              row = cast_types(row, column_types) if options[:symbolize_keys]
-              row = row.symbolize_keys if options[:cast]
-              rc = yield(row, column_types)
-            end
+            rc = block.call(row)
             break if has_do_until && rc == @options[:until]
             break if has_do_while && rc != @options[:while]
           end
@@ -109,7 +151,6 @@ def column_types
           warn "unknown OID: #{fname}(#{oid}) (#{sql})"
           OID::Identity.new
         }
-
       end
 
       @column_types = types

test/test_postgresql_cursor.rb

@@ -59,5 +59,24 @@ def test_exception
       assert_equal e.message, 'Oops'
     end
   end
+
+  def test_cursor
+    cursor = Product.all.each_row
+    assert cursor.respond_to?(:each)
+    r = cursor.map { |row| row["id"] }
+    assert_equal 1000, r.size
+    cursor = Product.each_row_by_sql("select * from products")
+    assert cursor.respond_to?(:each)
+    r = cursor.map { |row| row["id"] }
+    assert_equal 1000, r.size
+  end
+
+  def test_pluck
+    r = Product.pluck_rows(:id)
+    assert_equal 1000, r.size
+    r = Product.all.pluck_instances(:id)
+    assert_equal 1000, r.size
+    assert_equal Fixnum, r.first.class
+  end
 
 end