Result #[] with Array-like [offset, count] plus options in [offset, {opts hash}] and [offset, count, {opts hash}].

Author: sodabrew

ext/mysql2/result.c

@@ -786,46 +786,68 @@ static void rb_mysql_row_query_options(VALUE opts, ID *db_timezone, ID *app_time
   }
 }
 
-static VALUE rb_mysql_result_element(VALUE self, VALUE seek) {
+static VALUE rb_mysql_result_element(int argc, VALUE * argv, VALUE self) {
   result_each_args args;
   MYSQL_FIELD *fields = NULL;
-  long offset;
   ID db_timezone, app_timezone;
+  VALUE seek, count, row, rows;
+  long i, c_seek, c_count = 0;
   int symbolizeKeys, asArray, castBool, cacheRows, cast;
-  VALUE opts, (*fetch_row_func)(VALUE, MYSQL_FIELD *fields, const result_each_args *args);
+  VALUE defaults, block, opts, (*fetch_row_func)(VALUE, MYSQL_FIELD *fields, const result_each_args *args);
 
   GET_RESULT(self);
 
-  offset = NUM2LONG(seek);
+  defaults = rb_iv_get(self, "@query_options");
+  Check_Type(defaults, T_HASH);
+  rb_scan_args(argc, argv, "12&", &seek, &count, &opts, &block);
 
-  if (!wrapper->numberOfRows) {
-    wrapper->numberOfRows = mysql_num_rows(wrapper->result);
+  /* If the second arg is a hash, it's the opts and there's no count */
+  if (TYPE(count) == T_HASH) {
+    opts = count;
+    count = Qnil;
+  }
+
+  c_seek = NUM2LONG(seek);
+  if (!NIL_P(count)) {
+    c_count = NUM2LONG(count);
+    /* Special case: ary[x, 0] returns []*/
+    if (!c_count) return rb_ary_new();
+  }
+
+  if (!NIL_P(opts)) {
+    opts = rb_funcall(defaults, intern_merge, 1, opts);
+  } else {
+    opts = defaults;
   }
 
-  opts = rb_iv_get(self, "@query_options");
   rb_mysql_row_query_options(opts, &db_timezone, &app_timezone, &symbolizeKeys, &asArray, &castBool, &cast, &cacheRows);
 
   if (wrapper->is_streaming) {
     rb_raise(cMysql2Error, "Element reference operator #[] cannot be used in streaming mode.");
   }
 
+  if (!wrapper->numberOfRows) {
+    wrapper->numberOfRows = mysql_num_rows(wrapper->result);
+  }
+
   /* count back from the end if passed a negative number */
-  if (offset < 0) {
-    offset = wrapper->numberOfRows + offset;
+  if (c_seek < 0) {
+    c_seek = wrapper->numberOfRows + c_seek;
   }
 
   /* negative offset was too big */
-  if (offset < 0) {
+  if (c_seek < 0) {
     return Qnil;
-    /* rb_raise(cMysql2Error, "Out of range: offset %ld is beyond %lu rows (offset begins at 0).", offset, wrapper->numberOfRows); */
+    /* rb_raise(cMysql2Error, "Out of range: offset %ld is beyond %lu rows (offset begins at 0).", c_seek, wrapper->numberOfRows); */
   }
 
-  if (wrapper->numberOfRows <= (unsigned long)offset) {
-    return Qnil;
-    /* rb_raise(cMysql2Error, "Out of range: offset %ld is beyond %lu rows (offset begins at 0).", offset, wrapper->numberOfRows); */
+  if (wrapper->numberOfRows <= (unsigned long)c_seek) {
+    if (!c_count) return Qnil;
+    else return rb_ary_new();
+    /* rb_raise(cMysql2Error, "Out of range: offset %ld is beyond %lu rows (offset begins at 0).", c_seek, wrapper->numberOfRows); */
   }
 
-  mysql_data_seek(wrapper->result, offset);
+  mysql_data_seek(wrapper->result, c_seek);
 
   // Backward compat
   args.symbolizeKeys = symbolizeKeys;
@@ -835,15 +857,30 @@ static VALUE rb_mysql_result_element(VALUE self, VALUE seek) {
   args.cast = cast;
   args.db_timezone = db_timezone;
   args.app_timezone = app_timezone;
-  args.block_given = Qnil;
+  args.block_given = block;
 
   if (wrapper->stmt) {
     fetch_row_func = rb_mysql_result_fetch_row_stmt;
   } else {
     fetch_row_func = rb_mysql_result_fetch_row;
   }
 
-  return fetch_row_func(self, fields, &args);
+  if (!c_count) {
+    return fetch_row_func(self, fields, &args);
+  }
+
+  /* given ary = [1, 2, 3] then ary[1, 100] returns [2, 3] */
+  if ((unsigned long)(c_seek + c_count) > wrapper->numberOfRows) {
+    c_count = wrapper->numberOfRows - c_seek;
+  }
+
+  /* return an array! */
+  rows = rb_ary_new2(c_count);
+  for (i = 0; i < c_count; i++) {
+    row = fetch_row_func(self, fields, &args);
+    rb_ary_store(rows, i, row);
+  }
+  return rows;
 }
 
 static VALUE rb_mysql_result_each_(VALUE self,
@@ -1042,7 +1079,7 @@ void init_mysql2_result() {
   cDateTime = rb_const_get(rb_cObject, rb_intern("DateTime"));
 
   cMysql2Result = rb_define_class_under(mMysql2, "Result", rb_cObject);
-  rb_define_method(cMysql2Result, "[]", rb_mysql_result_element, 1);
+  rb_define_method(cMysql2Result, "[]", rb_mysql_result_element, -1);
   rb_define_method(cMysql2Result, "each", rb_mysql_result_each, -1);
   rb_define_method(cMysql2Result, "fields", rb_mysql_result_fetch_fields, 0);
   rb_define_method(cMysql2Result, "count", rb_mysql_result_count, 0);

spec/mysql2/result_spec.rb

@@ -100,31 +100,46 @@
 
   context "#[]" do
     it "should return results when accessed by [offset]" do
-      result = @client.query "SELECT 1 UNION SELECT 2"
-      result[1][result.fields.first].should eql(2)
-      result[0][result.fields.first].should eql(1)
+      result = @client.query "SELECT 1 AS col UNION SELECT 2 AS col"
+      result[1].should eql({"col" => 2})
+      result[0].should eql({"col" => 1})
     end
 
     it "should return results when accessed by negative [offset]" do
-      result = @client.query "SELECT 1 UNION SELECT 2"
-      result[-1][result.fields.first].should eql(2)
-      result[-2][result.fields.first].should eql(1)
+      result = @client.query "SELECT 1 AS col UNION SELECT 2 AS col"
+      result[-1].should eql({"col" => 2})
+      result[-2].should eql({"col" => 1})
+    end
+
+    it "should return array of results when accessed by [offset, count]" do
+      result = @client.query "SELECT 1 AS col UNION SELECT 2 AS col"
+      result[1, 1].should eql([{"col" => 2}])
+      result[-2, 10].should eql([{"col" => 1}, {"col" => 2}])
     end
 
     it "should return nil if we use too large [offset]" do
-      result = @client.query "SELECT 1 UNION SELECT 2"
+      result = @client.query "SELECT 1 AS col UNION SELECT 2 AS col"
       result[2].should be_nil
       result[200].should be_nil
     end
 
     it "should return nil if we use too negative [offset]" do
-      result = @client.query "SELECT 1 UNION SELECT 2"
+      result = @client.query "SELECT 1 AS col UNION SELECT 2 AS col"
       result[-3].should be_nil
       result[-300].should be_nil
     end
 
+    it "should accept hash args in [offset, {:foo => bar}] and [offset, count, {:foo => bar}]" do
+      result = @client.query "SELECT 1 AS col UNION SELECT 2 AS col"
+      result[1, {:symbolize_keys => true}].should eql({:col => 2})
+      result[1, 1, {:symbolize_keys => true}].should eql([{:col => 2}])
+      # This syntax does not work in Ruby 1.8:
+      # result[1, :symbolize_keys => true].should eql({:col => 2})
+      # result[1, 1, :symbolize_keys => true].should eql([{:col => 2}])
+    end
+
     it "should throw an exception if we use an [offset] in streaming mode" do
-      result = @client.query "SELECT 1 UNION SELECT 2", :stream => true
+      result = @client.query "SELECT 1 AS col UNION SELECT 2 AS col", :stream => true
       expect { result[0] }.to raise_exception(Mysql2::Error)
     end
   end