Statements attached to a discarded db raise an exception

sparklemotion · Sep 18, 2024 · c213875 · c213875
1 parent 1cabd3f
commit c213875
Show file tree

Hide file tree

Showing 3 changed files with 119 additions and 31 deletions.
diff --git a/ext/sqlite3/database.c b/ext/sqlite3/database.c
@@ -12,6 +12,44 @@
 
 VALUE cSqlite3Database;
 
+static void
+discard_db(sqlite3RubyPtr ctx)
+{
+    sqlite3_file *sfile;
+    int status;
+
+    // release as much heap memory as possible by deallocating non-essential memory
+    // allocations held by the database library. Memory used to cache database pages to
+    // improve performance is an example of non-essential memory.
+    // on my development machine, this reduces the lost memory from 152k to 69k.
+    sqlite3_db_release_memory(ctx->db);
+
+    // release file descriptors
+#ifdef HAVE_SQLITE3_DB_NAME
+    const char *db_name;
+    int j_db = 0;
+    while ((db_name = sqlite3_db_name(ctx->db, j_db)) != NULL) {
+        status = sqlite3_file_control(ctx->db, db_name, SQLITE_FCNTL_FILE_POINTER, &sfile);
+        if (status == 0 && sfile->pMethods != NULL) {
+            sfile->pMethods->xClose(sfile);
+        }
+        j_db++;
+    }
+#else
+    status = sqlite3_file_control(ctx->db, NULL, SQLITE_FCNTL_FILE_POINTER, &sfile);
+    if (status == 0 && sfile->pMethods != NULL) {
+        sfile->pMethods->xClose(sfile);
+    }
+#endif
+
+    status = sqlite3_file_control(ctx->db, NULL, SQLITE_FCNTL_JOURNAL_POINTER, &sfile);
+    if (status == 0 && sfile->pMethods != NULL) {
+        sfile->pMethods->xClose(sfile);
+    }
+
+    ctx->db = NULL;
+}
+
 static void
 close_or_discard_db(sqlite3RubyPtr ctx)
 {
@@ -21,41 +59,12 @@ close_or_discard_db(sqlite3RubyPtr ctx)
         if (isReadonly || ctx->owner == getpid()) {
             // Ordinary close.
             sqlite3_close_v2(ctx->db);
+            ctx->db = NULL;
         } else {
             // This is an open connection carried across a fork().
             // "Discard" it. See adr/2024-09-fork-safety.md
-            sqlite3_file *sfile;
-            int status;
-
-            // release as much heap memory as possible by deallocating non-essential memory
-            // allocations held by the database library. Memory used to cache database pages to
-            // improve performance is an example of non-essential memory.
-            sqlite3_db_release_memory(ctx->db);
-
-            // release file descriptors
-#ifdef HAVE_SQLITE3_DB_NAME
-            const char *db_name;
-            int j_db = 0;
-            while ((db_name = sqlite3_db_name(ctx->db, j_db)) != NULL) {
-                status = sqlite3_file_control(ctx->db, db_name, SQLITE_FCNTL_FILE_POINTER, &sfile);
-                if (status == 0 && sfile->pMethods != NULL) {
-                    sfile->pMethods->xClose(sfile);
-                }
-                j_db++;
-            }
-#else
-            status = sqlite3_file_control(ctx->db, NULL, SQLITE_FCNTL_FILE_POINTER, &sfile);
-            if (status == 0 && sfile->pMethods != NULL) {
-                sfile->pMethods->xClose(sfile);
-            }
-#endif
-
-            status = sqlite3_file_control(ctx->db, NULL, SQLITE_FCNTL_JOURNAL_POINTER, &sfile);
-            if (status == 0 && sfile->pMethods != NULL) {
-                sfile->pMethods->xClose(sfile);
-            }
+            discard_db(ctx);
         }
-        ctx->db = NULL;
     }
 }
 
@@ -189,6 +198,20 @@ sqlite3_rb_close(VALUE self)
     return self;
 }
 
