fix: use vu Context() instead of background Context

releases/v1.0.4.md

@@ -0,0 +1,10 @@
+xk6-sql `v1.0.4` is here 🎉!
+
+This release includes:
+
+## Bugfixes
+
+- [Use VU Context()](https://github.com/grafana/xk6-sql/issues/124): VU Context() is now used in `query()` and `exec()` functions instead of background context. Using background context is a potential problem if SQL operations are still running after the VU context is invalidated.
+
+
+

sql/module.go

@@ -2,6 +2,7 @@
 package sql
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -24,20 +25,23 @@ type rootModule struct{}
 
 // NewModuleInstance implements the modules.Module interface to return
 // a new instance for each VU.
-func (*rootModule) NewModuleInstance(_ modules.VU) modules.Instance {
+func (*rootModule) NewModuleInstance(vu modules.VU) modules.Instance {
 	instance := &module{}
 
 	instance.exports.Default = instance
 	instance.exports.Named = map[string]interface{}{
 		"open": instance.Open,
 	}
 
+	instance.vu = vu
+
 	return instance
 }
 
 // module represents an instance of the JavaScript module for every VU.
 type module struct {
 	exports modules.Exports
+	vu      modules.VU
 }
 
 // Exports is representation of ESM exports of a module.
@@ -97,17 +101,18 @@ func (mod *module) Open(driverID sobek.Value, connectionString string, opts *opt
 		return nil, err
 	}
 
-	return &Database{db: db}, nil
+	return &Database{db: db, ctx: mod.vu.Context}, nil
 }
 
 // Database is a database handle representing a pool of zero or more underlying connections.
 type Database struct {
-	db *sql.DB
+	db  *sql.DB
+	ctx func() context.Context
 }
 
 // Query executes a query that returns rows, typically a SELECT.
 func (dbase *Database) Query(query string, args ...interface{}) ([]KeyValue, error) {
-	rows, err := dbase.db.Query(query, args...)
+	rows, err := dbase.db.QueryContext(dbase.ctx(), query, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -152,7 +157,7 @@ func (dbase *Database) Query(query string, args ...interface{}) ([]KeyValue, err
 
 // Exec a query without returning any rows.
 func (dbase *Database) Exec(query string, args ...interface{}) (sql.Result, error) {
-	return dbase.db.Exec(query, args...)
+	return dbase.db.ExecContext(dbase.ctx(), query, args...)
 }
 
 // Close the database and prevents new queries from starting.

sql/sql_internal_test.go

@@ -5,10 +5,13 @@ import (
 
 	"github.com/grafana/sobek"
 	"github.com/stretchr/testify/require"
+	"go.k6.io/k6/js/modulestest"
 )
 
 func TestOpen(t *testing.T) { //nolint: paralleltest
-	mod, ok := New().NewModuleInstance(nil).(*module)
+	rt := modulestest.NewRuntime(t)
+
+	mod, ok := New().NewModuleInstance(rt.VU).(*module)
 
 	require.True(t, ok)