add ctx arguments to most methods

8 files changed, 154 insertions, 132 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 948bb2e..b56ec45 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,11 +7,14 @@ SPDX-License-Identifier: Apache-2.0
 
 Changes:
 
+- All `Store` methods that run DB operations now take a `context.Context` argument.
+  The previous decision to avoid `ctx` arguments was based on benchmarking using `b.Context()`,
+  where it turns out that we just benchmarked those particular stacked contexts being inefficient.
 - `Store.Insert()` now takes its arguments by-pointer. This is probably slightly less efficient,
   but significantly safer because autogenerated field values cannot be disregarded by accident.
 - Add `Store.Update()`.
 - Removed support for SQL dialects that rely on LastInsertId() for ID columns.
-  Using a RETURNING clause to collect autogenerated field values is objectively better in every way,
+  Using a RETURNING clause to collect autogenerated field values is better in many ways,
   and has been supported by both MariaDB and SQLite for at least six years.
   In practice, this only drops support specifically for Oracle MySQL.
 
diff --git a/README.md b/README.md
index aa3fec8..0bc2e9e 100644
--- a/README.md
+++ b/README.md
@@ -28,12 +28,6 @@ The design goals, ordered by priority (most important comes first), are:
 - A minimal amount of CPU usage.
 - As few library dependencies as possible.
 
-As a surprising consequence, this set of priorities forced this library to eschew `context.Context` arguments.
-Early benchmarking showed that replacing `QueryRow` with `QueryRowContext` increased allocations by up to 50% and memory allocated by up to 100%.
-The author of this library is still a fan of `context.Context` for things like HTTP requests to external services, where unpredictable delays make a structured cancellation facility vital.
-But this library optimizes for blazing fast OLTP workloads (with maybe a few OLAP queries every once in a while), where not being able to back out of a running query is not that big of a deal because query runtimes should always be short anyway.
-If in doubt, the lack of `context.Context` arguments can be counteracted by setting timeouts on your DB transactions during OLAP workloads.
-
 Explicit non-goals include:
 
 - A fully featured API for query construction:
