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