From b390914bca283068d2b61f34fb0e51a92ea75d66 Mon Sep 17 00:00:00 2001
From: Marco Manino <mrcmnn@gmail.com>
Date: Thu, 25 Jul 2024 15:45:23 +0200
Subject: [PATCH 1/2] Cleanup prepared statement during node close

---
 src/gateway.c | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/gateway.c b/src/gateway.c
index 0ba1f60f4..845e992ed 100644
--- a/src/gateway.c
+++ b/src/gateway.c
@@ -76,6 +76,15 @@ void gateway__leader_close(struct gateway *g, int reason)
 			 */
 			sqlite3_finalize(g->req->stmt);
 			g->req = NULL;
+		} else if (g->req->type == DQLITE_REQUEST_QUERY) {
+			/* In case the statement is a prepared one, it
+			 * will be finalized by the stmt__registry_close
+			 * call below. Nevertheless, we must signal that
+			 * the request is not in place anymore so that any
+			 * callback which is already in the queue will not
+			 * attempt to execute a finalized statement.
+			 */
+			g->req = NULL;
 		}
 	}
 	stmt__registry_close(&g->stmts);

From e6ac700c2191de7741bf3642e1e6a53a0005aae7 Mon Sep 17 00:00:00 2001
From: Marco Manino <mrcmnn@gmail.com>
Date: Fri, 26 Jul 2024 11:04:16 +0200
Subject: [PATCH 2/2] Add testing

---
 test/unit/test_gateway.c | 50 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/test/unit/test_gateway.c b/test/unit/test_gateway.c
index 425ee57f4..d0e1fb7c6 100644
--- a/test/unit/test_gateway.c
+++ b/test/unit/test_gateway.c
@@ -1631,6 +1631,56 @@ TEST_CASE(query, manyParams, NULL)
 	return MUNIT_OK;
 }
 
+
+/* Successfully query that yields a large number of rows that need to be split
+ * into several reponses. */
+TEST_CASE(query, close_while_in_flight, NULL)
+{
+	struct query_fixture *f = data;
+	unsigned i;
+	uint64_t stmt_id;
+	uint64_t n;
+	const char *column;
+	struct value value;
+	bool finished;
+	(void)params;
+	EXEC("BEGIN");
+
+	/* 16 = 8B header + 8B value (int) */
+	unsigned n_rows_buffer = max_rows_buffer(16);
+	/* Insert 1 less than 2 response buffers worth of rows, otherwise we
+	 * need 3 responses, of which the last one contains no rows. */
+	for (i = 0; i < ((2 * n_rows_buffer) - 1); i++) {
+		EXEC("INSERT INTO test(n) VALUES(123)");
+	}
+	EXEC("COMMIT");
+
+	PREPARE("SELECT n FROM test");
+	f->request.db_id = 0;
+	f->request.stmt_id = stmt_id;
+	ENCODE(&f->request, query);
+	HANDLE(QUERY);
+	ASSERT_CALLBACK(0, ROWS);
+
+	uint64__decode(f->cursor, &n);
+	munit_assert_int(n, ==, 1);
+	text__decode(f->cursor, &column);
+	munit_assert_string_equal(column, "n");
+
+	/* First response contains max amount of rows */
+	for (i = 0; i < n_rows_buffer; i++) {
+		DECODE_ROW(1, &value);
+		munit_assert_int(value.type, ==, SQLITE_INTEGER);
+		munit_assert_int(value.integer, ==, 123);
+	}
+
+	/* Simulate a gateway close */
+	gateway__close(f->gateway);
+	gateway__resume(f->gateway, &finished);
+
+	return MUNIT_OK;
+}
+
 /******************************************************************************
  *
  * finalize