Merge pull request #463 from fractaledmind/statement-timeout

Add `statement_timeout=` method

bin/test-gem-file-contents

@@ -123,7 +123,7 @@ describe File.basename(gemfile) do
 
       it "contains extension C and header files" do
         assert_equal(6, gemfile_contents.count { |f| File.fnmatch?("ext/**/*.c", f) })
-        assert_equal(6, gemfile_contents.count { |f| File.fnmatch?("ext/**/*.h", f) })
+        assert_equal(7, gemfile_contents.count { |f| File.fnmatch?("ext/**/*.h", f) })
       end
 
       it "includes C files in extra_rdoc_files" do
@@ -165,7 +165,7 @@ describe File.basename(gemfile) do
 
       it "contains extension C and header files" do
         assert_equal(6, gemfile_contents.count { |f| File.fnmatch?("ext/**/*.c", f) })
-        assert_equal(6, gemfile_contents.count { |f| File.fnmatch?("ext/**/*.h", f) })
+        assert_equal(7, gemfile_contents.count { |f| File.fnmatch?("ext/**/*.h", f) })
       end
 
       it "includes C files in extra_rdoc_files" do

ext/sqlite3/database.c

@@ -259,6 +259,44 @@ busy_handler(int argc, VALUE *argv, VALUE self)
     return self;
 }
 
+static int
+rb_sqlite3_statement_timeout(void *context)
+{
+    sqlite3RubyPtr ctx = (sqlite3RubyPtr)context;
+    struct timespec currentTime;
+    clock_gettime(CLOCK_MONOTONIC, &currentTime);
+
+    if (!timespecisset(&ctx->stmt_deadline)) {
+        // Set stmt_deadline if not already set
+        ctx->stmt_deadline = currentTime;
+    } else if (timespecafter(&currentTime, &ctx->stmt_deadline)) {
+        return 1;
+    }
+
+    return 0;
+}
+
+/* call-seq: db.statement_timeout = ms
+ *
+ * Indicates that if a query lasts longer than the indicated number of
+ * milliseconds, SQLite should interrupt that query and return an error.
+ * By default, SQLite does not interrupt queries. To restore the default
+ * behavior, send 0 as the +ms+ parameter.
+ */
+static VALUE
+set_statement_timeout(VALUE self, VALUE milliseconds)
+{
+    sqlite3RubyPtr ctx;
+    TypedData_Get_Struct(self, sqlite3Ruby, &database_type, ctx);
+
+    ctx->stmt_timeout = NUM2INT(milliseconds);
+    int n = NUM2INT(milliseconds) == 0 ? -1 : 1000;
+
+    sqlite3_progress_handler(ctx->db, n, rb_sqlite3_statement_timeout, (void *)ctx);
+
+    return self;
+}
+
 /* call-seq: last_insert_row_id
  *
  * Obtains the unique row ID of the last row to be inserted by this Database
@@ -869,6 +907,9 @@ init_sqlite3_database(void)
     rb_define_method(cSqlite3Database, "authorizer=", set_authorizer, 1);
     rb_define_method(cSqlite3Database, "busy_handler", busy_handler, -1);
     rb_define_method(cSqlite3Database, "busy_timeout=", set_busy_timeout, 1);
+#ifndef SQLITE_OMIT_PROGRESS_CALLBACK
+    rb_define_method(cSqlite3Database, "statement_timeout=", set_statement_timeout, 1);
+#endif
     rb_define_method(cSqlite3Database, "extended_result_codes=", set_extended_result_codes, 1);
     rb_define_method(cSqlite3Database, "transaction_active?", transaction_active_p, 0);
     rb_define_private_method(cSqlite3Database, "exec_batch", exec_batch, 2);

ext/sqlite3/database.h

@@ -6,6 +6,8 @@
 struct _sqlite3Ruby {
     sqlite3 *db;
     VALUE busy_handler;
+    int stmt_timeout;
+    struct timespec stmt_deadline;
 };
 
 typedef struct _sqlite3Ruby sqlite3Ruby;

ext/sqlite3/sqlite3_ruby.h

@@ -41,6 +41,7 @@ extern VALUE cSqlite3Blob;
 #include <statement.h>
 #include <exception.h>
 #include <backup.h>
+#include <timespec.h>
 
 int bignum_to_int64(VALUE big, sqlite3_int64 *result);
 

ext/sqlite3/statement.c

@@ -70,6 +70,7 @@ prepare(VALUE self, VALUE db, VALUE sql)
              );
 
     CHECK(db_ctx->db, status);
+    timespecclear(&db_ctx->stmt_deadline);
 
     return rb_str_new2(tail);
 }

ext/sqlite3/timespec.h

@@ -0,0 +1,20 @@
+#define	timespecclear(tsp)		(tsp)->tv_sec = (tsp)->tv_nsec = 0
+#define	timespecisset(tsp)		((tsp)->tv_sec || (tsp)->tv_nsec)
+#define	timespecisvalid(tsp)						\
+	((tsp)->tv_nsec >= 0 && (tsp)->tv_nsec < 1000000000L)
+#define	timespeccmp(tsp, usp, cmp)					\
+	(((tsp)->tv_sec == (usp)->tv_sec) ?				\
+	    ((tsp)->tv_nsec cmp (usp)->tv_nsec) :			\
+	    ((tsp)->tv_sec cmp (usp)->tv_sec))
+#define	timespecsub(tsp, usp, vsp)					\
+	do {								\
+		(vsp)->tv_sec = (tsp)->tv_sec - (usp)->tv_sec;		\
+		(vsp)->tv_nsec = (tsp)->tv_nsec - (usp)->tv_nsec;	\
+		if ((vsp)->tv_nsec < 0) {				\
+			(vsp)->tv_sec--;				\
+			(vsp)->tv_nsec += 1000000000L;			\
+		}							\
+	} while (0)
+#define timespecafter(tsp, usp)						\
+	(((tsp)->tv_sec > (usp)->tv_sec) ||				\
+	 ((tsp)->tv_sec == (usp)->tv_sec && (tsp)->tv_nsec > (usp)->tv_nsec))

lib/sqlite3/database.rb

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

sqlite3.gemspec

@@ -58,6 +58,7 @@ Gem::Specification.new do |s|
     "ext/sqlite3/sqlite3_ruby.h",
     "ext/sqlite3/statement.c",
     "ext/sqlite3/statement.h",
+    "ext/sqlite3/timespec.h",
     "lib/sqlite3.rb",
     "lib/sqlite3/constants.rb",
     "lib/sqlite3/database.rb",

test/test_integration_statement.rb

@@ -191,4 +191,20 @@ def test_committing_tx_with_statement_active
     end
     assert called
   end
+
+  def test_long_running_statements_get_interrupted_when_statement_timeout_set
+    @db.statement_timeout = 10
+    assert_raises(SQLite3::InterruptException) do
+      @db.execute <<~SQL
+        WITH RECURSIVE r(i) AS (
+          VALUES(0)
+          UNION ALL
+          SELECT i FROM r
+          LIMIT 100000
+        )
+        SELECT i FROM r ORDER BY i LIMIT 1;
+      SQL
+    end
+    @db.statement_timeout = 0
+  end
 end