Merge pull request #128 from grafana/127-timeout-support-in-exec-and-…

…query-functions

feat: timeout support

index.d.ts

@@ -220,6 +220,48 @@ export interface Database {
    * ```
    */
   exec(query: string, ...args: any[]): Result;
+  /**
+   * Execute a query (with a timeout) without returning any rows.
+   * The timeout parameter is a duration string, a possibly signed sequence of decimal numbers,
+   * each with optional fraction and a unit suffix, such as "300ms", "-1.5h" or "2h45m".
+   * Valid time units are "ns", "us" (or "µs"), "ms", "s", "m", "h".
+   * @param timeout the query timeout as a duration string
+   * @param query the query to execute
+   * @param args placeholder parameters in the query
+   * @returns summary of the executed SQL commands
+   * @example
+   *  ```ts file=examples/example.js
+   *  import sql from "k6/x/sql";
+   *
+   *  // the actual database driver should be used instead of ramsql
+   *  import driver from "k6/x/sql/driver/ramsql";
+   *
+   *  const db = sql.open(driver, "roster_db");
+   *
+   *  export function setup() {
+   *    db.exec(`
+   *      CREATE TABLE IF NOT EXISTS roster
+   *        (
+   *          id INTEGER PRIMARY KEY AUTOINCREMENT,
+   *          given_name VARCHAR NOT NULL,
+   *          family_name VARCHAR NOT NULL
+   *        );
+   *    `);
+   *
+   *    let result = db.execWithTimeout("10s", `
+   *      INSERT INTO roster
+   *        (given_name, family_name)
+   *      VALUES
+   *        ('Peter', 'Pan'),
+   *        ('Wendy', 'Darling'),
+   *        ('Tinker', 'Bell'),
+   *        ('James', 'Hook');
+   *    `);
+   *    console.log(`${result.rowsAffected()} rows inserted`);
+   *  }
+   * ```
+   */
+  execWithTimeout(timeout: string, query: string, ...args: any[]): Result;
   /**
    * Query executes a query that returns rows, typically a SELECT.
    * @param query the query to execute
@@ -243,6 +285,33 @@ export interface Database {
    * ```
    */
   query(query: string, ...args: any[]): Row[];
+  /**
+   * Query executes a query (with a timeout) that returns rows, typically a SELECT.
+   * The timeout parameter is a duration string, a possibly signed sequence of decimal numbers,
+   * each with optional fraction and a unit suffix, such as "300ms", "-1.5h" or "2h45m".
+   * Valid time units are "ns", "us" (or "µs"), "ms", "s", "m", "h".
+   * @param timeout the query timeout as a duration string
+   * @param query the query to execute
+   * @param args placeholder parameters in the query
+   * @returns rows of the query result
+   * @example
+   *  ```ts file=examples/example.js
+   *  import sql from "k6/x/sql";
+   *
+   *  // the actual database driver should be used instead of ramsql
+   *  import driver from "k6/x/sql/driver/ramsql";
+   *
+   *  const db = sql.open(driver, "roster_db");
+   *
+   *  export default function () {
+   *    let rows = db.queryWithTimeout("10s", "SELECT * FROM roster WHERE given_name = $1;", "Peter");
+   *    for (const row of results) {
+   *      console.log(`${row.family_name}, ${row.given_name}`);
+   *    }
+   *  }
+   * ```
+   */
+  queryWithTimeout(timeout: string, query: string, ...args: any[]): Row[];
 }
 
 /**

releases/v1.0.4.md

@@ -2,6 +2,10 @@ xk6-sql `v1.0.4` is here 🎉!
 
 This release includes:
 
+## New features
+
+- [Timeout support in `execWithTimeout()` and `queryWithTimeout()` functions](https://github.com/grafana/xk6-sql/issues/127): The `exec()` and `query()` functions wait for an unlimited amount of time for the result. This can lead to the test stalling, for example, in the event of a network problem. The API has been extended with timeout-handling counterparts of these functions: `execWithTimeout()` and `queryWithTimeout()`. The first parameter is the timeout. The timeout parameter is a duration string, a possibly signed sequence of decimal numbers, each with optional fraction and a unit suffix, such as `300ms`, `-1.5h` or `2h45m`. Valid time units are `ns`, `us` (or `µs`), `ms`, `s`, `m`, `h`.
+
 ## 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

@@ -6,6 +6,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/grafana/sobek"
 	"go.k6.io/k6/js/modules"
@@ -110,9 +111,19 @@ type Database struct {
 	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.QueryContext(dbase.ctx(), query, args...)
+func (dbase *Database) parseTimeout(timeout string) (context.Context, context.CancelFunc, error) {
+	dur, err := time.ParseDuration(timeout)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	ctx, cancel := context.WithTimeout(dbase.ctx(), dur)
+
+	return ctx, cancel, nil
+}
+
+func (dbase *Database) query(ctx context.Context, query string, args ...interface{}) ([]KeyValue, error) {
+	rows, err := dbase.db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -155,9 +166,44 @@ func (dbase *Database) Query(query string, args ...interface{}) ([]KeyValue, err
 	return result, nil
 }
 
-// Exec a query without returning any rows.
+// Query executes a query that returns rows, typically a SELECT.
+func (dbase *Database) Query(query string, args ...interface{}) ([]KeyValue, error) {
+	return dbase.query(dbase.ctx(), query, args...)
+}
+
+// QueryWithTimeout executes a query (with a timeout) that returns rows, typically a SELECT.
+// The timeout can be specified as a duration string.
+func (dbase *Database) QueryWithTimeout(timeout string, query string, args ...interface{}) ([]KeyValue, error) {
+	ctx, cancel, err := dbase.parseTimeout(timeout)
+	if err != nil {
+		return nil, err
+	}
+
+	defer cancel()
+
+	return dbase.query(ctx, query, args...)
+}
+
+func (dbase *Database) exec(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
+	return dbase.db.ExecContext(ctx, query, args...)
+}
+
+// Exec executes a query without returning any rows.
 func (dbase *Database) Exec(query string, args ...interface{}) (sql.Result, error) {
-	return dbase.db.ExecContext(dbase.ctx(), query, args...)
+	return dbase.exec(dbase.ctx(), query, args...)
+}
+
+// ExecWithTimeout executes a query (with a timeout) without returning any rows.
+// The timeout can be specified as a duration string.
+func (dbase *Database) ExecWithTimeout(timeout string, query string, args ...interface{}) (sql.Result, error) {
+	ctx, cancel, err := dbase.parseTimeout(timeout)
+	if err != nil {
+		return nil, err
+	}
+
+	defer cancel()
+
+	return dbase.exec(ctx, query, args...)
 }
 
 // Close the database and prevents new queries from starting.

sql/sql_internal_test.go

@@ -1,6 +1,8 @@
 package sql
 
 import (
+	"context"
+	"runtime"
 	"testing"
 
 	"github.com/grafana/sobek"
@@ -23,6 +25,16 @@ func TestOpen(t *testing.T) { //nolint: paralleltest
 	require.NoError(t, err)
 	require.NotNil(t, db)
 
+	const expr = `CREATE TABLE address (id BIGSERIAL PRIMARY KEY, street TEXT, street_number INT);`
+
+	if runtime.GOOS != "windows" {
+		_, err = db.ExecWithTimeout("1ns", expr)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+
+		_, err = db.QueryWithTimeout("1ns", expr)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	}
+
 	_, err = mod.Open(sobek.New().ToValue("foo"), "testdb", nil) // not a Symbol
 
 	require.Error(t, err)

sql/testdata/script.js

@@ -11,7 +11,7 @@ if (all_rows.length != 5) {
   throw new Error("Expected all five rows to be returned; got " + all_rows.length);
 }
 
-let one_row = db.query("SELECT * FROM test_table WHERE name = $1;", "name-2");
+let one_row = db.queryWithTimeout("10s", "SELECT * FROM test_table WHERE name = $1;", "name-2");
 if (one_row.length != 1) {
   throw new Error("Expected single row to be returned; got " + one_row.length);
 }