[BUG] contrib/database/sql: Context passed to PrepareContext is used in later queries

[WKBae]

Version of dd-trace-go

v1.54.0

Describe what happened:

Context passed to db.PrepareContext() is kept throughout the lifespan of stmt, yielding "context canceled" error when using the stmt after the context is canceled.

The following code shadows ctx with a context derived from s.ctx, which is already canceled.

dd-trace-go/contrib/database/sql/stmt.go

Lines 62 to 65 in 30bfd52

	func (s *tracedStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (rows driver.Rows, err error) {
	start := time.Now()
	if stmtQueryContext, ok := s.Stmt.(driver.StmtQueryContext); ok {
	ctx, end := startTraceTask(s.ctx, QueryTypeQuery)

Go documentation states that The provided context is used for the preparation of the statement, not for the execution of the statement.

Describe what you expected:
Context used in PrepareContext is not preserved; returned Stmt should be usable even after the prepare context has been expired.

Steps to reproduce the issue:

Code to reproduce: (actual driver implementation does not matter, tested MySQL and SQLite)

package main

import (
	"context"
	"database/sql"
	"fmt"

	"github.com/mattn/go-sqlite3"
	sqltrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
)

func connectSQLite() *sql.DB {
	sqltrace.Register("sqlite3", &sqlite3.SQLiteDriver{})
	db, err := sqltrace.Open("sqlite3", "file::memory:?cache=shared")
	if err != nil {
		panic(err)
	}
	return db
}

func connectSQLiteNoTrace() *sql.DB {
	db, err := sql.Open("sqlite3", "file::memory:?cache=shared")
	if err != nil {
		panic(err)
	}
	return db
}

func main() {
	db := connectSQLite()
	defer db.Close()

	// Prepare using ctx1
	ctx1, cancel := context.WithCancel(context.Background())
	stmt, err := db.PrepareContext(ctx1, "SELECT 1")
	if err != nil {
		panic(err)
	}
	defer stmt.Close()
	cancel()

	// Query stmt using ctx2
	ctx2 := context.Background()
	var result int
	err = stmt.QueryRowContext(ctx2).Scan(&result)
	if err != nil {
		panic(err) // panics with "context canceled"
	}

	fmt.Println("result:", result)
}

When using connectSQLiteNoTrace(), the result prints out correctly. However, when using connectSQLite(), it returns "context canceled" error on QueryRowContext.

Additional environment details (Version of Go, Operating System, etc.):
Tested on go1.20.5 darwin/arm64, but OS shouldn't matter.

[nsrip-dd]

Thank you for the report, and for the reproducer! I've sent #2173 to address this.