Merge pull request #446 from ged/scheduler-stream

Fix compatibility of res.stream_each_* methods

Author: larskanis

ext/pg.h

@@ -344,6 +344,7 @@ void pg_typemap_compact _(( void * ));
 
 PGconn *pg_get_pgconn _(( VALUE ));
 t_pg_connection *pg_get_connection _(( VALUE ));
+VALUE pgconn_block _(( int, VALUE *, VALUE ));
 
 VALUE pg_new_result _(( PGresult *, VALUE ));
 VALUE pg_new_result_autoclear _(( PGresult *, VALUE ));

ext/pg_connection.c

@@ -2972,7 +2972,7 @@ get_result_readable(PGconn *conn)
  * If +true+ is returned, +conn.is_busy+ will return +false+
  * and +conn.get_result+ will not block.
  */
-static VALUE
+VALUE
 pgconn_block( int argc, VALUE *argv, VALUE self ) {
 struct timeval timeout;
 struct timeval *ptimeout = NULL;

ext/pg_result.c

@@ -1468,6 +1468,11 @@ pgresult_stream_any(VALUE self, void (*yielder)(VALUE, int, int, void*), void* d
 
 yielder( self, ntuples, nfields, data );
 
+if( gvl_PQisBusy(pgconn) ){
+/* wait for input (without blocking) before reading each result */
+pgconn_block( 0, NULL, this->connection );
+}
+
 pgresult = gvl_PQgetResult(pgconn);
 if( pgresult == NULL )
 rb_raise( rb_eNoResultError, "no result received - possibly an intersection with another result retrieval");

spec/pg/scheduler_spec.rb

@@ -158,6 +158,19 @@ def run_with_scheduler(timeout=10)
 end
 end
 
+it "can use stream_each_* methods" do
+run_with_scheduler do |conn|
+conn.send_query( "SELECT generate_series(0,999);" )
+conn.set_single_row_mode
+
+start_time = Time.now
+res = conn.get_result
+rows = res.stream_each_row.to_a
+
+expect( rows ).to eq( (0..999).map{ [_1.to_s] } )
+end
+end
+
 it "can receive COPY data" do
 run_with_scheduler do |conn|
 rows = []