Merge pull request #448 from sodabrew/check_streaming_errors

Check for closed connection while streaming results

Author: sodabrew

ext/mysql2/result.c

@@ -430,6 +430,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;
+  const char * errstr;
   int symbolizeKeys = 0, asArray = 0, castBool = 0, cacheRows = 1, cast = 1, streaming = 0;
   MYSQL_FIELD * fields = NULL;
 
@@ -493,7 +494,7 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
   }
 
   if (wrapper->lastRowProcessed == 0) {
-    if(streaming) {
+    if (streaming) {
       /* We can't get number of rows if we're streaming, */
       /* until we've finished fetching all rows */
       wrapper->numberOfRows = 0;
@@ -509,7 +510,7 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
   }
 
   if (streaming) {
-    if(!wrapper->streamingComplete) {
+    if (!wrapper->streamingComplete) {
       VALUE row;
 
       fields = mysql_fetch_fields(wrapper->result);
@@ -527,6 +528,13 @@ static VALUE rb_mysql_result_each(int argc, VALUE * argv, VALUE self) {
 
       wrapper->numberOfRows = wrapper->lastRowProcessed;
       wrapper->streamingComplete = 1;
+
+      // Check for errors, the connection might have gone out from under us
+      // mysql_error returns an empty string if there is no error
+      errstr = mysql_error(wrapper->client_wrapper->client);
+      if (errstr[0]) {
+        rb_raise(cMysql2Error, "%s", errstr);
+      }
     } else {
       rb_raise(cMysql2Error, "You have already fetched all the rows for this query and streaming is true. (to reiterate you must requery).");
     }

spec/mysql2/result_spec.rb

@@ -6,35 +6,6 @@
     @result = @client.query "SELECT 1"
   end
 
-  it "should maintain a count while streaming" do
-    result = @client.query('SELECT 1')
-
-    result.count.should eql(1)
-    result.each.to_a
-    result.count.should eql(1)
-  end
-
-  it "should set the actual count of rows after streaming" do
-      @client.query "USE test"
-      result = @client.query("SELECT * FROM mysql2_test", :stream => true, :cache_rows => false)
-      result.count.should eql(0)
-      result.each {|r|  }
-      result.count.should eql(1)
-  end
-
-  it "should not yield nil at the end of streaming" do
-    result = @client.query('SELECT * FROM mysql2_test', :stream => true)
-    result.each { |r| r.should_not be_nil}
-  end
-
-  it "#count should be zero for rows after streaming when there were no results " do
-      @client.query "USE test"
-      result = @client.query("SELECT * FROM mysql2_test WHERE null_test IS NOT NULL", :stream => true, :cache_rows => false)
-      result.count.should eql(0)
-      result.each.to_a
-      result.count.should eql(0)
-  end
-
   it "should have included Enumerable" do
     Mysql2::Result.ancestors.include?(Enumerable).should be_true
   end
@@ -121,7 +92,6 @@
 
   context "#fields" do
     before(:each) do
-      @client.query "USE test"
       @test_result = @client.query("SELECT * FROM mysql2_test ORDER BY id DESC LIMIT 1")
     end
 
@@ -135,9 +105,59 @@
     end
   end
 
+  context "streaming" do
+    it "should maintain a count while streaming" do
+      result = @client.query('SELECT 1')
+
+      result.count.should eql(1)
+      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)
+      result.count.should eql(0)
+      result.each {|r|  }
+      result.count.should eql(1)
+    end
+
+    it "should not yield nil at the end of streaming" do
+      result = @client.query('SELECT * FROM mysql2_test', :stream => true, :cache_rows => false)
+      result.each { |r| r.should_not be_nil}
+    end
+
+    it "#count should be zero for rows after streaming when there were no results" do
+      result = @client.query("SELECT * FROM mysql2_test WHERE null_test IS NOT NULL", :stream => true, :cache_rows => false)
+      result.count.should eql(0)
+      result.each.to_a
+      result.count.should eql(0)
+    end
+
+    it "should raise an exception if streaming ended due to a timeout" do
+      # Create an extra client instance, since we're going to time it out
+      client = Mysql2::Client.new DatabaseCredentials['root']
+      client.query "CREATE TEMPORARY TABLE streamingTest (val BINARY(255))"
+
+      # Insert enough records to force the result set into multiple reads
+      # (the BINARY type is used simply because it forces full width results)
+      10000.times do |i|
+        client.query "INSERT INTO streamingTest (val) VALUES ('Foo #{i}')"
+      end
+
+      client.query "SET net_write_timeout = 1"
+      res = client.query "SELECT * FROM streamingTest", :stream => true, :cache_rows => false
+
+      lambda {
+        res.each_with_index do |row, i|
+          # Exhaust the first result packet then trigger a timeout
+          sleep 2 if i > 0 && i % 1000 == 0
+        end
+      }.should raise_error(Mysql2::Error, /Lost connection/)
+    end
+  end
+
   context "row data type mapping" do
     before(:each) do
-      @client.query "USE test"
       @test_result = @client.query("SELECT * FROM mysql2_test ORDER BY id DESC LIMIT 1").first
     end
 
@@ -323,7 +343,6 @@
           result['enum_test'].encoding.should eql(Encoding.find('utf-8'))
 
           client2 = Mysql2::Client.new(DatabaseCredentials['root'].merge(:encoding => 'ascii'))
-          client2.query "USE test"
           result = client2.query("SELECT * FROM mysql2_test ORDER BY id DESC LIMIT 1").first
           result['enum_test'].encoding.should eql(Encoding.find('us-ascii'))
           client2.close
@@ -353,7 +372,6 @@
           result['set_test'].encoding.should eql(Encoding.find('utf-8'))
 
           client2 = Mysql2::Client.new(DatabaseCredentials['root'].merge(:encoding => 'ascii'))
-          client2.query "USE test"
           result = client2.query("SELECT * FROM mysql2_test ORDER BY id DESC LIMIT 1").first
           result['set_test'].encoding.should eql(Encoding.find('us-ascii'))
           client2.close
@@ -436,7 +454,6 @@
               result[field].encoding.should eql(Encoding.find('utf-8'))
 
               client2 = Mysql2::Client.new(DatabaseCredentials['root'].merge(:encoding => 'ascii'))
-              client2.query "USE test"
               result = client2.query("SELECT * FROM mysql2_test ORDER BY id DESC LIMIT 1").first
               result[field].encoding.should eql(Encoding.find('us-ascii'))
               client2.close