Skip to content

Commit e13df02

Browse files
kardianosbradfitz
kardianos
authored and
bradfitz
committed
database/sql: add context methods
Add context methods to sql and sql/driver methods. If the driver doesn't implement context methods the connection pool will still handle timeouts when a query fails to return in time or when a connection is not available from the pool in time. There will be a follow-up CL that will add support for context values that specify transaction levels and modes that a driver can use. Fixes #15123 Change-Id: Ia99f3957aa3f177b23044dd99d4ec217491a30a7 Reviewed-on: https://go-review.googlesource.com/29381 Reviewed-by: Brad Fitzpatrick <[email protected]> Run-TryBot: Brad Fitzpatrick <[email protected]> TryBot-Result: Gobot Gobot <[email protected]>
1 parent 54a72d9 commit e13df02

File tree

5 files changed

+525
-100
lines changed

5 files changed

+525
-100
lines changed

‎src/database/sql/ctxutil.go

+231
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
// Copyright 2016 The Go Authors. All rights reserved.
2+
// Use of this source code is governed by a BSD-style
3+
// license that can be found in the LICENSE file.
4+
5+
package sql
6+
7+
import (
8+
"context"
9+
"database/sql/driver"
10+
"errors"
11+
)
12+
13+
func ctxDriverPrepare(ctx context.Context, ci driver.Conn, query string) (driver.Stmt, error) {
14+
if ciCtx, is := ci.(driver.ConnPrepareContext); is {
15+
return ciCtx.PrepareContext(ctx, query)
16+
}
17+
type R struct {
18+
err error
19+
panic interface{}
20+
si driver.Stmt
21+
}
22+
23+
rc := make(chan R, 1)
24+
go func() {
25+
r := R{}
26+
defer func() {
27+
if v := recover(); v != nil {
28+
r.panic = v
29+
}
30+
rc <- r
31+
}()
32+
r.si, r.err = ci.Prepare(query)
33+
}()
34+
select {
35+
case <-ctx.Done():
36+
go func() {
37+
<-rc
38+
close(rc)
39+
}()
40+
return nil, ctx.Err()
41+
case r := <-rc:
42+
if r.panic != nil {
43+
panic(r.panic)
44+
}
45+
return r.si, r.err
46+
}
47+
}
48+
49+
func ctxDriverExec(ctx context.Context, execer driver.Execer, query string, dargs []driver.Value) (driver.Result, error) {
50+
if execerCtx, is := execer.(driver.ExecerContext); is {
51+
return execerCtx.ExecContext(ctx, query, dargs)
52+
}
53+
type R struct {
54+
err error
55+
panic interface{}
56+
resi driver.Result
57+
}
58+
59+
rc := make(chan R, 1)
60+
go func() {
61+
r := R{}
62+
defer func() {
63+
if v := recover(); v != nil {
64+
r.panic = v
65+
}
66+
rc <- r
67+
}()
68+
r.resi, r.err = execer.Exec(query, dargs)
69+
}()
70+
select {
71+
case <-ctx.Done():
72+
go func() {
73+
<-rc
74+
close(rc)
75+
}()
76+
return nil, ctx.Err()
77+
case r := <-rc:
78+
if r.panic != nil {
79+
panic(r.panic)
80+
}
81+
return r.resi, r.err
82+
}
83+
}
84+
85+
func ctxDriverQuery(ctx context.Context, queryer driver.Queryer, query string, dargs []driver.Value) (driver.Rows, error) {
86+
if queryerCtx, is := queryer.(driver.QueryerContext); is {
87+
return queryerCtx.QueryContext(ctx, query, dargs)
88+
}
89+
type R struct {
90+
err error
91+
panic interface{}
92+
rowsi driver.Rows
93+
}
94+
95+
rc := make(chan R, 1)
96+
go func() {
97+
r := R{}
98+
defer func() {
99+
if v := recover(); v != nil {
100+
r.panic = v
101+
}
102+
rc <- r
103+
}()
104+
r.rowsi, r.err = queryer.Query(query, dargs)
105+
}()
106+
select {
107+
case <-ctx.Done():
108+
go func() {
109+
<-rc
110+
close(rc)
111+
}()
112+
return nil, ctx.Err()
113+
case r := <-rc:
114+
if r.panic != nil {
115+
panic(r.panic)
116+
}
117+
return r.rowsi, r.err
118+
}
119+
}
120+
121+
func ctxDriverStmtExec(ctx context.Context, si driver.Stmt, dargs []driver.Value) (driver.Result, error) {
122+
if siCtx, is := si.(driver.StmtExecContext); is {
123+
return siCtx.ExecContext(ctx, dargs)
124+
}
125+
type R struct {
126+
err error
127+
panic interface{}
128+
resi driver.Result
129+
}
130+
131+
rc := make(chan R, 1)
132+
go func() {
133+
r := R{}
134+
defer func() {
135+
if v := recover(); v != nil {
136+
r.panic = v
137+
}
138+
rc <- r
139+
}()
140+
r.resi, r.err = si.Exec(dargs)
141+
}()
142+
select {
143+
case <-ctx.Done():
144+
go func() {
145+
<-rc
146+
close(rc)
147+
}()
148+
return nil, ctx.Err()
149+
case r := <-rc:
150+
if r.panic != nil {
151+
panic(r.panic)
152+
}
153+
return r.resi, r.err
154+
}
155+
}
156+
157+
func ctxDriverStmtQuery(ctx context.Context, si driver.Stmt, dargs []driver.Value) (driver.Rows, error) {
158+
if siCtx, is := si.(driver.StmtQueryContext); is {
159+
return siCtx.QueryContext(ctx, dargs)
160+
}
161+
type R struct {
162+
err error
163+
panic interface{}
164+
rowsi driver.Rows
165+
}
166+
167+
rc := make(chan R, 1)
168+
go func() {
169+
r := R{}
170+
defer func() {
171+
if v := recover(); v != nil {
172+
r.panic = v
173+
}
174+
rc <- r
175+
}()
176+
r.rowsi, r.err = si.Query(dargs)
177+
}()
178+
select {
179+
case <-ctx.Done():
180+
go func() {
181+
<-rc
182+
close(rc)
183+
}()
184+
return nil, ctx.Err()
185+
case r := <-rc:
186+
if r.panic != nil {
187+
panic(r.panic)
188+
}
189+
return r.rowsi, r.err
190+
}
191+
}
192+
193+
var errLevelNotSupported = errors.New("sql: selected isolation level is not supported")
194+
195+
func ctxDriverBegin(ctx context.Context, ci driver.Conn) (driver.Tx, error) {
196+
if ciCtx, is := ci.(driver.ConnBeginContext); is {
197+
return ciCtx.BeginContext(ctx)
198+
}
199+
// TODO(kardianos): check the transaction level in ctx. If set and non-default
200+
// then return an error here as the BeginContext driver value is not supported.
201+
202+
type R struct {
203+
err error
204+
panic interface{}
205+
txi driver.Tx
206+
}
207+
rc := make(chan R, 1)
208+
go func() {
209+
r := R{}
210+
defer func() {
211+
if v := recover(); v != nil {
212+
r.panic = v
213+
}
214+
rc <- r
215+
}()
216+
r.txi, r.err = ci.Begin()
217+
}()
218+
select {
219+
case <-ctx.Done():
220+
go func() {
221+
<-rc
222+
close(rc)
223+
}()
224+
return nil, ctx.Err()
225+
case r := <-rc:
226+
if r.panic != nil {
227+
panic(r.panic)
228+
}
229+
return r.txi, r.err
230+
}
231+
}

