extralite 2.4 → 2.6

data/ext/extralite/common.c

@@ -31,7 +31,7 @@ static inline VALUE get_column_value(sqlite3_stmt *stmt, int col, int type) {
   return Qnil;
 }
 
-void bind_parameter_value(sqlite3_stmt *stmt, int pos, VALUE value);
+int bind_parameter_value(sqlite3_stmt *stmt, int pos, VALUE value);
 
 static inline void bind_key_value(sqlite3_stmt *stmt, VALUE k, VALUE v) {
   switch (TYPE(k)) {
@@ -72,24 +72,24 @@ void bind_struct_parameter_values(sqlite3_stmt *stmt, VALUE struct_obj) {
   RB_GC_GUARD(members);
 }
 
-inline void bind_parameter_value(sqlite3_stmt *stmt, int pos, VALUE value) {
+inline int bind_parameter_value(sqlite3_stmt *stmt, int pos, VALUE value) {
   switch (TYPE(value)) {
     case T_NIL:
       sqlite3_bind_null(stmt, pos);
-      return;
+      return 1;
     case T_FIXNUM:
     case T_BIGNUM:
       sqlite3_bind_int64(stmt, pos, NUM2LL(value));
-      return;
+      return 1;
     case T_FLOAT:
       sqlite3_bind_double(stmt, pos, NUM2DBL(value));
-      return;
+      return 1;
     case T_TRUE:
       sqlite3_bind_int(stmt, pos, 1);
-      return;
+      return 1;
     case T_FALSE:
       sqlite3_bind_int(stmt, pos, 0);
-      return;
+      return 1;
     case T_SYMBOL:
       value = rb_sym2str(value);
     case T_STRING:
@@ -97,13 +97,20 @@ inline void bind_parameter_value(sqlite3_stmt *stmt, int pos, VALUE value) {
         sqlite3_bind_blob(stmt, pos, RSTRING_PTR(value), RSTRING_LEN(value), SQLITE_TRANSIENT);
       else
         sqlite3_bind_text(stmt, pos, RSTRING_PTR(value), RSTRING_LEN(value), SQLITE_TRANSIENT);
-      return;
+      return 1;
+    case T_ARRAY:
+      {
+        int count = RARRAY_LEN(value);
+        for (int i = 0; i < count; i++)
+          bind_parameter_value(stmt, pos + i, RARRAY_AREF(value, i));
+        return count;
+      }
     case T_HASH:
       bind_hash_parameter_values(stmt, value);
-      return;
+      return 0;
     case T_STRUCT:
       bind_struct_parameter_values(stmt, value);
-      return;
+      return 0;
     default:
       rb_raise(cParameterError, "Cannot bind parameter at position %d of type %"PRIsVALUE"",
         pos, rb_class_name(rb_obj_class(value)));
@@ -111,14 +118,18 @@ inline void bind_parameter_value(sqlite3_stmt *stmt, int pos, VALUE value) {
 }
 
 inline void bind_all_parameters(sqlite3_stmt *stmt, int argc, VALUE *argv) {
-  for (int i = 0; i < argc; i++) bind_parameter_value(stmt, i + 1, argv[i]);
+  int pos = 1;
+  for (int i = 0; i < argc; i++) {
+    pos += bind_parameter_value(stmt, pos, argv[i]);
+  }
 }
 
 inline void bind_all_parameters_from_object(sqlite3_stmt *stmt, VALUE obj) {
   if (TYPE(obj) == T_ARRAY) {
+    int pos = 1;
     int count = RARRAY_LEN(obj);
     for (int i = 0; i < count; i++)
-      bind_parameter_value(stmt, i + 1, RARRAY_AREF(obj, i));
+      pos += bind_parameter_value(stmt, pos, RARRAY_AREF(obj, i));
   }
   else
     bind_parameter_value(stmt, 1, obj);
@@ -194,9 +205,9 @@ statements. It will release the GVL while the statements are being prepared and
 executed. All statements excluding the last one are executed. The last statement
 is not executed, but instead handed back to the caller for looping over results.
 */
-void prepare_multi_stmt(sqlite3 *db, sqlite3_stmt **stmt, VALUE sql) {
+void prepare_multi_stmt(enum gvl_mode mode, sqlite3 *db, sqlite3_stmt **stmt, VALUE sql) {
   prepare_stmt_ctx ctx = {db, stmt, RSTRING_PTR(sql), RSTRING_LEN(sql), 0};
-  gvl_call(GVL_RELEASE, prepare_multi_stmt_impl, (void *)&ctx);
+  gvl_call(mode, prepare_multi_stmt_impl, (void *)&ctx);
   RB_GC_GUARD(sql);
 
   switch (ctx.rc) {
@@ -236,9 +247,9 @@ end:
   return NULL;
 }
 
-void prepare_single_stmt(sqlite3 *db, sqlite3_stmt **stmt, VALUE sql) {
+void prepare_single_stmt(enum gvl_mode mode, sqlite3 *db, sqlite3_stmt **stmt, VALUE sql) {
   prepare_stmt_ctx ctx = {db, stmt, RSTRING_PTR(sql), RSTRING_LEN(sql), 0};
-  gvl_call(GVL_RELEASE, prepare_single_stmt_impl, (void *)&ctx);
+  gvl_call(mode, prepare_single_stmt_impl, (void *)&ctx);
   RB_GC_GUARD(sql);
 
   switch (ctx.rc) {
@@ -423,8 +434,80 @@ VALUE safe_query_single_value(query_ctx *ctx) {
   return value;
 }
 
-VALUE safe_execute_multi(query_ctx *ctx) {
+enum batch_mode {
+  BATCH_EXECUTE,
+  BATCH_QUERY_ARY,
+  BATCH_QUERY_HASH,
+  BATCH_QUERY_SINGLE_COLUMN
+};
+
+static inline VALUE batch_iterate_hash(query_ctx *ctx) {
+  VALUE rows = rb_ary_new();
+  VALUE row = Qnil;
+  int column_count = sqlite3_column_count(ctx->stmt);
+  VALUE column_names = get_column_names(ctx->stmt, column_count);
+
+  while (stmt_iterate(ctx)) {
+    row = row_to_hash(ctx->stmt, column_count, column_names);
+    rb_ary_push(rows, row);
+  }
+
+  RB_GC_GUARD(column_names);
+  RB_GC_GUARD(rows);
+  return rows;
+}
+
+static inline VALUE batch_iterate_ary(query_ctx *ctx) {
+  VALUE rows = rb_ary_new();
+  VALUE row = Qnil;
+  int column_count = sqlite3_column_count(ctx->stmt);
+
+  while (stmt_iterate(ctx)) {
+    row = row_to_ary(ctx->stmt, column_count);
+    rb_ary_push(rows, row);
+  }
+
+  RB_GC_GUARD(rows);
+  return rows;
+}
+
+static inline VALUE batch_iterate_single_column(query_ctx *ctx) {
+  VALUE rows = rb_ary_new();
+  VALUE value = Qnil;
+  int column_count = sqlite3_column_count(ctx->stmt);
+  if (column_count != 1) rb_raise(cError, "Expected query result to have 1 column");
+
+  while (stmt_iterate(ctx)) {
+    value = get_column_value(ctx->stmt, 0, sqlite3_column_type(ctx->stmt, 0));
+    rb_ary_push(rows, value);
+  }
+
+  RB_GC_GUARD(rows);
+  return rows;
+}
+
+static inline void batch_iterate(query_ctx *ctx, enum batch_mode mode, VALUE *rows) {
+  switch (mode) {
+    case BATCH_EXECUTE:
+      while (stmt_iterate(ctx));
+      break;
+    case BATCH_QUERY_ARY:
+      *rows = batch_iterate_ary(ctx);
+      break;
+    case BATCH_QUERY_HASH:
+      *rows = batch_iterate_hash(ctx);
+      break;
+    case BATCH_QUERY_SINGLE_COLUMN:
+      *rows = batch_iterate_single_column(ctx);
+      break;
+  }
+}
+
+static inline VALUE batch_run_array(query_ctx *ctx, enum batch_mode mode) {
   int count = RARRAY_LEN(ctx->params);
+  int block_given = rb_block_given_p();
+  VALUE results = (mode != BATCH_EXECUTE) && !block_given ? rb_ary_new() : Qnil;
+  VALUE rows = Qnil;
   int changes = 0;
 
   for (int i = 0; i < count; i++) {
@@ -432,11 +515,135 @@ VALUE safe_execute_multi(query_ctx *ctx) {
     sqlite3_clear_bindings(ctx->stmt);
     bind_all_parameters_from_object(ctx->stmt, RARRAY_AREF(ctx->params, i));
 
-    while (stmt_iterate(ctx));
+    batch_iterate(ctx, mode, &rows);
     changes += sqlite3_changes(ctx->sqlite3_db);
+
+    if (mode != BATCH_EXECUTE) {
+      if (block_given)
+        rb_yield(rows);
+      else
+        rb_ary_push(results, rows);
+    }
+  }
+
+  RB_GC_GUARD(rows);
+  RB_GC_GUARD(results);
+
+  if (mode == BATCH_EXECUTE || block_given)
+    return INT2FIX(changes);
+  else
+    return results;
+}
+
+struct batch_execute_each_ctx {
+  query_ctx *ctx;
+  enum batch_mode mode;
+  int block_given;
+  VALUE results;
+  int changes;
+};
+
+static VALUE batch_run_each_iter(RB_BLOCK_CALL_FUNC_ARGLIST(yield_value, vctx)) {
+  struct batch_execute_each_ctx *each_ctx = (struct batch_execute_each_ctx*)vctx;
+  VALUE rows = Qnil;
+
+  sqlite3_reset(each_ctx->ctx->stmt);
+  sqlite3_clear_bindings(each_ctx->ctx->stmt);
+  bind_all_parameters_from_object(each_ctx->ctx->stmt, yield_value);
+
+  batch_iterate(each_ctx->ctx, each_ctx->mode, &rows);
+  each_ctx->changes += sqlite3_changes(each_ctx->ctx->sqlite3_db);
+
+  if (each_ctx->mode != BATCH_EXECUTE) {
+    if (each_ctx->block_given)
+      rb_yield(rows);
+    else
+      rb_ary_push(each_ctx->results, rows);
+  }
+  RB_GC_GUARD(rows);
+
+  return Qnil;
+}
+
+static inline VALUE batch_run_each(query_ctx *ctx, enum batch_mode mode) {
+  struct batch_execute_each_ctx each_ctx = {
+    .ctx          = ctx,
+    .mode         = mode,
+    .block_given  = rb_block_given_p(),
+    .results      = ((mode != BATCH_EXECUTE) && !rb_block_given_p() ? rb_ary_new() : Qnil),
+    .changes      = 0
+  };
+  rb_block_call(ctx->params, ID_each, 0, 0, batch_run_each_iter, (VALUE)&each_ctx);
+
+  if (mode == BATCH_EXECUTE || each_ctx.block_given)
+    return INT2FIX(each_ctx.changes);
+  else
+    return each_ctx.results;
+}
+
+static inline VALUE batch_run_proc(query_ctx *ctx, enum batch_mode mode) {
+  VALUE params = Qnil;
+  int block_given = rb_block_given_p();
+  VALUE results = (mode != BATCH_EXECUTE) && !block_given ? rb_ary_new() : Qnil;
+  VALUE rows = Qnil;
+  int changes = 0;
+
+  while (1) {
+    params = rb_funcall(ctx->params, ID_call, 0);
+    if (NIL_P(params)) break;
+
+    sqlite3_reset(ctx->stmt);
+    sqlite3_clear_bindings(ctx->stmt);
+    bind_all_parameters_from_object(ctx->stmt, params);
+
+    batch_iterate(ctx, mode, &rows);
+    changes += sqlite3_changes(ctx->sqlite3_db);
+
+    if (mode != BATCH_EXECUTE) {
+      if (block_given)
+        rb_yield(rows);
+      else
+        rb_ary_push(results, rows);
+    }
   }
 
-  return INT2FIX(changes);
+  RB_GC_GUARD(rows);
+  RB_GC_GUARD(results);
+  RB_GC_GUARD(params);
+
+  if (mode == BATCH_EXECUTE || block_given)
+    return INT2FIX(changes);
+  else
+    return results;
+}
+
+static inline VALUE batch_run(query_ctx *ctx, enum batch_mode mode) {
+  if (TYPE(ctx->params) == T_ARRAY)
+    return batch_run_array(ctx, mode);
+  
+  if (rb_respond_to(ctx->params, ID_each))
+    return batch_run_each(ctx, mode);
+  
+  if (rb_respond_to(ctx->params, ID_call))
+    return batch_run_proc(ctx, mode);
+  
+  rb_raise(cParameterError, "Invalid parameter source supplied to #batch_execute");
+}
+
+VALUE safe_batch_execute(query_ctx *ctx) {
+  return batch_run(ctx, BATCH_EXECUTE);
+}
+
+VALUE safe_batch_query(query_ctx *ctx) {
+  return batch_run(ctx, BATCH_QUERY_HASH);
+}
+
+VALUE safe_batch_query_ary(query_ctx *ctx) {
+  return batch_run(ctx, BATCH_QUERY_ARY);
+}
+
+VALUE safe_batch_query_single_column(query_ctx *ctx) {
+  return batch_run(ctx, BATCH_QUERY_SINGLE_COLUMN);
 }
 
 VALUE safe_query_columns(query_ctx *ctx) {