Merge branch 'progress-handler' into statement-timeout

Author: fractaledmind

ext/sqlite3/database.c

@@ -251,6 +251,51 @@ busy_handler(int argc, VALUE *argv, VALUE self)
     return self;
 }
 
+static int
+rb_sqlite3_progress_handler(void *ctx)
+{
+  VALUE self = (VALUE)(ctx);
+  VALUE handle = rb_iv_get(self, "@progress_handler");
+  VALUE result = rb_funcall(handle, rb_intern("call"), 0);
+
+  if (Qfalse == result) return 1;
+
+  return 0;
+}
+
+/* call-seq:
+ *    progress_handler([n]) { ... }
+ *    progress_handler([n,] Class.new { def call; end }.new)
+ *
+ * Register a progress handler with this database instance.
+ * This handler will be invoked periodically during a long-running query or operation.
+ * If the handler returns +false+, the operation will be interrupted; otherwise, it continues.
+ * The parameter 'n' specifies the number of SQLite virtual machine instructions between invocations.
+ * If 'n' is not provided, the default value is 1.
+ */
+static VALUE
+progress_handler(int argc, VALUE *argv, VALUE self)
+{
+  sqlite3RubyPtr ctx;
+  VALUE block, n_value;
+
+  TypedData_Get_Struct(self, sqlite3Ruby, &database_type, ctx);
+  REQUIRE_OPEN_DB(ctx);
+
+  rb_scan_args(argc, argv, "02", &n_value, &block);
+
+  int n = NIL_P(n_value) ? 1 : NUM2INT(n_value);
+  if(NIL_P(block) && rb_block_given_p()) block = rb_block_proc();
+
+  rb_iv_set(self, "@progress_handler", block);
+
+  sqlite3_progress_handler(
+      ctx->db, n, NIL_P(block) ? NULL : rb_sqlite3_progress_handler, (void *)self);
+
+  return self;
+}
+
+
 /* call-seq: last_insert_row_id
  *
  * Obtains the unique row ID of the last row to be inserted by this Database
@@ -866,6 +911,9 @@ init_sqlite3_database(void)
     rb_define_method(cSqlite3Database, "changes", changes, 0);
     rb_define_method(cSqlite3Database, "authorizer=", set_authorizer, 1);
     rb_define_method(cSqlite3Database, "busy_handler", busy_handler, -1);
+#ifndef SQLITE_OMIT_PROGRESS_CALLBACK
+    rb_define_method(cSqlite3Database, "progress_handler", progress_handler, -1);
+#endif
     rb_define_method(cSqlite3Database, "busy_timeout=", set_busy_timeout, 1);
     rb_define_method(cSqlite3Database, "extended_result_codes=", set_extended_result_codes, 1);
     rb_define_method(cSqlite3Database, "transaction_active?", transaction_active_p, 0);

lib/sqlite3/database.rb

@@ -124,6 +124,7 @@ def initialize file, options = {}, zvfs = nil
       @authorizer       = nil
       @encoding         = nil
       @busy_handler     = nil
+      @progress_handler = nil
       @collations       = {}
       @functions        = {}
       @results_as_hash  = options[:results_as_hash]

test/test_integration.rb

@@ -470,7 +470,7 @@ def test_transaction_active
 
   def test_transaction_implicit_rollback
     assert !@db.transaction_active?
-    @db.transaction 
+    @db.transaction
     @db.execute('create table bar (x CHECK(1 = 0))')
     assert @db.transaction_active?
     assert_raises( SQLite3::ConstraintException ) do
@@ -504,4 +504,56 @@ def test_bind_array_parameter
       [ 1, "foo" ] )
     assert_equal "foo", result
   end
+
+  def test_progress_handler_used
+    progress_calls = []
+    @db.progress_handler do
+      progress_calls << nil
+      true
+    end
+    @db.execute "create table test1(a, b)"
+
+    assert_operator 1, :<, progress_calls.size
+  end
+
+  def test_progress_handler_opcode_arg
+    progress_calls = []
+    handler = Proc.new do
+      progress_calls << nil
+      true
+    end
+    @db.progress_handler(1, handler)
+    @db.execute "create table test1(a, b)"
+    first_count = progress_calls.size
+
+    progress_calls = []
+    @db.progress_handler(10, handler)
+    @db.execute "create table test2(a, b)"
+    second_count = progress_calls.size
+
+    assert_operator first_count, :>=, second_count
+  end
+
+  def test_progress_handler_interrupts_operation
+    @db.progress_handler do
+      false
+    end
+
+    assert_raises(SQLite3::InterruptException) do
+      @db.execute "create table test1(a, b)"
+    end
+  end
+
+  def test_clear_handler
+    progress_calls = []
+    @db.progress_handler do
+      progress_calls << nil
+      true
+    end
+    @db.progress_handler(nil)
+
+    @db.execute "create table test1(a, b)"
+
+    assert_equal 0, progress_calls.size
+  end
 end