diff --git a/bin/test-gem-file-contents b/bin/test-gem-file-contents index cbe98c24..f4ad822e 100755 --- a/bin/test-gem-file-contents +++ b/bin/test-gem-file-contents @@ -112,7 +112,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 @@ -154,7 +154,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 diff --git a/ext/sqlite3/database.c b/ext/sqlite3/database.c index 36b75787..9f6f134b 100644 --- a/ext/sqlite3/database.c +++ b/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, ¤tTime); + + if (!timespecisset(&ctx->stmt_deadline)) { + // Set stmt_deadline if not already set + ctx->stmt_deadline = currentTime; + } else if (timespecafter(¤tTime, &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 @@ -875,6 +913,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); diff --git a/ext/sqlite3/database.h b/ext/sqlite3/database.h index 56833020..3123f4fe 100644 --- a/ext/sqlite3/database.h +++ b/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; diff --git a/ext/sqlite3/sqlite3_ruby.h b/ext/sqlite3/sqlite3_ruby.h index bcf53e63..088d3cd5 100644 --- a/ext/sqlite3/sqlite3_ruby.h +++ b/ext/sqlite3/sqlite3_ruby.h @@ -41,6 +41,7 @@ extern VALUE cSqlite3Blob; #include #include #include +#include int bignum_to_int64(VALUE big, sqlite3_int64 *result); diff --git a/ext/sqlite3/statement.c b/ext/sqlite3/statement.c index 9fb7e4a1..d049dec1 100644 --- a/ext/sqlite3/statement.c +++ b/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); } diff --git a/ext/sqlite3/timespec.h b/ext/sqlite3/timespec.h new file mode 100644 index 00000000..322fe758 --- /dev/null +++ b/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)) diff --git a/lib/sqlite3/database.rb b/lib/sqlite3/database.rb index e8977896..5b97226a 100644 --- a/lib/sqlite3/database.rb +++ b/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] diff --git a/sqlite3.gemspec b/sqlite3.gemspec index d0746454..5555e6c5 100644 --- a/sqlite3.gemspec +++ b/sqlite3.gemspec @@ -61,6 +61,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", diff --git a/test/test_integration_statement.rb b/test/test_integration_statement.rb index e1f86a91..2e180acd 100644 --- a/test/test_integration_statement.rb +++ b/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