1package pq
2
3import (
4	"context"
5	"database/sql"
6	"database/sql/driver"
7	"runtime"
8	"strings"
9	"testing"
10	"time"
11)
12
13func TestMultipleSimpleQuery(t *testing.T) {
14	db := openTestConn(t)
15	defer db.Close()
16
17	rows, err := db.Query("select 1; set time zone default; select 2; select 3")
18	if err != nil {
19		t.Fatal(err)
20	}
21	defer rows.Close()
22
23	var i int
24	for rows.Next() {
25		if err := rows.Scan(&i); err != nil {
26			t.Fatal(err)
27		}
28		if i != 1 {
29			t.Fatalf("expected 1, got %d", i)
30		}
31	}
32	if !rows.NextResultSet() {
33		t.Fatal("expected more result sets", rows.Err())
34	}
35	for rows.Next() {
36		if err := rows.Scan(&i); err != nil {
37			t.Fatal(err)
38		}
39		if i != 2 {
40			t.Fatalf("expected 2, got %d", i)
41		}
42	}
43
44	// Make sure that if we ignore a result we can still query.
45
46	rows, err = db.Query("select 4; select 5")
47	if err != nil {
48		t.Fatal(err)
49	}
50	defer rows.Close()
51
52	for rows.Next() {
53		if err := rows.Scan(&i); err != nil {
54			t.Fatal(err)
55		}
56		if i != 4 {
57			t.Fatalf("expected 4, got %d", i)
58		}
59	}
60	if !rows.NextResultSet() {
61		t.Fatal("expected more result sets", rows.Err())
62	}
63	for rows.Next() {
64		if err := rows.Scan(&i); err != nil {
65			t.Fatal(err)
66		}
67		if i != 5 {
68			t.Fatalf("expected 5, got %d", i)
69		}
70	}
71	if rows.NextResultSet() {
72		t.Fatal("unexpected result set")
73	}
74}
75
76const contextRaceIterations = 100
77
78func TestContextCancelExec(t *testing.T) {
79	db := openTestConn(t)
80	defer db.Close()
81
82	ctx, cancel := context.WithCancel(context.Background())
83
84	// Delay execution for just a bit until db.ExecContext has begun.
85	defer time.AfterFunc(time.Millisecond*10, cancel).Stop()
86
87	// Not canceled until after the exec has started.
88	if _, err := db.ExecContext(ctx, "select pg_sleep(1)"); err == nil {
89		t.Fatal("expected error")
90	} else if err.Error() != "pq: canceling statement due to user request" {
91		t.Fatalf("unexpected error: %s", err)
92	}
93
94	// Context is already canceled, so error should come before execution.
95	if _, err := db.ExecContext(ctx, "select pg_sleep(1)"); err == nil {
96		t.Fatal("expected error")
97	} else if err.Error() != "context canceled" {
98		t.Fatalf("unexpected error: %s", err)
99	}
100
101	for i := 0; i < contextRaceIterations; i++ {
102		func() {
103			ctx, cancel := context.WithCancel(context.Background())
104			defer cancel()
105			if _, err := db.ExecContext(ctx, "select 1"); err != nil {
106				t.Fatal(err)
107			}
108		}()
109
110		if _, err := db.Exec("select 1"); err != nil {
111			t.Fatal(err)
112		}
113	}
114}
115
116func TestContextCancelQuery(t *testing.T) {
117	db := openTestConn(t)
118	defer db.Close()
119
120	ctx, cancel := context.WithCancel(context.Background())
121
122	// Delay execution for just a bit until db.QueryContext has begun.
123	defer time.AfterFunc(time.Millisecond*10, cancel).Stop()
124
125	// Not canceled until after the exec has started.
126	if _, err := db.QueryContext(ctx, "select pg_sleep(1)"); err == nil {
127		t.Fatal("expected error")
128	} else if err.Error() != "pq: canceling statement due to user request" {
129		t.Fatalf("unexpected error: %s", err)
130	}
131
132	// Context is already canceled, so error should come before execution.
133	if _, err := db.QueryContext(ctx, "select pg_sleep(1)"); err == nil {
134		t.Fatal("expected error")
135	} else if err.Error() != "context canceled" {
136		t.Fatalf("unexpected error: %s", err)
137	}
138
139	for i := 0; i < contextRaceIterations; i++ {
140		func() {
141			ctx, cancel := context.WithCancel(context.Background())
142			rows, err := db.QueryContext(ctx, "select 1")
143			cancel()
144			if err != nil {
145				t.Fatal(err)
146			} else if err := rows.Close(); err != nil && err != driver.ErrBadConn && err != context.Canceled {
147				t.Fatal(err)
148			}
149		}()
150
151		if rows, err := db.Query("select 1"); err != nil {
152			t.Fatal(err)
153		} else if err := rows.Close(); err != nil {
154			t.Fatal(err)
155		}
156	}
157}
158
159// TestIssue617 tests that a failed query in QueryContext doesn't lead to a
160// goroutine leak.
161func TestIssue617(t *testing.T) {
162	db := openTestConn(t)
163	defer db.Close()
164
165	const N = 10
166
167	numGoroutineStart := runtime.NumGoroutine()
168	for i := 0; i < N; i++ {
169		func() {
170			ctx, cancel := context.WithCancel(context.Background())
171			defer cancel()
172			_, err := db.QueryContext(ctx, `SELECT * FROM DOESNOTEXIST`)
173			pqErr, _ := err.(*Error)
174			// Expecting "pq: relation \"doesnotexist\" does not exist" error.
175			if err == nil || pqErr == nil || pqErr.Code != "42P01" {
176				t.Fatalf("expected undefined table error, got %v", err)
177			}
178		}()
179	}
180
181	// Give time for goroutines to terminate
182	delayTime := time.Millisecond * 50
183	waitTime := time.Second
184	iterations := int(waitTime / delayTime)
185
186	var numGoroutineFinish int
187	for i := 0; i < iterations; i++ {
188		time.Sleep(delayTime)
189
190		numGoroutineFinish = runtime.NumGoroutine()
191
192		// We use N/2 and not N because the GC and other actors may increase or
193		// decrease the number of goroutines.
194		if numGoroutineFinish-numGoroutineStart < N/2 {
195			return
196		}
197	}
198
199	t.Errorf("goroutine leak detected, was %d, now %d", numGoroutineStart, numGoroutineFinish)
200}
201
202func TestContextCancelBegin(t *testing.T) {
203	db := openTestConn(t)
204	defer db.Close()
205
206	ctx, cancel := context.WithCancel(context.Background())
207	tx, err := db.BeginTx(ctx, nil)
208	if err != nil {
209		t.Fatal(err)
210	}
211
212	// Delay execution for just a bit until tx.Exec has begun.
213	defer time.AfterFunc(time.Millisecond*10, cancel).Stop()
214
215	// Not canceled until after the exec has started.
216	if _, err := tx.Exec("select pg_sleep(1)"); err == nil {
217		t.Fatal("expected error")
218	} else if err.Error() != "pq: canceling statement due to user request" {
219		t.Fatalf("unexpected error: %s", err)
220	}
221
222	// Transaction is canceled, so expect an error.
223	if _, err := tx.Query("select pg_sleep(1)"); err == nil {
224		t.Fatal("expected error")
225	} else if err != sql.ErrTxDone {
226		t.Fatalf("unexpected error: %s", err)
227	}
228
229	// Context is canceled, so cannot begin a transaction.
230	if _, err := db.BeginTx(ctx, nil); err == nil {
231		t.Fatal("expected error")
232	} else if err.Error() != "context canceled" {
233		t.Fatalf("unexpected error: %s", err)
234	}
235
236	for i := 0; i < contextRaceIterations; i++ {
237		func() {
238			ctx, cancel := context.WithCancel(context.Background())
239			tx, err := db.BeginTx(ctx, nil)
240			cancel()
241			if err != nil {
242				t.Fatal(err)
243			} else if err := tx.Rollback(); err != nil &&
244				err.Error() != "pq: canceling statement due to user request" &&
245				err != sql.ErrTxDone && err != driver.ErrBadConn && err != context.Canceled {
246				t.Fatal(err)
247			}
248		}()
249
250		if tx, err := db.Begin(); err != nil {
251			t.Fatal(err)
252		} else if err := tx.Rollback(); err != nil {
253			t.Fatal(err)
254		}
255	}
256}
257
258func TestTxOptions(t *testing.T) {
259	db := openTestConn(t)
260	defer db.Close()
261	ctx := context.Background()
262
263	tests := []struct {
264		level     sql.IsolationLevel
265		isolation string
266	}{
267		{
268			level:     sql.LevelDefault,
269			isolation: "",
270		},
271		{
272			level:     sql.LevelReadUncommitted,
273			isolation: "read uncommitted",
274		},
275		{
276			level:     sql.LevelReadCommitted,
277			isolation: "read committed",
278		},
279		{
280			level:     sql.LevelRepeatableRead,
281			isolation: "repeatable read",
282		},
283		{
284			level:     sql.LevelSerializable,
285			isolation: "serializable",
286		},
287	}
288
289	for _, test := range tests {
290		for _, ro := range []bool{true, false} {
291			tx, err := db.BeginTx(ctx, &sql.TxOptions{
292				Isolation: test.level,
293				ReadOnly:  ro,
294			})
295			if err != nil {
296				t.Fatal(err)
297			}
298
299			var isolation string
300			err = tx.QueryRow("select current_setting('transaction_isolation')").Scan(&isolation)
301			if err != nil {
302				t.Fatal(err)
303			}
304
305			if test.isolation != "" && isolation != test.isolation {
306				t.Errorf("wrong isolation level: %s != %s", isolation, test.isolation)
307			}
308
309			var isRO string
310			err = tx.QueryRow("select current_setting('transaction_read_only')").Scan(&isRO)
311			if err != nil {
312				t.Fatal(err)
313			}
314
315			if ro != (isRO == "on") {
316				t.Errorf("read/[write,only] not set: %t != %s for level %s",
317					ro, isRO, test.isolation)
318			}
319
320			tx.Rollback()
321		}
322	}
323
324	_, err := db.BeginTx(ctx, &sql.TxOptions{
325		Isolation: sql.LevelLinearizable,
326	})
327	if err == nil {
328		t.Fatal("expected LevelLinearizable to fail")
329	}
330	if !strings.Contains(err.Error(), "isolation level not supported") {
331		t.Errorf("Expected error to mention isolation level, got %q", err)
332	}
333}
334