‎src/database/sql/driver/driver.go

+45-1
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@
88
// Most code should use package sql.
99
package driver
1010

11-
import "errors"
11+
import (
12+
"context"
13+
"errors"
14+
)
1215

1316
// Value is a value that drivers must be able to handle.
1417
// It is either nil or an instance of one of these types:
@@ -65,6 +68,12 @@ type Execer interface {
6568
Exec(query string, args []Value) (Result, error)
6669
}
6770

71+
// ExecerContext is like execer, but must honor the context timeout and return
72+
// when the context is cancelled.
73+
type ExecerContext interface {
74+
ExecContext(ctx context.Context, query string, args []Value) (Result, error)
75+
}
76+
6877
// Queryer is an optional interface that may be implemented by a Conn.
6978
//
7079
// If a Conn does not implement Queryer, the sql package's DB.Query will
@@ -76,6 +85,12 @@ type Queryer interface {
7685
Query(query string, args []Value) (Rows, error)
7786
}
7887

88+
// QueryerContext is like Queryer, but most honor the context timeout and return
89+
// when the context is cancelled.
90+
type QueryerContext interface {
91+
QueryContext(ctx context.Context, query string, args []Value) (Rows, error)
92+
}
93+
7994
// Conn is a connection to a database. It is not used concurrently
8095
// by multiple goroutines.
8196
//
@@ -98,6 +113,23 @@ type Conn interface {
98113
Begin() (Tx, error)
99114
}
100115

116+
// ConnPrepareContext enhances the Conn interface with context.
117+
type ConnPrepareContext interface {
118+
// PrepareContext returns a prepared statement, bound to this connection.
119+
// context is for the preparation of the statement,
120+
// it must not store the context within the statement itself.
121+
PrepareContext(ctx context.Context, query string) (Stmt, error)
122+
}
123+
124+
// ConnBeginContext enhances the Conn interface with context.
125+
type ConnBeginContext interface {
126+
// BeginContext starts and returns a new transaction.
127+
// the provided context should be used to roll the transaction back
128+
// if it is cancelled. If there is an isolation level in context
129+
// that is not supported by the driver an error must be returned.
130+
BeginContext(ctx context.Context) (Tx, error)
131+
}
132+
101133
// Result is the result of a query execution.
102134
type Result interface {
103135
// LastInsertId returns the database's auto-generated ID
@@ -139,6 +171,18 @@ type Stmt interface {
139171
Query(args []Value) (Rows, error)
140172
}
141173

174+
// StmtExecContext enhances the Stmt interface by providing Exec with context.
175+
type StmtExecContext interface {
176+
// ExecContext must honor the context timeout and return when it is cancelled.
177+
ExecContext(ctx context.Context, args []Value) (Result, error)
178+
}
179+
180+
// StmtQueryContext enhances the Stmt interface by providing Query with context.
181+
type StmtQueryContext interface {
182+
// QueryContext must honor the context timeout and return when it is cancelled.
183+
QueryContext(ctx context.Context, args []Value) (Rows, error)
184+
}
185+
142186
// ColumnConverter may be optionally implemented by Stmt if the
143187
// statement is aware of its own columns' types and can convert from
144188
// any type to a driver Value.

0 commit comments

Comments
 (0)