Merge pull request brianmario#223 from benofsky/master

Basic implementation of :streaming

Author: brianmario

README.md

@@ -212,6 +212,21 @@ This is especially helpful since it saves the cost of creating the row in Ruby i
 If you only plan on using each row once, then it's much more efficient to disable this behavior by setting the `:cache_rows` option to false.
 This would be helpful if you wanted to iterate over the results in a streaming manner. Meaning the GC would cleanup rows you don't need anymore as you're iterating over the result set.
 
+### Streaming
+
+`Mysql2::Client` can optionally only fetch rows from the server on demand by setting `:stream => true`. This is handy when handling very large result sets which might not fit in memory on the client.
+
+``` ruby
+result = client.query("SELECT * FROM really_big_Table", :stream => true)
+```
+
+There are a few things that need to be kept in mind while using streaming:
+
+* `:cache_rows` is ignored currently. (if you want to use `:cache_rows` you probably don't want to be using `:stream`)
+* You must fetch all rows in the result set of your query before you can make new queries. (i.e. with `Mysql2::Result#each`)
+
+Read more about the consequences of using `mysql_use_result` (what streaming is implemented with) here: http://dev.mysql.com/doc/refman/5.0/en/mysql-use-result.html.
+
 ## ActiveRecord
 
 To use the ActiveRecord driver (with or without rails), all you should need to do is have this gem installed and set the adapter in your database.yml to "mysql2".

ext/mysql2/client.c

@@ -9,7 +9,7 @@
 VALUE cMysql2Client;
 extern VALUE mMysql2, cMysql2Error;
 static VALUE intern_encoding_from_charset;
-static VALUE sym_id, sym_version, sym_async, sym_symbolize_keys, sym_as, sym_array;
+static VALUE sym_id, sym_version, sym_async, sym_symbolize_keys, sym_as, sym_array, sym_stream;
 static ID intern_merge, intern_error_number_eql, intern_sql_state_eql;
 
 #define REQUIRE_OPEN_DB(wrapper) \
@@ -278,21 +278,32 @@ static VALUE nogvl_read_query_result(void *ptr) {
   return res == 0 ? Qtrue : Qfalse;
 }
 
-/* mysql_store_result may (unlikely) read rows off the socket */
-static VALUE nogvl_store_result(void *ptr) {
+static VALUE nogvl_do_result(void *ptr, char use_result) {
   mysql_client_wrapper *wrapper;
   MYSQL_RES *result;
 
   wrapper = (mysql_client_wrapper *)ptr;
-  result = mysql_store_result(wrapper->client);
+  if(use_result) {
+    result = mysql_use_result(wrapper->client);
+  } else {
+    result = mysql_store_result(wrapper->client);
+  }
 
   // once our result is stored off, this connection is
   // ready for another command to be issued
   wrapper->active = 0;
-
   return (VALUE)result;
 }
 
+/* mysql_store_result may (unlikely) read rows off the socket */
+static VALUE nogvl_store_result(void *ptr) {
+  return nogvl_do_result(ptr, 0);
+}
+
+static VALUE nogvl_use_result(void *ptr) {
+  return nogvl_do_result(ptr, 1);
+}
+
 static VALUE rb_mysql_client_async_result(VALUE self) {
   MYSQL_RES * result;
   VALUE resultObj;
@@ -312,7 +323,12 @@ static VALUE rb_mysql_client_async_result(VALUE self) {
     return rb_raise_mysql2_error(wrapper);
   }
 
-  result = (MYSQL_RES *)rb_thread_blocking_region(nogvl_store_result, wrapper, RUBY_UBF_IO, 0);
+  VALUE is_streaming = rb_hash_aref(rb_iv_get(self, "@query_options"), sym_stream);
+  if(is_streaming == Qtrue) {
+    result = (MYSQL_RES *)rb_thread_blocking_region(nogvl_use_result, wrapper, RUBY_UBF_IO, 0);
+  } else {
+    result = (MYSQL_RES *)rb_thread_blocking_region(nogvl_store_result, wrapper, RUBY_UBF_IO, 0);
+  }
 
   if (result == NULL) {
     if (mysql_errno(wrapper->client) != 0) {
@@ -788,6 +804,7 @@ void init_mysql2_client() {
   sym_symbolize_keys  = ID2SYM(rb_intern("symbolize_keys"));
   sym_as              = ID2SYM(rb_intern("as"));
   sym_array           = ID2SYM(rb_intern("array"));
+  sym_stream          = ID2SYM(rb_intern("stream"));
 
   intern_merge = rb_intern("merge");
   intern_error_number_eql = rb_intern("error_number=");

ext/mysql2/client.h

@@ -39,4 +39,5 @@ typedef struct {
   MYSQL *client;
 } mysql_client_wrapper;
 
-#endif
+#endif
+

ext/mysql2/result.c

@@ -55,7 +55,7 @@ static VALUE intern_encoding_from_charset;
 static ID intern_new, intern_utc, intern_local, intern_encoding_from_charset_code,
           intern_localtime, intern_local_offset, intern_civil, intern_new_offset;
 static VALUE sym_symbolize_keys, sym_as, sym_array, sym_database_timezone, sym_application_timezone,
-          sym_local, sym_utc, sym_cast_booleans, sym_cache_rows, sym_cast;
+          sym_local, sym_utc, sym_cast_booleans, sym_cache_rows, sym_cast, sym_stream;
 static ID intern_merge;
 
 static void rb_mysql_result_mark(void * wrapper) {
@@ -392,7 +392,7 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
   ID db_timezone, app_timezone, dbTz, appTz;
   mysql2_result_wrapper * wrapper;
   unsigned long i;
-  int symbolizeKeys = 0, asArray = 0, castBool = 0, cacheRows = 1, cast = 1;
+  int symbolizeKeys = 0, asArray = 0, castBool = 0, cacheRows = 1, cast = 1, streaming = 0;
 
   GetMysql2Result(self, wrapper);
 
@@ -423,6 +423,14 @@ 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) {
+    rb_warn("cacheRows is ignored if streaming is true");
+  }
+
   dbTz = rb_hash_aref(opts, sym_database_timezone);
   if (dbTz == sym_local) {
     db_timezone = intern_local;
@@ -445,48 +453,77 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
   }
 
   if (wrapper->lastRowProcessed == 0) {
-    wrapper->numberOfRows = mysql_num_rows(wrapper->result);
-    if (wrapper->numberOfRows == 0) {
+    if(streaming) {
+      // We can't get number of rows if we're streaming,
+      // until we've finished fetching all rows
+      wrapper->numberOfRows = 0;
       wrapper->rows = rb_ary_new();
-      return wrapper->rows;
+    } 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);
     }
-    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
-    for (i = 0; i < wrapper->numberOfRows; i++) {
-      rb_yield(rb_ary_entry(wrapper->rows, i));
-    }
-  } else {
-    unsigned long rowsProcessed = 0;
-    rowsProcessed = RARRAY_LEN(wrapper->rows);
-    for (i = 0; i < wrapper->numberOfRows; i++) {
+  if (streaming) {
+    if(!wrapper->streamingComplete) {
       VALUE row;
-      if (cacheRows && i < rowsProcessed) {
-        row = rb_ary_entry(wrapper->rows, i);
-      } else {
+
+      do {
         row = rb_mysql_result_fetch_row(self, db_timezone, app_timezone, symbolizeKeys, asArray, castBool, cast);
-        if (cacheRows) {
-          rb_ary_store(wrapper->rows, i, row);
+
+        if (block != Qnil) {
+          rb_yield(row);
+          wrapper->lastRowProcessed++;
         }
-        wrapper->lastRowProcessed++;
+      } while(row != Qnil);
+
+      rb_mysql_result_free_result(wrapper);
+
+      wrapper->numberOfRows = wrapper->lastRowProcessed;
+      wrapper->streamingComplete = 1;
+    } else {
+      rb_raise(cMysql2Error, "You have already fetched all the rows for this query and streaming is true. (to reiterate you must requery).");
+    }
+  } else {
+    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
+      for (i = 0; i < wrapper->numberOfRows; i++) {
+        rb_yield(rb_ary_entry(wrapper->rows, i));
       }
+    } else {
+      unsigned long rowsProcessed = 0;
+      rowsProcessed = RARRAY_LEN(wrapper->rows);
+      for (i = 0; i < wrapper->numberOfRows; i++) {
+        VALUE row;
+        if (cacheRows && i < rowsProcessed) {
+          row = rb_ary_entry(wrapper->rows, i);
+        } else {
+          row = rb_mysql_result_fetch_row(self, db_timezone, app_timezone, symbolizeKeys, asArray, castBool, cast);
+          if (cacheRows) {
+            rb_ary_store(wrapper->rows, i, row);
+          }
+          wrapper->lastRowProcessed++;
+        }
+
+        if (row == Qnil) {
+          // we don't need the mysql C dataset around anymore, peace it
+          rb_mysql_result_free_result(wrapper);
+          return Qnil;
+        }
 
-      if (row == Qnil) {
+        if (block != Qnil) {
+          rb_yield(row);
+        }
+      }
+      if (wrapper->lastRowProcessed == wrapper->numberOfRows) {
         // we don't need the mysql C dataset around anymore, peace it
         rb_mysql_result_free_result(wrapper);
-        return Qnil;
       }
-
-      if (block != Qnil) {
-        rb_yield(row);
-      }
-    }
-    if (wrapper->lastRowProcessed == wrapper->numberOfRows) {
-      // we don't need the mysql C dataset around anymore, peace it
-      rb_mysql_result_free_result(wrapper);
     }
   }
 
@@ -514,6 +551,7 @@ VALUE rb_mysql_result_to_obj(MYSQL_RES * r) {
   wrapper->fields = Qnil;
   wrapper->rows = Qnil;
   wrapper->encoding = Qnil;
+  wrapper->streamingComplete = 0;
   rb_obj_call_init(obj, 0, NULL);
   return obj;
 }
@@ -551,6 +589,7 @@ void init_mysql2_result() {
   sym_application_timezone  = ID2SYM(rb_intern("application_timezone"));
   sym_cache_rows     = ID2SYM(rb_intern("cache_rows"));
   sym_cast           = ID2SYM(rb_intern("cast"));
+  sym_stream         = ID2SYM(rb_intern("stream"));
 
   opt_decimal_zero = rb_str_new2("0.0");
   rb_global_variable(&opt_decimal_zero); //never GC

ext/mysql2/result.h

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

spec/mysql2/client_spec.rb

@@ -92,6 +92,22 @@ def connect *args
   end
 
   context "#query" do
+    it "should let you query again if iterating is finished when streaming" do
+      @client.query("SELECT 1 UNION SELECT 2", :stream => true, :cache_rows => false).each {}
+
+      expect {
+        @client.query("SELECT 1 UNION SELECT 2", :stream => true, :cache_rows => false)
+      }.to_not raise_exception(Mysql2::Error)
+    end
+
+    it "should not let you query again if iterating is not finished when streaming" do
+      @client.query("SELECT 1 UNION SELECT 2", :stream => true, :cache_rows => false).first
+
+      expect {
+        @client.query("SELECT 1 UNION SELECT 2", :stream => true, :cache_rows => false)
+      }.to raise_exception(Mysql2::Error)
+    end
+
     it "should only accept strings as the query parameter" do
       lambda {
         @client.query ["SELECT 'not right'"]

spec/mysql2/result_spec.rb

@@ -73,6 +73,25 @@
       result = @client.query "SELECT 1", :cache_rows => false
       result.first.object_id.should_not eql(result.first.object_id)
     end
+
+    it "should yield different value for #first if streaming" do
+      result = @client.query "SELECT 1 UNION SELECT 2", :stream => true, :cache_rows => false
+      result.first.should_not eql(result.first)
+    end
+
+    it "should yield the same value for #first if streaming is disabled" do
+      result = @client.query "SELECT 1 UNION SELECT 2", :stream => false
+      result.first.should eql(result.first)
+    end
+
+    it "should throw an exception if we try to iterate twice when streaming is enabled" do
+      result = @client.query "SELECT 1 UNION SELECT 2", :stream => true, :cache_rows => false
+
+      expect {
+        result.each {}
+        result.each {}
+      }.to raise_exception(Mysql2::Error)
+    end
   end
 
   context "#fields" do