diff --git a/benchmark/benchmark_test.go b/benchmark/benchmark_test.go
index 98d5818..439ecaf 100644
--- a/benchmark/benchmark_test.go
+++ b/benchmark/benchmark_test.go
@@ -4,6 +4,7 @@
 package main_test
 
 import (
+	"context"
 	"crypto/sha256"
 	"database/sql"
 	"fmt"
@@ -20,6 +21,9 @@ import (
 	"gorm.io/gorm"
 )
 
+// Do not use b.Context() within benchmarks, or you will merely demonstrate that using a deep stack of Context objects is expensive.
+var noctx = context.Background()
+
 // This is not a real benchmark (obviously).
 // Its purpose is to be the first line that is printed, while having one of the longest names,
 // so that all other results are aligned with it and the table looks nice.
@@ -91,12 +95,12 @@ func BenchmarkSelectMany(b *testing.B) {
 			precomputedQuery := store.MustPrepareSelectQueryWhere(partialQuery)
 
 			selectWithOblast := func(b *testing.B) {
-				records := must.Return(store.Select(db, query))(b)
+				records := must.Return(store.Select(noctx, db, query))(b)
 				assert.Equal(b, len(records), batchSize)
 			}
 
 			selectWithOblastWhere := func(b *testing.B) {
-				records := must.Return(precomputedQuery.Select(db))(b)
+				records := must.Return(precomputedQuery.Select(noctx, db))(b)
 				assert.Equal(b, len(records), batchSize)
 			}
 
@@ -185,12 +189,12 @@ func BenchmarkSelectOne(b *testing.B) {
 	precomputedQuery := store.MustPrepareSelectQueryWhere(partialQuery)
 
 	selectWithOblast := func(b *testing.B) {
-		r := must.Return(store.SelectOne(db, query))(b)
+		r := must.Return(store.SelectOne(noctx, db, query))(b)
 		assert.Equal(b, r.ID, recordID)
 	}
 
 	selectWithOblastWhere := func(b *testing.B) {
-		r := must.Return(precomputedQuery.SelectOne(db))(b)
+		r := must.Return(precomputedQuery.SelectOne(noctx, db))(b)
 		assert.Equal(b, r.ID, recordID)
 	}
 
@@ -273,13 +277,13 @@ func BenchmarkInsertAndDelete(b *testing.B) {
 					records[idx] = OblastEntry{Message: "hello"}
 					recordsForInsert[idx] = &records[idx]
 				}
-				must.Succeed(b, store.Insert(db, recordsForInsert...))
+				must.Succeed(b, store.Insert(noctx, db, recordsForInsert...))
 				for _, r := range records {
 					if r.ID == 0 {
 						b.Errorf("ID was not filled!")
 					}
 				}
-				must.Succeed(b, store.Delete(db, records...))
+				must.Succeed(b, store.Delete(noctx, db, records...))
 			}
 
 			insertAndDeleteWithGorp := func(b *testing.B) {
@@ -429,7 +433,7 @@ func BenchmarkUpdate(b *testing.B) {
 				recordsForOblast[idx] = OblastEntry{Message: "hello"}
 				recordsForOblastForInsert[idx] = &recordsForOblast[idx]
 			}
-			must.Succeed(b, store.Insert(db, recordsForOblastForInsert...))
+			must.Succeed(b, store.Insert(noctx, db, recordsForOblastForInsert...))
 			recordsForGorp := make([]any, batchSize)
 			for idx, r := range recordsForOblast {
 				recordsForGorp[idx] = new(GorpEntry(r))
@@ -444,7 +448,7 @@ func BenchmarkUpdate(b *testing.B) {
 				for idx := range recordsForOblast {
 					recordsForOblast[idx].Message = message
 				}
-				must.Succeed(b, store.Update(db, recordsForOblast...))
+				must.Succeed(b, store.Update(noctx, db, recordsForOblast...))
 			}
 			updateWithGorp := func(b *testing.B, message string) {
 				for _, r := range recordsForGorp {
diff --git a/oblast.go b/oblast.go
index 059911b..ba43798 100644
--- a/oblast.go
+++ b/oblast.go
@@ -96,6 +96,7 @@
 package oblast // import "go.xyrillian.de/oblast"
 
 import (
+	"context"
 	"database/sql"
 	"database/sql/driver"
 	"fmt"
@@ -134,10 +135,10 @@ func StructTagKeyIs(key string) PlanOption {
 // Handle is an interface for functions providing direct DB access.
 // It covers methods provided by both *sql.DB and *sql.Tx, thus allowing functions using it to be used both within and outside of transactions.
 type Handle interface {
-	Exec(query string, args ...any) (sql.Result, error)
-	Prepare(query string) (*sql.Stmt, error)
-	Query(query string, args ...any) (*sql.Rows, error)
-	QueryRow(query string, args ...any) *sql.Row
+	ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
+	PrepareContext(ctx context.Context, query string) (*sql.Stmt, error)
+	QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
+	QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
 }
 
 // TODO: investigate if we can extend type Handle to cover types github.com/jackc/pgx.{Conn,Tx}
diff --git a/query.go b/query.go
index 73c2233..b9bc328 100644
--- a/query.go
+++ b/query.go
@@ -4,6 +4,7 @@
 package oblast
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"reflect"
@@ -29,7 +30,7 @@ type preparedStatement struct {
 }
 
 // prepare behaves like [Handle.Prepare].
-func prepare(db Handle, query, operation string, inputSize int) (preparedStatement, error) {
+func prepare(ctx context.Context, db Handle, query, operation string, inputSize int) (preparedStatement, error) {
 	if query == "" {
 		return preparedStatement{}, fmt.Errorf("cannot execute %s() because query could not be autogenerated", operation)
 	}
@@ -37,7 +38,7 @@ func prepare(db Handle, query, operation string, inputSize int) (preparedStateme
 	if inputSize < PrepareThreshold {
 		return preparedStatement{db, query, nil}, nil
 	}
-	stmt, err := db.Prepare(query)
+	stmt, err := db.PrepareContext(ctx, query)
 	if err != nil {
 		return preparedStatement{}, fmt.Errorf("during Prepare(): %w", err)
 	}
@@ -52,20 +53,20 @@ func (s preparedStatement) Close() error {
 	return s.stmt.Close()
 }
 
-// Exec behaves like [sql.Stmt.Exec].
-func (s preparedStatement) Exec(args ...any) (sql.Result, error) {
+// ExecContext behaves like [sql.Stmt.ExecContext].
+func (s preparedStatement) ExecContext(ctx context.Context, args ...any) (sql.Result, error) {
 	if s.stmt == nil {
-		return s.db.Exec(s.query, args...)
+		return s.db.ExecContext(ctx, s.query, args...)
 	}
-	return s.stmt.Exec(args...)
+	return s.stmt.ExecContext(ctx, args...)
 }
 
-// QueryRow behaves like [sql.Stmt.QueryRow].
-func (s preparedStatement) QueryRow(args ...any) *sql.Row {
+// QueryRow behaves like [sql.Stmt.QueryRowContext].
+func (s preparedStatement) QueryRowContext(ctx context.Context, args ...any) *sql.Row {
 	if s.stmt == nil {
-		return s.db.QueryRow(s.query, args...)
+		return s.db.QueryRowContext(ctx, s.query, args...)
 	}
-	return s.stmt.QueryRow(args...)
+	return s.stmt.QueryRowContext(ctx, args...)
 }
 
 // Insert executes an SQL INSERT statement for each of the provided records.
@@ -77,18 +78,18 @@ func (s preparedStatement) QueryRow(args ...any) *sql.Row {
 //
 // Returns an error if any of the `records` has a non-zero value in any column marked as `db:",auto"`.
 // Records that already exist in the database should be handled with [Store.Update] instead.
-func (s Store[R]) Insert(db Handle, records ...*R) error {
+func (s Store[R]) Insert(ctx context.Context, db Handle, records ...*R) error {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
-	stmt, err := prepare(db, s.plan.Insert.Query, "Insert", len(records))
+	stmt, err := prepare(ctx, db, s.plan.Insert.Query, "Insert", len(records))
 	if err != nil {
 		return err
 	}
-	return s.insertUsing(stmt, db, records)
+	return s.insertUsing(ctx, stmt, db, records)
 }
 
-func (s Store[R]) insertUsing(stmt preparedStatement, db Handle, records []*R) error {
+func (s Store[R]) insertUsing(ctx context.Context, stmt preparedStatement, db Handle, records []*R) error {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
@@ -101,7 +102,7 @@ func (s Store[R]) insertUsing(stmt preparedStatement, db Handle, records []*R) e
 
 	for idx, r := range records {
 		v := reflect.ValueOf(r).Elem()
-		err := insertRecord(v, idx, stmt, argumentIndexes, argumentSlots, scanIndexes, scanSlots)
+		err := insertRecord(ctx, v, idx, stmt, argumentIndexes, argumentSlots, scanIndexes, scanSlots)
 		if err != nil {
 			return newIOError(err, "Stmt.Close", stmt.Close())
 		}
@@ -110,7 +111,7 @@ func (s Store[R]) insertUsing(stmt preparedStatement, db Handle, records []*R) e
 	return newIOError(nil, "Stmt.Close", stmt.Close())
 }
 
-func insertRecord(v reflect.Value, recordIndex int, stmt preparedStatement, argumentIndexes [][]int, argumentSlots []any, scanIndexes [][]int, scanSlots []any) error {
+func insertRecord(ctx context.Context, v reflect.Value, recordIndex int, stmt preparedStatement, argumentIndexes [][]int, argumentSlots []any, scanIndexes [][]int, scanSlots []any) error {
 	for idx, index := range argumentIndexes {
 		argumentSlots[idx] = v.FieldByIndex(index).Interface()
 	}
@@ -123,10 +124,10 @@ func insertRecord(v reflect.Value, recordIndex int, stmt preparedStatement, argu
 	}
 	var err error
 	if len(scanSlots) == 0 {
-		_, err = stmt.Exec(argumentSlots...)
+		_, err = stmt.ExecContext(ctx, argumentSlots...)
 	} else {
-		// TODO: using QueryRow for inserting is somehow extremely expensive in terms of allocs; other libraries are doing better by limiting themselves to Exec() + LastInsertId()
-		err = stmt.QueryRow(argumentSlots...).Scan(scanSlots...)
+		// TODO: using QueryRow for inserting is extremely expensive because database/sql allocates a Rows instance under the hood; other libraries are doing better by limiting themselves to ExecContext() + LastInsertId()
+		err = stmt.QueryRowContext(ctx, argumentSlots...).Scan(scanSlots...)
 	}
 	if err != nil {
 		return fmt.Errorf("while inserting record with idx = %d: %w", recordIndex, err)
@@ -138,11 +139,11 @@ func insertRecord(v reflect.Value, recordIndex int, stmt preparedStatement, argu
 // Returns [MissingRecordError] if any of the records does not exist in the database, that is, if for any of the records, the database contains no row with the same primary key values.
 //
 // Returns an error if [NewStore] was called without the [TableNameIs] or [PrimaryKeyIs] options, which are both required to generate a query for this method.
-func (s Store[R]) Update(db Handle, records ...R) error {
+func (s Store[R]) Update(ctx context.Context, db Handle, records ...R) error {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
-	stmt, err := prepare(db, s.plan.Update.Query, "Update", len(records))
+	stmt, err := prepare(ctx, db, s.plan.Update.Query, "Update", len(records))
 	if err != nil {
 		return err
 	}
@@ -154,7 +155,7 @@ func (s Store[R]) Update(db Handle, records ...R) error {
 
 	for idx := range records {
 		v := reflect.ValueOf(&records[idx]).Elem()
-		rowsAffected, err := updateRecord(v, idx, stmt, argumentIndexes, argumentSlots)
+		rowsAffected, err := updateRecord(ctx, v, idx, stmt, argumentIndexes, argumentSlots)
 		if err == nil && rowsAffected == 0 {
 			err = MissingRecordError[R]{records[idx], s.plan}
 		}
@@ -165,11 +166,11 @@ func (s Store[R]) Update(db Handle, records ...R) error {
 	return newIOError(nil, "Stmt.Close", stmt.Close())
 }
 
-func updateRecord(v reflect.Value, recordIndex int, stmt preparedStatement, argumentIndexes [][]int, argumentSlots []any) (int64, error) {
+func updateRecord(ctx context.Context, v reflect.Value, recordIndex int, stmt preparedStatement, argumentIndexes [][]int, argumentSlots []any) (int64, error) {
 	for idx, index := range argumentIndexes {
 		argumentSlots[idx] = v.FieldByIndex(index).Interface()
 	}
-	result, err := stmt.Exec(argumentSlots...)
+	result, err := stmt.ExecContext(ctx, argumentSlots...)
 	if err != nil {
 		return 0, fmt.Errorf("while updating record with idx = %d: %w", recordIndex, err)
 	}
@@ -183,11 +184,11 @@ func updateRecord(v reflect.Value, recordIndex int, stmt preparedStatement, argu
 // Delete executes an SQL DELETE statement for each of the provided records, using their primary keys to locate the respective table rows.
 //
 // Returns an error if [NewStore] was called without the [TableNameIs] or [PrimaryKeyIs] options, which are both required to generate a query for this method.
-func (s Store[R]) Delete(db Handle, records ...R) error {
+func (s Store[R]) Delete(ctx context.Context, db Handle, records ...R) error {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
-	stmt, err := prepare(db, s.plan.Delete.Query, "Delete", len(records))
+	stmt, err := prepare(ctx, db, s.plan.Delete.Query, "Delete", len(records))
 	if err != nil {
 		return err
 	}
@@ -199,7 +200,7 @@ func (s Store[R]) Delete(db Handle, records ...R) error {
 
 	for idx := range records {
 		v := reflect.ValueOf(&records[idx]).Elem()
-		err := deleteRecord(v, idx, stmt, argumentIndexes, argumentSlots)
+		err := deleteRecord(ctx, v, idx, stmt, argumentIndexes, argumentSlots)
 		if err != nil {
 			return newIOError(err, "Stmt.Close", stmt.Close())
 		}
@@ -208,11 +209,11 @@ func (s Store[R]) Delete(db Handle, records ...R) error {
 	return newIOError(nil, "Stmt.Close", stmt.Close())
 }
 
-func deleteRecord(v reflect.Value, recordIndex int, stmt preparedStatement, argumentIndexes [][]int, argumentSlots []any) error {
+func deleteRecord(ctx context.Context, v reflect.Value, recordIndex int, stmt preparedStatement, argumentIndexes [][]int, argumentSlots []any) error {
 	for idx, index := range argumentIndexes {
 		argumentSlots[idx] = v.FieldByIndex(index).Interface()
 	}
-	_, err := stmt.Exec(argumentSlots...)
+	_, err := stmt.ExecContext(ctx, argumentSlots...)
 	if err != nil {
 		return fmt.Errorf("while deleting record with idx = %d: %w", recordIndex, err)
 	}
@@ -227,24 +228,24 @@ func deleteRecord(v reflect.Value, recordIndex int, stmt preparedStatement, argu
 //     Returns an error if [NewStore] was called without the [TableNameIs] or [PrimaryKeyIs] options, which are both required to generate the respective queries for this method.
 //   - For record types that do not have fields declared with the "auto" tag, an INSERT ... ON CONFLICT statement is used.
 //     Returns an error if [NewStore] was called without the [TableNameIs] option, which is required to generate a query for this method.
-func (s Store[R]) Upsert(db Handle, records ...*R) error {
+func (s Store[R]) Upsert(ctx context.Context, db Handle, records ...*R) error {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
 	if len(s.plan.AutoColumnNames) == 0 {
-		stmt, err := prepare(db, s.plan.Upsert.Query, "Upsert", len(records))
+		stmt, err := prepare(ctx, db, s.plan.Upsert.Query, "Upsert", len(records))
 		if err != nil {
 			return err
 		}
-		return s.insertUsing(stmt, db, records)
+		return s.insertUsing(ctx, stmt, db, records)
 	}
 
 	// TODO: respect PrepareThreshold (or not? may be too much bookkeeping overhead for not a whole lot of benefit)
-	insertStmt, err := prepare(db, s.plan.Insert.Query, "Insert", 0)
+	insertStmt, err := prepare(ctx, db, s.plan.Insert.Query, "Insert", 0)
 	if err != nil {
 		return err
 	}
-	updateStmt, err := prepare(db, s.plan.Update.Query, "Update", 0)
+	updateStmt, err := prepare(ctx, db, s.plan.Update.Query, "Update", 0)
 	if err != nil {
 		return err
 	}
@@ -266,10 +267,10 @@ func (s Store[R]) Upsert(db Handle, records ...*R) error {
 		}
 
 		if isInsert {
-			err = insertRecord(v, idx, insertStmt, insertArgumentIndexes, insertArgumentSlots, insertScanIndexes, insertScanSlots)
+			err = insertRecord(ctx, v, idx, insertStmt, insertArgumentIndexes, insertArgumentSlots, insertScanIndexes, insertScanSlots)
 		} else {
 			var rowsAffected int64
-			rowsAffected, err = updateRecord(v, idx, updateStmt, updateArgumentIndexes, updateArgumentSlots)
+			rowsAffected, err = updateRecord(ctx, v, idx, updateStmt, updateArgumentIndexes, updateArgumentSlots)
 			if err == nil && rowsAffected == 0 {
 				err = MissingRecordError[R]{*r, s.plan}
 			}
diff --git a/query_test.go b/query_test.go
index 6f73642..05a3af2 100644
--- a/query_test.go
+++ b/query_test.go
@@ -16,6 +16,7 @@ import (
 )
 
 func TestInsertBasic(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -40,7 +41,7 @@ func TestInsertBasic(t *testing.T) {
 					AndReturnColumns("id").
 					WithRow(int64(42 + idx))
 			}
-			must.Succeed(t, store.Insert(db, records...))
+			must.Succeed(t, store.Insert(ctx, db, records...))
 			for idx, r := range records {
 				assert.Equal(t, r.ID, int64(42+idx))
 			}
@@ -49,6 +50,7 @@ func TestInsertBasic(t *testing.T) {
 }
 
 func TestUpdateBasic(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -72,12 +74,13 @@ func TestUpdateBasic(t *testing.T) {
 					ExpectExecWithArgs(r.Name, r.ID).
 					AndReturnRowsAffected(1)
 			}
-			must.Succeed(t, store.Update(db, records...))
+			must.Succeed(t, store.Update(ctx, db, records...))
 		})
 	}
 }
 
 func TestDeleteBasic(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -101,12 +104,13 @@ func TestDeleteBasic(t *testing.T) {
 					ExpectExecWithArgs(r.ID).
 					AndReturnRowsAffected(1)
 			}
-			must.Succeed(t, store.Delete(db, records...))
+			must.Succeed(t, store.Delete(ctx, db, records...))
 		})
 	}
 }
 
 func TestUpsertBasicWithAutoColumn(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -141,7 +145,7 @@ func TestUpsertBasicWithAutoColumn(t *testing.T) {
 		{Name: "third needs insert"},
 		{ID: 4, Name: "fourth needs update"},
 	}
-	must.Succeed(t, store.Upsert(db, records...))
+	must.Succeed(t, store.Upsert(ctx, db, records...))
 
 	assert.SliceDeepEqual(t, records,
 		&basicRecord{ID: 1, Name: "first needs insert"},
@@ -152,6 +156,7 @@ func TestUpsertBasicWithAutoColumn(t *testing.T) {
 }
 
 func TestWriteQueriesNotPossible(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -165,21 +170,22 @@ func TestWriteQueriesNotPossible(t *testing.T) {
 	)
 
 	r := basicRecord{Name: "foo"}
-	err := store.Insert(db, &r)
+	err := store.Insert(ctx, db, &r)
 	assert.ErrEqual(t, err, "cannot execute Insert() because query could not be autogenerated")
 
-	err = store.Upsert(db, &r)
+	err = store.Upsert(ctx, db, &r)
 	assert.ErrEqual(t, err, "cannot execute Insert() because query could not be autogenerated")
 
 	r.ID = 42
-	err = store.Update(db, r)
+	err = store.Update(ctx, db, r)
 	assert.ErrEqual(t, err, "cannot execute Update() because query could not be autogenerated")
 
-	err = store.Delete(db, r)
+	err = store.Delete(ctx, db, r)
 	assert.ErrEqual(t, err, "cannot execute Delete() because query could not be autogenerated")
 }
 
 func TestWriteQueriesFailDuringPrepare(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -201,7 +207,7 @@ func TestWriteQueriesFailDuringPrepare(t *testing.T) {
 			recordsForInsert[idx] = &basicRecord{Name: "foo"}
 		}
 
-		err := store.Insert(db, recordsForInsert...)
+		err := store.Insert(ctx, db, recordsForInsert...)
 		baseError := `unexpected query: INSERT INTO "basic_records" ("name") VALUES (?) RETURNING "id"`
 		if batchSize < oblast.PrepareThreshold {
 			assert.ErrEqual(t, err, "while inserting record with idx = 0: "+baseError)
@@ -209,7 +215,7 @@ func TestWriteQueriesFailDuringPrepare(t *testing.T) {
 			assert.ErrEqual(t, err, "during Prepare(): "+baseError)
 		}
 
-		err = store.Update(db, records...)
+		err = store.Update(ctx, db, records...)
 		baseError = `unexpected query: UPDATE "basic_records" SET "name" = ? WHERE "id" = ?`
 		if batchSize < oblast.PrepareThreshold {
 			assert.ErrEqual(t, err, "while updating record with idx = 0: "+baseError)
@@ -217,7 +223,7 @@ func TestWriteQueriesFailDuringPrepare(t *testing.T) {
 			assert.ErrEqual(t, err, "during Prepare(): "+baseError)
 		}
 
-		err = store.Delete(db, records...)
+		err = store.Delete(ctx, db, records...)
 		baseError = `unexpected query: DELETE FROM "basic_records" WHERE "id" = ?`
 		if batchSize < oblast.PrepareThreshold {
 			assert.ErrEqual(t, err, "while deleting record with idx = 0: "+baseError)
@@ -228,6 +234,7 @@ func TestWriteQueriesFailDuringPrepare(t *testing.T) {
 }
 
 func TestUpdateOrUpsertFailsOnMissingRecord(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -245,7 +252,7 @@ func TestUpdateOrUpsertFailsOnMissingRecord(t *testing.T) {
 	md.ForQuery(`UPDATE "basic_records" SET "name" = ? WHERE "id" = ?`).
 		ExpectExecWithArgs("changed", 42).
 		AndReturnRowsAffected(0)
-	err := store.Update(db, basicRecord{ID: 42, Name: "changed"})
+	err := store.Update(ctx, db, basicRecord{ID: 42, Name: "changed"})
 	assert.ErrEqual(t, err, "could not UPDATE record that does not exist in the database: id = 42")
 	_, hasCorrectType := err.(oblast.MissingRecordError[basicRecord]) //nolint:errorlint // we explicitly do not want a wrapped error
 	assert.Equal(t, hasCorrectType, true)
@@ -255,13 +262,14 @@ func TestUpdateOrUpsertFailsOnMissingRecord(t *testing.T) {
 	md.ForQuery(`UPDATE "basic_records" SET "name" = ? WHERE "id" = ?`).
 		ExpectExecWithArgs("changed", 42).
 		AndReturnRowsAffected(0)
-	err = store.Upsert(db, &basicRecord{ID: 42, Name: "changed"})
+	err = store.Upsert(ctx, db, &basicRecord{ID: 42, Name: "changed"})
 	assert.ErrEqual(t, err, "could not UPDATE record that does not exist in the database: id = 42")
 	_, hasCorrectType = err.(oblast.MissingRecordError[basicRecord]) //nolint:errorlint // we explicitly do not want a wrapped error
 	assert.Equal(t, hasCorrectType, true)
 }
 
 func TestInsertFailsOnFilledAutoField(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -279,11 +287,12 @@ func TestInsertFailsOnFilledAutoField(t *testing.T) {
 		ExpectQueryWithArgs("existing").
 		AndReturnColumns("id").
 		WithRow(42)
-	err := store.Insert(db, &basicRecord{ID: 23, Name: "third"})
+	err := store.Insert(ctx, db, &basicRecord{ID: 23, Name: "third"})
 	assert.ErrEqual(t, err, `refusing to INSERT record with idx = 0 that already has non-zero values in its "auto" columns`)
 }
 
 func TestInsertAndUpsertWithNoAutoColumns(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -301,7 +310,7 @@ func TestInsertAndUpsertWithNoAutoColumns(t *testing.T) {
 	md.ForQuery(`INSERT INTO "foo_bar_relations" ("foo_id", "bar_id") VALUES (?, ?)`).
 		ExpectExecWithArgs(23, 42).
 		AndReturnRowsAffected(1)
-	must.Succeed(t, store.Insert(db, &relation{23, 42}))
+	must.Succeed(t, store.Insert(ctx, db, &relation{23, 42}))
 
 	// test Upsert()
 	md.ForQuery(`INSERT INTO "foo_bar_relations" ("foo_id", "bar_id") VALUES (?, ?) ON CONFLICT ("foo_id", "bar_id") DO NOTHING`).
@@ -310,10 +319,11 @@ func TestInsertAndUpsertWithNoAutoColumns(t *testing.T) {
 	md.ForQuery(`INSERT INTO "foo_bar_relations" ("foo_id", "bar_id") VALUES (?, ?) ON CONFLICT ("foo_id", "bar_id") DO NOTHING`).
 		ExpectExecWithArgs(3, 4).
 		AndReturnRowsAffected(1)
-	must.Succeed(t, store.Upsert(db, &relation{1, 2}, &relation{3, 4}))
+	must.Succeed(t, store.Upsert(ctx, db, &relation{1, 2}, &relation{3, 4}))
 }
 
 func TestUpsertFailsOnMixedAutoFieldState(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -333,6 +343,6 @@ func TestUpsertFailsOnMixedAutoFieldState(t *testing.T) {
 		Name:      "foo",
 		CreatedAt: time.Time{}, // this looks like we need to INSERT
 	}
-	err := store.Upsert(db, &brokenRecord)
+	err := store.Upsert(ctx, db, &brokenRecord)
 	assert.ErrEqual(t, err, `cannot decide whether to INSERT or UPDATE record with idx = 0: some "auto" columns are zero, others are not`)
 }
diff --git a/select.go b/select.go
index 2cd1e44..6838f8d 100644
--- a/select.go
+++ b/select.go
@@ -4,6 +4,7 @@
 package oblast
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -14,11 +15,11 @@ import (
 // according to the column names reported by the database as part of the result set.
 //
 // An error is returned if any column name in the result set does not correspond to an addressable field in R.
-func (s Store[R]) Select(db Handle, query string, args ...any) ([]R, error) {
+func (s Store[R]) Select(ctx context.Context, db Handle, query string, args ...any) ([]R, error) {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
-	rows, indexes, err := startSelectQuery(db, s.plan, query, args...)
+	rows, indexes, err := startSelectQuery(ctx, db, s.plan, query, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -50,11 +51,11 @@ func (s Store[R]) Select(db Handle, query string, args ...any) ([]R, error) {
 // Besides a condition for the WHERE clause, it may contain additional clauses, such as ORDER BY or LIMIT.
 //
 // Returns an error if [NewStore] was called without the [TableNameIs] option, which is required to generate a query for this method.
-func (s Store[R]) SelectWhere(db Handle, partialQuery string, args ...any) ([]R, error) {
+func (s Store[R]) SelectWhere(ctx context.Context, db Handle, partialQuery string, args ...any) ([]R, error) {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
-	rows, indexes, err := startSelectWhereQuery(db, s.plan, partialQuery, args...)
+	rows, indexes, err := startSelectWhereQuery(ctx, db, s.plan, partialQuery, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -73,8 +74,8 @@ func (s Store[R]) SelectWhere(db Handle, partialQuery string, args ...any) ([]R,
 	return result, newIOError(err, "Rows.Err", rows.Err())
 }
 
-func startSelectQuery(db Handle, plan plan, query string, args ...any) (*sql.Rows, [][]int, error) {
-	rows, err := db.Query(query, args...)
+func startSelectQuery(ctx context.Context, db Handle, plan plan, query string, args ...any) (*sql.Rows, [][]int, error) {
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, nil, fmt.Errorf("during Query(): %w", err)
 	}
@@ -100,12 +101,12 @@ func startSelectQuery(db Handle, plan plan, query string, args ...any) (*sql.Row
 	return rows, indexes, nil
 }
 
-func startSelectWhereQuery(db Handle, plan plan, partialQuery string, args ...any) (rows *sql.Rows, indexes [][]int, err error) {
+func startSelectWhereQuery(ctx context.Context, db Handle, plan plan, partialQuery string, args ...any) (rows *sql.Rows, indexes [][]int, err error) {
 	if plan.Select.Query == "" {
 		return nil, nil, errors.New("cannot execute SelectWhere() because query could not be autogenerated")
 	}
 	query := plan.Select.Query + partialQuery
-	rows, err = db.Query(query, args...)
+	rows, err = db.QueryContext(ctx, query, args...)
 	if err != nil {
 		err = fmt.Errorf("during Query(): %w", err)
 	}
@@ -134,12 +135,12 @@ func collectRow(rows *sql.Rows, plan plan, v reflect.Value, slots []any, indexes
 //
 // Warning: Because of limitations in the interface of database/sql, this function is built on [Store.Select] and cannot be any faster than it.
 // For maximum performance, use [Store.SelectOneWhere] which avoids the overhead of potentially having to read multiple rows.
-func (s Store[R]) SelectOne(db Handle, query string, args ...any) (result R, err error) {
+func (s Store[R]) SelectOne(ctx context.Context, db Handle, query string, args ...any) (result R, err error) {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
 	var results []R
-	results, err = s.Select(db, query, args...)
+	results, err = s.Select(ctx, db, query, args...)
 	if err == nil {
 		if len(results) == 0 {
 			err = sql.ErrNoRows
@@ -155,23 +156,23 @@ func (s Store[R]) SelectOne(db Handle, query string, args ...any) (result R, err
 //
 // This method is more efficient than [Store.SelectOne] on CPU runtime, but has a slight memory allocation overhead per call from query preparation.
 // This can be avoided by using [Store.PrepareSelectQueryWhere] instead.
-func (s Store[R]) SelectOneWhere(db Handle, partialQuery string, args ...any) (result R, err error) {
+func (s Store[R]) SelectOneWhere(ctx context.Context, db Handle, partialQuery string, args ...any) (result R, err error) {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
-	err = selectOneWhere(db, s.plan, reflect.ValueOf(&result).Elem(), partialQuery, args)
+	err = selectOneWhere(ctx, db, s.plan, reflect.ValueOf(&result).Elem(), partialQuery, args)
 	return
 }
 
-func selectOneWhere(db Handle, plan plan, v reflect.Value, partialQuery string, args []any) error {
+func selectOneWhere(ctx context.Context, db Handle, plan plan, v reflect.Value, partialQuery string, args []any) error {
 	if plan.Select.Query == "" {
 		return errors.New("cannot execute SelectOneWhere() because query could not be autogenerated")
 	}
 	query := plan.Select.Query + partialQuery
-	return selectOne(db, plan, v, query, args)
+	return selectOne(ctx, db, plan, v, query, args)
 }
 
-func selectOne(db Handle, plan plan, v reflect.Value, query string, args []any) error {
+func selectOne(ctx context.Context, db Handle, plan plan, v reflect.Value, query string, args []any) error {
 	for _, index := range plan.IndexesOfTransparentPointerStructs {
 		f := v.FieldByIndex(index)
 		f.Set(reflect.New(f.Type().Elem()))
@@ -180,7 +181,7 @@ func selectOne(db Handle, plan plan, v reflect.Value, query string, args []any)
 	for idx, index := range plan.Select.ScanIndexes {
 		slots[idx] = v.FieldByIndex(index).Addr().Interface()
 	}
-	return db.QueryRow(query, args...).Scan(slots...)
+	return db.QueryRowContext(ctx, query, args...).Scan(slots...)
 }
 
 // PrepareSelectQueryWhere performs the same query string preparation as [Store.SelectWhere] or [Store.SelectOneWhere].
@@ -219,11 +220,11 @@ type PreparedSelectQuery[R any] struct {
 }
 
 // Select behaves the same as [Store.SelectWhere], but uses the query that was precomputed when q was constructed.
-func (q PreparedSelectQuery[R]) Select(db Handle, args ...any) ([]R, error) {
+func (q PreparedSelectQuery[R]) Select(ctx context.Context, db Handle, args ...any) ([]R, error) {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
-	rows, indexes, err := startSelectQuery(db, q.store.plan, q.query, args...)
+	rows, indexes, err := startSelectQuery(ctx, db, q.store.plan, q.query, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -243,10 +244,10 @@ func (q PreparedSelectQuery[R]) Select(db Handle, args ...any) ([]R, error) {
 }
 
 // SelectOne behaves the same as [Store.SelectOneWhere], but uses the query that was precomputed when q was constructed.
-func (q PreparedSelectQuery[R]) SelectOne(db Handle, args ...any) (result R, err error) {
+func (q PreparedSelectQuery[R]) SelectOne(ctx context.Context, db Handle, args ...any) (result R, err error) {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
-	err = selectOne(db, q.store.plan, reflect.ValueOf(&result).Elem(), q.query, args)
+	err = selectOne(ctx, db, q.store.plan, reflect.ValueOf(&result).Elem(), q.query, args)
 	return
 }
diff --git a/select_test.go b/select_test.go
index f38fbdd..d4b1970 100644
--- a/select_test.go
+++ b/select_test.go
@@ -16,6 +16,7 @@ import (
 )
 
 func TestSelectReturningSomeRecords(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -35,7 +36,7 @@ func TestSelectReturningSomeRecords(t *testing.T) {
 			AndReturnColumns("name", "id").
 			WithRow("foo", 1).
 			WithRow("bar", 2)
-		records := must.Return(store.Select(db, `SELECT * FROM basic_records WHERE id < ?`, 3))(t)
+		records := must.Return(store.Select(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3))(t)
 		assert.SliceEqual(t, records,
 			basicRecord{1, "foo"},
 			basicRecord{2, "bar"},
@@ -48,7 +49,7 @@ func TestSelectReturningSomeRecords(t *testing.T) {
 			AndReturnColumns("id", "name").
 			WithRow(1, "ffoo").
 			WithRow(2, "bbar")
-		records := must.Return(store.SelectWhere(db, `id < ?`, 3))(t)
+		records := must.Return(store.SelectWhere(ctx, db, `id < ?`, 3))(t)
 		assert.SliceEqual(t, records,
 			basicRecord{1, "ffoo"},
 			basicRecord{2, "bbar"},
@@ -62,7 +63,7 @@ func TestSelectReturningSomeRecords(t *testing.T) {
 			WithRow(1, "fffoo").
 			WithRow(2, "bbbar")
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		records := must.Return(query.Select(db, 3))(t)
+		records := must.Return(query.Select(ctx, db, 3))(t)
 		assert.SliceEqual(t, records,
 			basicRecord{1, "fffoo"},
 			basicRecord{2, "bbbar"},
@@ -75,7 +76,7 @@ func TestSelectReturningSomeRecords(t *testing.T) {
 			AndReturnColumns("name", "id").
 			WithRow("ffffoo", 1).
 			WithRow("bbbbar", 2)
-		record := must.Return(store.SelectOne(db, `SELECT * FROM basic_records WHERE id < ?`, 3))(t)
+		record := must.Return(store.SelectOne(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3))(t)
 		assert.Equal(t, record, basicRecord{1, "ffffoo"})
 	})
 
@@ -85,7 +86,7 @@ func TestSelectReturningSomeRecords(t *testing.T) {
 			AndReturnColumns("id", "name").
 			WithRow(1, "fffffoo").
 			WithRow(2, "bbbbbar")
-		record := must.Return(store.SelectOneWhere(db, `id < ?`, 3))(t)
+		record := must.Return(store.SelectOneWhere(ctx, db, `id < ?`, 3))(t)
 		assert.Equal(t, record, basicRecord{1, "fffffoo"})
 	})
 
@@ -96,12 +97,13 @@ func TestSelectReturningSomeRecords(t *testing.T) {
 			WithRow(1, "ffffffoo").
 			WithRow(2, "bbbbbbar")
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		record := must.Return(query.SelectOne(db, 3))(t)
+		record := must.Return(query.SelectOne(ctx, db, 3))(t)
 		assert.Equal(t, record, basicRecord{1, "ffffffoo"})
 	})
 }
 
 func TestSelectReturningNoRecords(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -119,7 +121,7 @@ func TestSelectReturningNoRecords(t *testing.T) {
 		md.ForQuery(`SELECT * FROM basic_records WHERE id < ?`).
 			ExpectQueryWithArgs(3).
 			AndReturnColumns("name", "id")
-		records := must.Return(store.Select(db, `SELECT * FROM basic_records WHERE id < ?`, 3))(t)
+		records := must.Return(store.Select(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3))(t)
 		assert.SliceEqual(t, records, nil...)
 	})
 
@@ -127,7 +129,7 @@ func TestSelectReturningNoRecords(t *testing.T) {
 		md.ForQuery(`SELECT "id", "name" FROM "basic_records" WHERE id < ?`).
 			ExpectQueryWithArgs(3).
 			AndReturnColumns("id", "name")
-		records := must.Return(store.SelectWhere(db, `id < ?`, 3))(t)
+		records := must.Return(store.SelectWhere(ctx, db, `id < ?`, 3))(t)
 		assert.SliceEqual(t, records, nil...)
 	})
 
@@ -136,7 +138,7 @@ func TestSelectReturningNoRecords(t *testing.T) {
 			ExpectQueryWithArgs(3).
 			AndReturnColumns("id", "name")
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		records := must.Return(query.Select(db, 3))(t)
+		records := must.Return(query.Select(ctx, db, 3))(t)
 		assert.SliceEqual(t, records, nil...)
 	})
 
@@ -144,7 +146,7 @@ func TestSelectReturningNoRecords(t *testing.T) {
 		md.ForQuery(`SELECT * FROM basic_records WHERE id < ?`).
 			ExpectQueryWithArgs(3).
 			AndReturnColumns("name", "id")
-		_, err := store.SelectOne(db, `SELECT * FROM basic_records WHERE id < ?`, 3)
+		_, err := store.SelectOne(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3)
 		assert.ErrEqual(t, err, sql.ErrNoRows.Error())
 	})
 
@@ -152,7 +154,7 @@ func TestSelectReturningNoRecords(t *testing.T) {
 		md.ForQuery(`SELECT "id", "name" FROM "basic_records" WHERE id < ?`).
 			ExpectQueryWithArgs(3).
 			AndReturnColumns("id", "name")
-		_, err := store.SelectOneWhere(db, `id < ?`, 3)
+		_, err := store.SelectOneWhere(ctx, db, `id < ?`, 3)
 		assert.ErrEqual(t, err, sql.ErrNoRows.Error())
 	})
 
@@ -161,12 +163,13 @@ func TestSelectReturningNoRecords(t *testing.T) {
 			ExpectQueryWithArgs(3).
 			AndReturnColumns("id", "name")
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		_, err := query.SelectOne(db, 3)
+		_, err := query.SelectOne(ctx, db, 3)
 		assert.ErrEqual(t, err, sql.ErrNoRows.Error())
 	})
 }
 
 func TestSelectIntoUnexpectedField(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -193,18 +196,19 @@ func TestSelectIntoUnexpectedField(t *testing.T) {
 
 	t.Run("using Store.Select", func(t *testing.T) {
 		commonSetup()
-		_, err := store.Select(db, `SELECT * FROM basic_records WHERE id < ?`, 3)
+		_, err := store.Select(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3)
 		assert.ErrEqual(t, err, expectedError)
 	})
 
 	t.Run("using Store.SelectOne", func(t *testing.T) {
 		commonSetup()
-		_, err := store.SelectOne(db, `SELECT * FROM basic_records WHERE id < ?`, 3)
+		_, err := store.SelectOne(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3)
 		assert.ErrEqual(t, err, expectedError)
 	})
 }
 
 func TestSelectWithScanError(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -229,44 +233,45 @@ func TestSelectWithScanError(t *testing.T) {
 
 	t.Run("using Store.Select", func(t *testing.T) {
 		commonSetup(`SELECT * FROM basic_records WHERE id < ?`)
-		_, err := store.Select(db, `SELECT * FROM basic_records WHERE id < ?`, 3)
+		_, err := store.Select(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3)
 		assert.ErrEqual(t, err, expectedError)
 	})
 
 	t.Run("using Store.SelectWhere", func(t *testing.T) {
 		commonSetup(`SELECT "id", "created_at" FROM "basic_records" WHERE id < ?`)
-		_, err := store.SelectWhere(db, `id < ?`, 3)
+		_, err := store.SelectWhere(ctx, db, `id < ?`, 3)
 		assert.ErrEqual(t, err, expectedError)
 	})
 
 	t.Run("using PreparedSelectQuery.Select", func(t *testing.T) {
 		commonSetup(`SELECT "id", "created_at" FROM "basic_records" WHERE id < ?`)
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		_, err := query.Select(db, 3)
+		_, err := query.Select(ctx, db, 3)
 		assert.ErrEqual(t, err, expectedError)
 	})
 
 	t.Run("using Store.SelectOne", func(t *testing.T) {
 		commonSetup(`SELECT * FROM basic_records WHERE id < ?`)
-		_, err := store.SelectOne(db, `SELECT * FROM basic_records WHERE id < ?`, 3)
+		_, err := store.SelectOne(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3)
 		assert.ErrEqual(t, err, expectedError)
 	})
 
 	t.Run("using Store.SelectOneWhere", func(t *testing.T) {
 		commonSetup(`SELECT "id", "created_at" FROM "basic_records" WHERE id < ?`)
-		_, err := store.SelectOneWhere(db, `id < ?`, 3)
+		_, err := store.SelectOneWhere(ctx, db, `id < ?`, 3)
 		assert.ErrEqual(t, err, expectedError)
 	})
 
 	t.Run("using PreparedSelectQuery.SelectOne", func(t *testing.T) {
 		commonSetup(`SELECT "id", "created_at" FROM "basic_records" WHERE id < ?`)
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		_, err := query.SelectOne(db, 3)
+		_, err := query.SelectOne(ctx, db, 3)
 		assert.ErrEqual(t, err, expectedError)
 	})
 }
 
 func TestSelectIntoEmbeddedTypes(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -299,7 +304,7 @@ func TestSelectIntoEmbeddedTypes(t *testing.T) {
 
 	t.Run("using Store.Select", func(t *testing.T) {
 		commonSetup(`SELECT * FROM composite_records`)
-		records := must.Return(store.Select(db, `SELECT * FROM composite_records`))(t)
+		records := must.Return(store.Select(ctx, db, `SELECT * FROM composite_records`))(t)
 		assert.SliceDeepEqual(t, records,
 			compositeRecord{1, HasCreatedAt{time.Unix(1, 0)}, &HasUpdatedAt{new(time.Unix(3, 0))}},
 			compositeRecord{2, HasCreatedAt{time.Unix(2, 0)}, &HasUpdatedAt{nil}},
@@ -308,7 +313,7 @@ func TestSelectIntoEmbeddedTypes(t *testing.T) {
 
 	t.Run("using Store.SelectWhere", func(t *testing.T) {
 		commonSetup(`SELECT "id", "created_at", "updated_at" FROM "composite_records" WHERE TRUE`)
-		records := must.Return(store.SelectWhere(db, `TRUE`))(t)
+		records := must.Return(store.SelectWhere(ctx, db, `TRUE`))(t)
 		assert.SliceDeepEqual(t, records,
 			compositeRecord{1, HasCreatedAt{time.Unix(1, 0)}, &HasUpdatedAt{new(time.Unix(3, 0))}},
 			compositeRecord{2, HasCreatedAt{time.Unix(2, 0)}, &HasUpdatedAt{nil}},
@@ -318,7 +323,7 @@ func TestSelectIntoEmbeddedTypes(t *testing.T) {
 	t.Run("using PreparedSelectQuery.Select", func(t *testing.T) {
 		commonSetup(`SELECT "id", "created_at", "updated_at" FROM "composite_records" WHERE TRUE`)
 		query := store.MustPrepareSelectQueryWhere(`TRUE`)
-		records := must.Return(query.Select(db))(t)
+		records := must.Return(query.Select(ctx, db))(t)
 		assert.SliceDeepEqual(t, records,
 			compositeRecord{1, HasCreatedAt{time.Unix(1, 0)}, &HasUpdatedAt{new(time.Unix(3, 0))}},
 			compositeRecord{2, HasCreatedAt{time.Unix(2, 0)}, &HasUpdatedAt{nil}},
@@ -327,7 +332,7 @@ func TestSelectIntoEmbeddedTypes(t *testing.T) {
 
 	t.Run("using Store.SelectOne", func(t *testing.T) {
 		commonSetup(`SELECT * FROM composite_records`)
-		record := must.Return(store.SelectOne(db, `SELECT * FROM composite_records`))(t)
+		record := must.Return(store.SelectOne(ctx, db, `SELECT * FROM composite_records`))(t)
 		assert.DeepEqual(t, record,
 			compositeRecord{1, HasCreatedAt{time.Unix(1, 0)}, &HasUpdatedAt{new(time.Unix(3, 0))}},
 		)
@@ -335,7 +340,7 @@ func TestSelectIntoEmbeddedTypes(t *testing.T) {
 
 	t.Run("using Store.SelectOneWhere", func(t *testing.T) {
 		commonSetup(`SELECT "id", "created_at", "updated_at" FROM "composite_records" WHERE TRUE`)
-		record := must.Return(store.SelectOneWhere(db, `TRUE`))(t)
+		record := must.Return(store.SelectOneWhere(ctx, db, `TRUE`))(t)
 		assert.DeepEqual(t, record,
 			compositeRecord{1, HasCreatedAt{time.Unix(1, 0)}, &HasUpdatedAt{new(time.Unix(3, 0))}},
 		)
@@ -344,7 +349,7 @@ func TestSelectIntoEmbeddedTypes(t *testing.T) {
 	t.Run("using PreparedSelectQuery.SelectOne", func(t *testing.T) {
 		commonSetup(`SELECT "id", "created_at", "updated_at" FROM "composite_records" WHERE TRUE`)
 		query := store.MustPrepareSelectQueryWhere(`TRUE`)
-		record := must.Return(query.SelectOne(db))(t)
+		record := must.Return(query.SelectOne(ctx, db))(t)
 		assert.DeepEqual(t, record,
 			compositeRecord{1, HasCreatedAt{time.Unix(1, 0)}, &HasUpdatedAt{new(time.Unix(3, 0))}},
 		)
@@ -352,6 +357,7 @@ func TestSelectIntoEmbeddedTypes(t *testing.T) {
 }
 
 func TestSelectCapturingQueryError(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -366,39 +372,40 @@ func TestSelectCapturingQueryError(t *testing.T) {
 	)
 
 	t.Run("using Store.Select", func(t *testing.T) {
-		_, err := store.Select(db, `SELECT * FROM basic_records WHERE id < ?`, 3)
+		_, err := store.Select(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3)
 		assert.ErrEqual(t, err, "during Query(): unexpected query: SELECT * FROM basic_records WHERE id < ?")
 	})
 
 	t.Run("using Store.SelectWhere", func(t *testing.T) {
-		_, err := store.SelectWhere(db, `id < ?`, 3)
+		_, err := store.SelectWhere(ctx, db, `id < ?`, 3)
 		assert.ErrEqual(t, err, `during Query(): unexpected query: SELECT "id", "name" FROM "basic_records" WHERE id < ?`)
 	})
 
 	t.Run("using PreparedSelectQuery.Select", func(t *testing.T) {
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		_, err := query.Select(db, 3)
+		_, err := query.Select(ctx, db, 3)
 		assert.ErrEqual(t, err, `during Query(): unexpected query: SELECT "id", "name" FROM "basic_records" WHERE id < ?`)
 	})
 
 	t.Run("using Store.SelectOne", func(t *testing.T) {
-		_, err := store.SelectOne(db, `SELECT * FROM basic_records WHERE id < ?`, 3)
+		_, err := store.SelectOne(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3)
 		assert.ErrEqual(t, err, "during Query(): unexpected query: SELECT * FROM basic_records WHERE id < ?")
 	})
 
 	t.Run("using Store.SelectOneWhere", func(t *testing.T) {
-		_, err := store.SelectOneWhere(db, `id < ?`, 3)
+		_, err := store.SelectOneWhere(ctx, db, `id < ?`, 3)
 		assert.ErrEqual(t, err, `unexpected query: SELECT "id", "name" FROM "basic_records" WHERE id < ?`)
 	})
 
 	t.Run("using PreparedSelectQuery.SelectOne", func(t *testing.T) {
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		_, err := query.SelectOne(db, 3)
+		_, err := query.SelectOne(ctx, db, 3)
 		assert.ErrEqual(t, err, `unexpected query: SELECT "id", "name" FROM "basic_records" WHERE id < ?`)
 	})
 }
 
 func TestSelectCapturingCloseError(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -423,44 +430,45 @@ func TestSelectCapturingCloseError(t *testing.T) {
 
 	t.Run("using Store.Select", func(t *testing.T) {
 		commonSetup(`SELECT * FROM basic_records WHERE id < ?`)
-		_, err := store.Select(db, `SELECT * FROM basic_records WHERE id < ?`, 3)
+		_, err := store.Select(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3)
 		assert.ErrEqual(t, err, "during Rows.Err(): datacenter on fire")
 	})
 
 	t.Run("using Store.SelectWhere", func(t *testing.T) {
 		commonSetup(`SELECT "id", "name" FROM "basic_records" WHERE id < ?`)
-		_, err := store.SelectWhere(db, `id < ?`, 3)
+		_, err := store.SelectWhere(ctx, db, `id < ?`, 3)
 		assert.ErrEqual(t, err, "during Rows.Err(): datacenter on fire")
 	})
 
 	t.Run("using PreparedSelectQuery.Select", func(t *testing.T) {
 		commonSetup(`SELECT "id", "name" FROM "basic_records" WHERE id < ?`)
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		_, err := query.Select(db, 3)
+		_, err := query.Select(ctx, db, 3)
 		assert.ErrEqual(t, err, "during Rows.Err(): datacenter on fire")
 	})
 
 	t.Run("using Store.SelectOne", func(t *testing.T) {
 		commonSetup(`SELECT * FROM basic_records WHERE id < ?`)
-		_, err := store.SelectOne(db, `SELECT * FROM basic_records WHERE id < ?`, 3)
+		_, err := store.SelectOne(ctx, db, `SELECT * FROM basic_records WHERE id < ?`, 3)
 		assert.ErrEqual(t, err, "during Rows.Err(): datacenter on fire")
 	})
 
 	t.Run("using Store.SelectOneWhere", func(t *testing.T) {
 		commonSetup(`SELECT "id", "name" FROM "basic_records" WHERE id < ?`)
-		_, err := store.SelectOneWhere(db, `id < ?`, 3)
+		_, err := store.SelectOneWhere(ctx, db, `id < ?`, 3)
 		assert.ErrEqual(t, err, "datacenter on fire")
 	})
 
 	t.Run("using PreparedSelectQuery.SelectOne", func(t *testing.T) {
 		commonSetup(`SELECT "id", "name" FROM "basic_records" WHERE id < ?`)
 		query := store.MustPrepareSelectQueryWhere(`id < ?`)
-		_, err := query.SelectOne(db, 3)
+		_, err := query.SelectOne(ctx, db, 3)
 		assert.ErrEqual(t, err, "datacenter on fire")
 	})
 }
 
 func TestSelectNotPossibleWithoutTableName(t *testing.T) {
+	ctx := t.Context()
 	md := mock.NewDriver()
 	db := sql.OpenDB(md)
 
@@ -471,12 +479,12 @@ func TestSelectNotPossibleWithoutTableName(t *testing.T) {
 	store := oblast.MustNewStore[basicRecord](oblast.SqliteDialect())
 
 	t.Run("using Store.SelectWhere", func(t *testing.T) {
-		_, err := store.SelectWhere(db, `id < ?`, 3)
+		_, err := store.SelectWhere(ctx, db, `id < ?`, 3)
 		assert.ErrEqual(t, err, "cannot execute SelectWhere() because query could not be autogenerated")
 	})
 
 	t.Run("using Store.SelectOneWhere", func(t *testing.T) {
-		_, err := store.SelectOneWhere(db, `id < ?`, 3)
+		_, err := store.SelectOneWhere(ctx, db, `id < ?`, 3)
 		assert.ErrEqual(t, err, "cannot execute SelectOneWhere() because query could not be autogenerated")
 	})