Add statement_timeout= method

ext/sqlite3/database.c

@@ -259,6 +259,45 @@ 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 (ctx->stmt_deadline == 0) {
+        ctx->stmt_deadline = currentTime.tv_sec * 1e9 + currentTime.tv_nsec + (long long)ctx->stmt_timeout * 1e6;
+    } else {
+        long long timeDiff = ctx->stmt_deadline - (currentTime.tv_sec * 1e9 + currentTime.tv_nsec);
+
+        if (timeDiff < 0) { 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
@@ -875,6 +914,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;
+    long long stmt_deadline;
 };
 
 typedef struct _sqlite3Ruby sqlite3Ruby;

ext/sqlite3/statement.c

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

lib/sqlite3/database.rb

@@ -119,6 +119,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_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