From 53fe3d1ab5e97e62b87e80714c075da5e4ef25b8 Mon Sep 17 00:00:00 2001
From: Anshul Khandelwal <12948312+k-anshul@users.noreply.github.com>
Date: Thu, 11 Jan 2024 18:25:19 -0600
Subject: [PATCH] interrupt query on ctx cancel/timeout (#143)

* interrupt query on ctx cancel/timeout

* adding unit test

* adding unit test - reduce timeout
---
 duckdb_test.go | 16 ++++++++++++++++
 statement.go   |  2 ++
 2 files changed, 18 insertions(+)

diff --git a/duckdb_test.go b/duckdb_test.go
index bd278d16..d9303973 100644
--- a/duckdb_test.go
+++ b/duckdb_test.go
@@ -1100,6 +1100,22 @@ func TestParquetExtension(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestQueryTimeout(t *testing.T) {
+	db := openDB(t)
+	defer db.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*250)
+	defer cancel()
+
+	now := time.Now()
+	_, err := db.ExecContext(ctx, "CREATE TABLE test AS SELECT * FROM range(10000000) t1, range(1000000) t2;")
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+
+	// a very defensive time check, but should be good enough
+	// the query takes much longer than 10 seconds
+	require.Less(t, time.Since(now), 10*time.Second)
+}
+
 func openDB(t *testing.T) *sql.DB {
 	db, err := sql.Open("duckdb", "")
 	require.NoError(t, err)
diff --git a/statement.go b/statement.go
index ad777ec2..23e2dad4 100644
--- a/statement.go
+++ b/statement.go
@@ -207,6 +207,8 @@ func (s *stmt) execute(ctx context.Context, args []driver.NamedValue) (*C.duckdb
 		select {
 		// if context is cancelled or deadline exceeded, don't execute further
 		case <-ctx.Done():
+			// also need to interrupt to cancel the query
+			C.duckdb_interrupt(*s.c.con)
 			return nil, ctx.Err()
 		default:
 			// continue