fix: timeout on query should return context.DeadlineExceeded (#4060)

This commit is contained in:
Kevin Wan
2024-04-10 12:17:39 +08:00
committed by GitHub
parent d1f24ab70f
commit a66ae0d4c4
2 changed files with 20 additions and 1 deletions

View File

@@ -296,7 +296,7 @@ func (db *commonSqlConn) queryRows(ctx context.Context, scanner func(*sql.Rows)
return query(ctx, conn, func(rows *sql.Rows) error { return query(ctx, conn, func(rows *sql.Rows) error {
e := scanner(rows) e := scanner(rows)
if e != nil { if e != nil && !errors.Is(e, context.DeadlineExceeded) {
scanFailed = true scanFailed = true
} }
return e return e

View File

@@ -4,6 +4,7 @@ import (
"context" "context"
"database/sql" "database/sql"
"errors" "errors"
"strconv"
"testing" "testing"
"time" "time"
@@ -290,6 +291,24 @@ func TestStmtBreaker(t *testing.T) {
}) })
} }
func TestQueryScanTimeout(t *testing.T) {
dbtest.RunTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel()
row := sqlmock.NewRows([]string{"foo"})
for i := 0; i < 10000; i++ {
row = row.AddRow("bar" + strconv.Itoa(i))
}
var val []struct {
Foo int
Bar string
}
conn := NewSqlConnFromDB(db)
err := conn.QueryRowsCtx(ctx, &val, "any")
assert.ErrorIs(t, err, context.DeadlineExceeded)
})
}
type mockedSessionConn struct { type mockedSessionConn struct {
lastInsertId int64 lastInsertId int64
rowsAffected int64 rowsAffected int64