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