Merge pull request #581 from sodabrew/streaming_procedure_specs

Fixes for streaming and stored procedures with multiple result sets

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, asArray, castBool, cacheRows, cast;
   MYSQL_FIELD * fields = NULL;
 
   GetMysql2Result(self, wrapper);
@@ -459,31 +459,13 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
     opts = defaults;
   }
 
-  if (rb_hash_aref(opts, sym_symbolize_keys) == Qtrue) {
-    symbolizeKeys = 1;
-  }
-
-  if (rb_hash_aref(opts, sym_as) == sym_array) {
-    asArray = 1;
-  }
-
-  if (rb_hash_aref(opts, sym_cast_booleans) == Qtrue) {
-    castBool = 1;
-  }
-
-  if (rb_hash_aref(opts, sym_cache_rows) == Qfalse) {
-    cacheRows = 0;
-  }
-
-  if (rb_hash_aref(opts, sym_cast) == Qfalse) {
-    cast = 0;
-  }
+  symbolizeKeys = RTEST(rb_hash_aref(opts, sym_symbolize_keys));
+  asArray       = rb_hash_aref(opts, sym_as) == sym_array;
+  castBool      = RTEST(rb_hash_aref(opts, sym_cast_booleans));
+  cacheRows     = RTEST(rb_hash_aref(opts, sym_cache_rows));
+  cast          = RTEST(rb_hash_aref(opts, sym_cast));
 
-  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");
   }
 
@@ -508,40 +490,28 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
     app_timezone = Qnil;
   }
 
-  if (wrapper->lastRowProcessed == 0) {
-    if (streaming) {
-      /* We can't get number of rows if we're streaming, */
-      /* until we've finished fetching all rows */
-      wrapper->numberOfRows = 0;
+  if (wrapper->is_streaming) {
+    /* When streaming, we will only yield rows, not return them. */
+    if (wrapper->rows == Qnil) {
       wrapper->rows = rb_ary_new();
-    } else {
-      wrapper->numberOfRows = mysql_num_rows(wrapper->result);
-      if (wrapper->numberOfRows == 0) {
-        wrapper->rows = rb_ary_new();
-        return wrapper->rows;
-      }
-      wrapper->rows = rb_ary_new2(wrapper->numberOfRows);
     }
-  }
 
-  if (streaming) {
     if (!wrapper->streamingComplete) {
       VALUE row;
 
       fields = mysql_fetch_fields(wrapper->result);
 
       do {
         row = rb_mysql_result_fetch_row(self, db_timezone, app_timezone, symbolizeKeys, asArray, castBool, cast, fields);
-
-        if (block != Qnil && row != Qnil) {
-          rb_yield(row);
-          wrapper->lastRowProcessed++;
+        if (row != Qnil) {
+          wrapper->numberOfRows++;
+          if (block != Qnil) {
+            rb_yield(row);
+          }
         }
       } while(row != Qnil);
 
       rb_mysql_result_free_result(wrapper);
-
-      wrapper->numberOfRows = wrapper->lastRowProcessed;
       wrapper->streamingComplete = 1;
 
       // Check for errors, the connection might have gone out from under us
@@ -554,6 +524,15 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
       rb_raise(cMysql2Error, "You have already fetched all the rows for this query and streaming is true. (to reiterate you must requery).");
     }
   } else {
+    if (wrapper->lastRowProcessed == 0) {
+      wrapper->numberOfRows = mysql_num_rows(wrapper->result);
+      if (wrapper->numberOfRows == 0) {
+        wrapper->rows = rb_ary_new();
+        return wrapper->rows;
+      }
+      wrapper->rows = rb_ary_new2(wrapper->numberOfRows);
+    }
+
     if (cacheRows && wrapper->lastRowProcessed == wrapper->numberOfRows) {
       /* we've already read the entire dataset from the C result into our */
       /* internal array. Lets hand that over to the user since it's ready to go */
@@ -601,14 +580,17 @@ static VALUE rb_mysql_result_count(VALUE self) {
   mysql2_result_wrapper *wrapper;
 
   GetMysql2Result(self, wrapper);
+  if (wrapper->is_streaming) {
+    /* This is an unsigned long per result.h */
+    return ULONG2NUM(wrapper->numberOfRows);
+  }
+
   if (wrapper->resultFreed) {
-    if (wrapper->streamingComplete){
-      return LONG2NUM(wrapper->numberOfRows);
-    } else {
-      return LONG2NUM(RARRAY_LEN(wrapper->rows));
-    }
+    /* Ruby arrays have platform signed long length */
+    return LONG2NUM(RARRAY_LEN(wrapper->rows));
   } else {
-    return INT2FIX(mysql_num_rows(wrapper->result));
+    /* MySQL returns an unsigned 64-bit long here */
+    return ULL2NUM(mysql_num_rows(wrapper->result));
   }
 }
 
@@ -634,6 +616,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;

spec/mysql2/client_spec.rb

@@ -614,6 +614,21 @@ def connect *args
 
         @multi_client.more_results?.should be_false
       end
+
+      it "#more_results? should work with stored procedures" do
+        @multi_client.query("DROP PROCEDURE IF EXISTS test_proc")
+        @multi_client.query("CREATE PROCEDURE test_proc() BEGIN SELECT 1 AS 'set_1'; SELECT 2 AS 'set_2'; END")
+        @multi_client.query("CALL test_proc()").first.should eql({ 'set_1' => 1 })
+        @multi_client.more_results?.should be_true
+
+        @multi_client.next_result
+        @multi_client.store_result.first.should eql({ 'set_2' => 2 })
+
+        @multi_client.next_result
+        @multi_client.store_result.should be_nil # this is the result from CALL itself
+
+        @multi_client.more_results?.should be_false
+      end
     end
   end
 

spec/mysql2/result_spec.rb

@@ -107,18 +107,19 @@
 
   context "streaming" do
     it "should maintain a count while streaming" do
-      result = @client.query('SELECT 1')
-
-      result.count.should eql(1)
+      result = @client.query('SELECT 1', :stream => true, :cache_rows => false)
+      result.count.should eql(0)
       result.each.to_a
       result.count.should eql(1)
     end
 
-    it "should set the actual count of rows after streaming" do
-      result = @client.query("SELECT * FROM mysql2_test", :stream => true, :cache_rows => false)
+    it "should retain the count when mixing first and each" do
+      result = @client.query("SELECT 1 UNION SELECT 2", :stream => true, :cache_rows => false)
       result.count.should eql(0)
-      result.each {|r|  }
+      result.first
       result.count.should eql(1)
+      result.each.to_a
+      result.count.should eql(2)
     end
 
     it "should not yield nil at the end of streaming" do