Since streaming is a query option, process it in rb_mysql_result_to_obj instead of rb_mysql_result_each

Author: sodabrew

ext/mysql2/result.c

@@ -446,7 +446,7 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
   mysql2_result_wrapper * wrapper;
   unsigned long i;
   const char * errstr;
-  int symbolizeKeys = 0, asArray = 0, castBool = 0, cacheRows = 1, cast = 1, streaming = 0;
+  int symbolizeKeys = 0, asArray = 0, castBool = 0, cacheRows = 1, cast = 1;
   MYSQL_FIELD * fields = NULL;
 
   GetMysql2Result(self, wrapper);
@@ -479,11 +479,7 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
     cast = 0;
   }
 
-  if (rb_hash_aref(opts, sym_stream) == Qtrue) {
-    streaming = 1;
-  }
-
-  if (streaming && cacheRows) {
+  if (wrapper->is_streaming && cacheRows) {
     rb_warn("cacheRows is ignored if streaming is true");
   }
 
@@ -509,7 +505,7 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
   }
 
   if (wrapper->lastRowProcessed == 0) {
-    if (streaming) {
+    if (wrapper->is_streaming) {
       /* We can't get number of rows if we're streaming, */
       /* until we've finished fetching all rows */
       wrapper->numberOfRows = 0;
@@ -524,7 +520,7 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
     }
   }
 
-  if (streaming) {
+  if (wrapper->is_streaming) {
     if (!wrapper->streamingComplete) {
       VALUE row;
 
@@ -601,12 +597,12 @@ static VALUE rb_mysql_result_count(VALUE self) {
   mysql2_result_wrapper *wrapper;
 
   GetMysql2Result(self, wrapper);
+  if (wrapper->is_streaming) {
+    return LONG2NUM(wrapper->numberOfRows);
+  }
+
   if (wrapper->resultFreed) {
-    if (wrapper->streamingComplete){
-      return LONG2NUM(wrapper->numberOfRows);
-    } else {
-      return LONG2NUM(RARRAY_LEN(wrapper->rows));
-    }
+    return LONG2NUM(RARRAY_LEN(wrapper->rows));
   } else {
     return INT2FIX(mysql_num_rows(wrapper->result));
   }
@@ -634,6 +630,10 @@ VALUE rb_mysql_result_to_obj(VALUE client, VALUE encoding, VALUE options, MYSQL_
 
   rb_iv_set(obj, "@query_options", options);
 
+  /* Options that cannot be changed in results.each(...) { |row| }
+   * should be processed here. */
+  wrapper->is_streaming = (rb_hash_aref(options, sym_stream) == Qtrue ? 1 : 0);
+
   return obj;
 }
 

ext/mysql2/result.h

@@ -12,6 +12,7 @@ typedef struct {
   unsigned int numberOfFields;
   unsigned long numberOfRows;
   unsigned long lastRowProcessed;
+  char is_streaming;
   char streamingComplete;
   char resultFreed;
   MYSQL_RES *result;