+/* private method, primarily for testing */
+static VALUE
+sqlite3_rb_discard(VALUE self)
+{
+    sqlite3RubyPtr ctx;
+    TypedData_Get_Struct(self, sqlite3Ruby, &database_type, ctx);
+
+    discard_db(ctx);
+
+    rb_iv_set(self, "-aggregators", Qnil);
+
+    return self;
+}
+
 /* call-seq: db.closed?
  *
  * Returns +true+ if this database instance has been closed (see #close).
@@ -942,6 +965,7 @@ init_sqlite3_database(void)
     rb_define_private_method(cSqlite3Database, "open16", rb_sqlite3_open16, 1);
     rb_define_method(cSqlite3Database, "collation", collation, 2);
     rb_define_method(cSqlite3Database, "close", sqlite3_rb_close, 0);
+    rb_define_private_method(cSqlite3Database, "discard", sqlite3_rb_discard, 0);
     rb_define_method(cSqlite3Database, "closed?", closed_p, 0);
     rb_define_method(cSqlite3Database, "total_changes", total_changes, 0);
     rb_define_method(cSqlite3Database, "trace", trace, -1);

diff --git a/ext/sqlite3/statement.c b/ext/sqlite3/statement.c
@@ -4,6 +4,20 @@
   if(!_ctxt->st) \
     rb_raise(rb_path2class("SQLite3::Exception"), "cannot use a closed statement");
 
+static void
+require_open_db(VALUE stmt_rb)
+{
+    VALUE closed_p = rb_funcall(
+                         rb_iv_get(stmt_rb, "@connection"),
+                         rb_intern("closed?"), 0);
+
+    if (RTEST(closed_p)) {
+        rb_raise(rb_path2class("SQLite3::Exception"),
+                 "cannot use a statement associated with a closed database");
+    }
+}
+
+
 VALUE cSqlite3Statement;
 
 static void
@@ -121,6 +135,7 @@ step(VALUE self)
 
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
 
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     if (ctx->done_p) { return Qnil; }
@@ -216,6 +231,8 @@ bind_param(VALUE self, VALUE key, VALUE value)
     int index;
 
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     switch (TYPE(key)) {
@@ -308,6 +325,8 @@ reset_bang(VALUE self)
     sqlite3StmtRubyPtr ctx;
 
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     sqlite3_reset(ctx->st);
@@ -328,6 +347,8 @@ clear_bindings_bang(VALUE self)
     sqlite3StmtRubyPtr ctx;
 
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     sqlite3_clear_bindings(ctx->st);
@@ -360,6 +381,8 @@ column_count(VALUE self)
 {
     sqlite3StmtRubyPtr ctx;
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     return INT2NUM(sqlite3_column_count(ctx->st));
@@ -391,6 +414,8 @@ column_name(VALUE self, VALUE index)
     const char *name;
 
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     name = sqlite3_column_name(ctx->st, (int)NUM2INT(index));
@@ -414,6 +439,8 @@ column_decltype(VALUE self, VALUE index)
     const char *name;
 
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     name = sqlite3_column_decltype(ctx->st, (int)NUM2INT(index));
@@ -431,6 +458,8 @@ bind_parameter_count(VALUE self)
 {
     sqlite3StmtRubyPtr ctx;
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     return INT2NUM(sqlite3_bind_parameter_count(ctx->st));
@@ -538,7 +567,10 @@ stats_as_hash(VALUE self)
 {
     sqlite3StmtRubyPtr ctx;
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
+
     VALUE arg = rb_hash_new();
 
     stmt_stat_internal(arg, ctx->st);
@@ -554,6 +586,8 @@ stat_for(VALUE self, VALUE key)
 {
     sqlite3StmtRubyPtr ctx;
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     if (SYMBOL_P(key)) {
@@ -574,6 +608,8 @@ memused(VALUE self)
 {
     sqlite3StmtRubyPtr ctx;
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     return INT2NUM(sqlite3_stmt_status(ctx->st, SQLITE_STMTSTATUS_MEMUSED, 0));
@@ -591,6 +627,8 @@ database_name(VALUE self, VALUE index)
 {
     sqlite3StmtRubyPtr ctx;
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     return SQLITE3_UTF8_STR_NEW2(
@@ -608,6 +646,8 @@ get_sql(VALUE self)
 {
     sqlite3StmtRubyPtr ctx;
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     return rb_obj_freeze(SQLITE3_UTF8_STR_NEW2(sqlite3_sql(ctx->st)));
@@ -626,6 +666,8 @@ get_expanded_sql(VALUE self)
     VALUE rb_expanded_sql;
 
     TypedData_Get_Struct(self, sqlite3StmtRuby, &statement_type, ctx);
+
+    require_open_db(self);
     REQUIRE_OPEN_STMT(ctx);
 
     expanded_sql = sqlite3_expanded_sql(ctx->st);

diff --git a/test/test_database.rb b/test/test_database.rb
@@ -869,5 +869,27 @@ def test_close_does_not_discard_readonly_connections
         FileUtils.rm_f("test.db")
       end
     end
+
+    def test_a_discarded_connection_with_statements
+      skip("discard leaks memory") if i_am_running_in_valgrind
+
+      begin
+        db = SQLite3::Database.new("test.db")
+        db.execute("create table foo (bar int)")
+        db.execute("insert into foo values (1)")
+        stmt = db.prepare("select * from foo")
+
+        db.send(:discard)
+
+        e = assert_raises(SQLite3::Exception) { stmt.execute }
+        assert_match(/cannot use a statement associated with a closed database/, e.message)
+
+        assert_nothing_raised { stmt.close }
+        assert_predicate(stmt, :closed?)
+      ensure
+        db.close
+        FileUtils.rm_f("test.db")
+      end
+    end
   end
 end