Merge pull request trilogy-libraries#264 from byroot/abandon_results

Implement Trilogy#abandon_results!

Author: byroot

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Support `caching_sha2_password` over TCP without TLS by requesting the server RSA public key when needed. #26
 - Now raise an explicit error when a single connection is being used concurrently by multiple threads or fibers. #226.
+- `Trilogy#abandon_results!` as an optimized alternative to `client.next_result while client.more_results_exist?`.
 
 ### Fixed
 

contrib/ruby/ext/trilogy-ruby/cext.c

@@ -1092,6 +1092,31 @@ static VALUE rb_trilogy_more_results_exist(VALUE self)
     }
 }
 
+static VALUE rb_trilogy_abandon_results(VALUE self)
+{
+    struct trilogy_ctx *ctx = get_open_ctx(self);
+
+    long count = 0;
+    while (ctx->conn.server_status & TRILOGY_SERVER_STATUS_MORE_RESULTS_EXISTS) {
+        count++;
+        int rc = trilogy_drain_results(&ctx->conn);
+        while (rc == TRILOGY_AGAIN) {
+            rc = trilogy_sock_wait_read(ctx->conn.socket);
+            if (rc != TRILOGY_OK) {
+                handle_trilogy_error(ctx, rc, "trilogy_sock_wait_read");
+            }
+
+            rc = trilogy_drain_results(&ctx->conn);
+        }
+
+        if (rc != TRILOGY_OK) {
+            handle_trilogy_error(ctx, rc, "trilogy_drain_results");
+        }
+    }
+
+    return LONG2NUM(count);
+}
+
 static VALUE rb_trilogy_query(VALUE self, VALUE query)
 {
     struct trilogy_ctx *ctx = get_open_ctx(self);
@@ -1356,6 +1381,7 @@ RUBY_FUNC_EXPORTED void Init_cext(void)
     rb_define_method(Trilogy, "server_version", rb_trilogy_server_version, 0);
     rb_define_method(Trilogy, "more_results_exist?", rb_trilogy_more_results_exist, 0);
     rb_define_method(Trilogy, "next_result", rb_trilogy_next_result, 0);
+    rb_define_method(Trilogy, "abandon_results!", rb_trilogy_abandon_results, 0);
     rb_define_method(Trilogy, "set_server_option", rb_trilogy_set_server_option, 1);
     rb_define_const(Trilogy, "TLS_VERSION_10", INT2NUM(TRILOGY_TLS_VERSION_10));
     rb_define_const(Trilogy, "TLS_VERSION_11", INT2NUM(TRILOGY_TLS_VERSION_11));

contrib/ruby/test/client_test.rb

@@ -210,6 +210,21 @@ def test_trilogy_multiple_results
     refute client.more_results_exist?
   end
 
+  def test_trilogy_abandon_results
+    client = new_tcp_client(multi_statement: true)
+    create_test_table(client)
+
+    client.query("INSERT INTO trilogy_test (int_test) VALUES ('4')")
+    client.query("INSERT INTO trilogy_test (int_test) VALUES ('3')")
+    client.query("INSERT INTO trilogy_test (int_test) VALUES ('1')")
+
+    results = []
+
+    assert_equal [[1, 4]], client.query("SELECT id, int_test FROM trilogy_test WHERE id = 1; SELECT id, int_test FROM trilogy_test WHERE id IN (2, 3); SELECT id, int_test FROM trilogy_test").to_a
+    assert_equal 2, client.abandon_results!
+    assert_equal [[2]], client.query("SELECT 2").to_a
+  end
+
   def test_trilogy_multiple_results_doesnt_allow_multi_statement_queries
     client = new_tcp_client
     create_test_table(client)

src/client.c

@@ -1018,6 +1018,7 @@ int trilogy_drain_results(trilogy_conn_t *conn)
         }
 
         if (current_packet_type(conn) == TRILOGY_PACKET_EOF && conn->packet_buffer.len < 9) {
+            read_eof_packet(conn);
             return TRILOGY_OK;
         }
     }