From 1eef15245973e34e5859b13b5efb72a7b096f203 Mon Sep 17 00:00:00 2001
From: Brian Lopez <seniorlopez@gmail.com>
Date: Tue, 30 Mar 2010 09:14:27 -0700
Subject: [PATCH] allow passing a block to fetch_rows

---
 mysql_duce_ext.c | 42 ++++++++++++++++++++++++++++--------------
 mysql_duce_ext.h |  2 +-
 2 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/mysql_duce_ext.c b/mysql_duce_ext.c
index d5bed6d..9c86823 100644
--- a/mysql_duce_ext.c
+++ b/mysql_duce_ext.c
@@ -97,7 +97,6 @@ static VALUE rb_mysql_result_fetch_row(VALUE self) {
   fields = mysql_fetch_fields(result);
 
   rowHash = rb_hash_new();
-  // TODO: yield to a passed block
   for (i = 0; i < numFields; i++) {
     VALUE key = rb_str_new(fields[i].name, fields[i].name_length);
     if (row[i]) {
@@ -109,27 +108,42 @@ static VALUE rb_mysql_result_fetch_row(VALUE self) {
   return rowHash;
 }
 
-static VALUE rb_mysql_result_fetch_rows(VALUE self) {
-  VALUE dataset;
+static VALUE rb_mysql_result_fetch_rows(int argc, VALUE * argv, VALUE self) {
+  VALUE dataset, block;
   MYSQL_RES * result;
   unsigned long numRows, i;
-  
+
   GetMySQLResult(self, result);
-  
+
+  rb_scan_args(argc, argv, "0&", &block);
+
   numRows = mysql_num_rows(result);
   if (numRows == 0) {
     return Qnil;
   }
-  
-  dataset = rb_ary_new2(numRows);
-  for (i = 0; i < numRows; i++) {
-    VALUE row = rb_mysql_result_fetch_row(self);
-    if (row == Qnil) {
-      return Qnil;
+
+  // TODO: allow yielding datasets of configurable size
+  // like find_in_batches from AR...
+  if (block != Qnil) {
+    for (i = 0; i < numRows; i++) {
+      VALUE row = rb_mysql_result_fetch_row(self);
+      if (row == Qnil) {
+        return Qnil;
+      }
+      rb_yield(row);
     }
-    rb_ary_store(dataset, i, row);
+  } else {
+    dataset = rb_ary_new2(numRows);
+    for (i = 0; i < numRows; i++) {
+      VALUE row = rb_mysql_result_fetch_row(self);
+      if (row == Qnil) {
+        return Qnil;
+      }
+      rb_ary_store(dataset, i, row);
+    }
+    return dataset;
   }
-  return dataset;
+  return Qnil;
 }
 
 /* Ruby Extension initializer */
@@ -143,7 +157,7 @@ void Init_mysql_duce_ext() {
 
   cMySQLResult = rb_define_class_under(mMySQL, "Result", rb_cObject);
   rb_define_method(cMySQLResult, "fetch_row", rb_mysql_result_fetch_row, 0);
-  rb_define_method(cMySQLResult, "fetch_rows", rb_mysql_result_fetch_rows, 0);
+  rb_define_method(cMySQLResult, "fetch_rows", rb_mysql_result_fetch_rows, -1);
 
 #ifdef HAVE_RUBY_ENCODING_H
   utf8Encoding = rb_enc_find_index("UTF-8");
diff --git a/mysql_duce_ext.h b/mysql_duce_ext.h
index 162305f..266a4f1 100644
--- a/mysql_duce_ext.h
+++ b/mysql_duce_ext.h
@@ -27,4 +27,4 @@ VALUE cMySQLResult;
 static VALUE rb_mysql_result_to_obj(MYSQL_RES * res);
 void rb_mysql_result_free(void * result);
 static VALUE rb_mysql_result_fetch_row(VALUE self);
-static VALUE rb_mysql_result_fetch_rows(VALUE self);
\ No newline at end of file
+static VALUE rb_mysql_result_fetch_rows(int argc, VALUE * argv, VALUE self);
\ No newline at end of file