copy_test.go 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463
  1. package pq
  2. import (
  3. "bytes"
  4. "database/sql"
  5. "database/sql/driver"
  6. "strings"
  7. "testing"
  8. )
  9. func TestCopyInStmt(t *testing.T) {
  10. stmt := CopyIn("table name")
  11. if stmt != `COPY "table name" () FROM STDIN` {
  12. t.Fatal(stmt)
  13. }
  14. stmt = CopyIn("table name", "column 1", "column 2")
  15. if stmt != `COPY "table name" ("column 1", "column 2") FROM STDIN` {
  16. t.Fatal(stmt)
  17. }
  18. stmt = CopyIn(`table " name """`, `co"lumn""`)
  19. if stmt != `COPY "table "" name """"""" ("co""lumn""""") FROM STDIN` {
  20. t.Fatal(stmt)
  21. }
  22. }
  23. func TestCopyInSchemaStmt(t *testing.T) {
  24. stmt := CopyInSchema("schema name", "table name")
  25. if stmt != `COPY "schema name"."table name" () FROM STDIN` {
  26. t.Fatal(stmt)
  27. }
  28. stmt = CopyInSchema("schema name", "table name", "column 1", "column 2")
  29. if stmt != `COPY "schema name"."table name" ("column 1", "column 2") FROM STDIN` {
  30. t.Fatal(stmt)
  31. }
  32. stmt = CopyInSchema(`schema " name """`, `table " name """`, `co"lumn""`)
  33. if stmt != `COPY "schema "" name """"""".`+
  34. `"table "" name """"""" ("co""lumn""""") FROM STDIN` {
  35. t.Fatal(stmt)
  36. }
  37. }
  38. func TestCopyInMultipleValues(t *testing.T) {
  39. db := openTestConn(t)
  40. defer db.Close()
  41. txn, err := db.Begin()
  42. if err != nil {
  43. t.Fatal(err)
  44. }
  45. defer txn.Rollback()
  46. _, err = txn.Exec("CREATE TEMP TABLE temp (a int, b varchar)")
  47. if err != nil {
  48. t.Fatal(err)
  49. }
  50. stmt, err := txn.Prepare(CopyIn("temp", "a", "b"))
  51. if err != nil {
  52. t.Fatal(err)
  53. }
  54. longString := strings.Repeat("#", 500)
  55. for i := 0; i < 500; i++ {
  56. _, err = stmt.Exec(int64(i), longString)
  57. if err != nil {
  58. t.Fatal(err)
  59. }
  60. }
  61. _, err = stmt.Exec()
  62. if err != nil {
  63. t.Fatal(err)
  64. }
  65. err = stmt.Close()
  66. if err != nil {
  67. t.Fatal(err)
  68. }
  69. var num int
  70. err = txn.QueryRow("SELECT COUNT(*) FROM temp").Scan(&num)
  71. if err != nil {
  72. t.Fatal(err)
  73. }
  74. if num != 500 {
  75. t.Fatalf("expected 500 items, not %d", num)
  76. }
  77. }
  78. func TestCopyInRaiseStmtTrigger(t *testing.T) {
  79. db := openTestConn(t)
  80. defer db.Close()
  81. if getServerVersion(t, db) < 90000 {
  82. var exists int
  83. err := db.QueryRow("SELECT 1 FROM pg_language WHERE lanname = 'plpgsql'").Scan(&exists)
  84. if err == sql.ErrNoRows {
  85. t.Skip("language PL/PgSQL does not exist; skipping TestCopyInRaiseStmtTrigger")
  86. } else if err != nil {
  87. t.Fatal(err)
  88. }
  89. }
  90. txn, err := db.Begin()
  91. if err != nil {
  92. t.Fatal(err)
  93. }
  94. defer txn.Rollback()
  95. _, err = txn.Exec("CREATE TEMP TABLE temp (a int, b varchar)")
  96. if err != nil {
  97. t.Fatal(err)
  98. }
  99. _, err = txn.Exec(`
  100. CREATE OR REPLACE FUNCTION pg_temp.temptest()
  101. RETURNS trigger AS
  102. $BODY$ begin
  103. raise notice 'Hello world';
  104. return new;
  105. end $BODY$
  106. LANGUAGE plpgsql`)
  107. if err != nil {
  108. t.Fatal(err)
  109. }
  110. _, err = txn.Exec(`
  111. CREATE TRIGGER temptest_trigger
  112. BEFORE INSERT
  113. ON temp
  114. FOR EACH ROW
  115. EXECUTE PROCEDURE pg_temp.temptest()`)
  116. if err != nil {
  117. t.Fatal(err)
  118. }
  119. stmt, err := txn.Prepare(CopyIn("temp", "a", "b"))
  120. if err != nil {
  121. t.Fatal(err)
  122. }
  123. longString := strings.Repeat("#", 500)
  124. _, err = stmt.Exec(int64(1), longString)
  125. if err != nil {
  126. t.Fatal(err)
  127. }
  128. _, err = stmt.Exec()
  129. if err != nil {
  130. t.Fatal(err)
  131. }
  132. err = stmt.Close()
  133. if err != nil {
  134. t.Fatal(err)
  135. }
  136. var num int
  137. err = txn.QueryRow("SELECT COUNT(*) FROM temp").Scan(&num)
  138. if err != nil {
  139. t.Fatal(err)
  140. }
  141. if num != 1 {
  142. t.Fatalf("expected 1 items, not %d", num)
  143. }
  144. }
  145. func TestCopyInTypes(t *testing.T) {
  146. db := openTestConn(t)
  147. defer db.Close()
  148. txn, err := db.Begin()
  149. if err != nil {
  150. t.Fatal(err)
  151. }
  152. defer txn.Rollback()
  153. _, err = txn.Exec("CREATE TEMP TABLE temp (num INTEGER, text VARCHAR, blob BYTEA, nothing VARCHAR)")
  154. if err != nil {
  155. t.Fatal(err)
  156. }
  157. stmt, err := txn.Prepare(CopyIn("temp", "num", "text", "blob", "nothing"))
  158. if err != nil {
  159. t.Fatal(err)
  160. }
  161. _, err = stmt.Exec(int64(1234567890), "Héllö\n ☃!\r\t\\", []byte{0, 255, 9, 10, 13}, nil)
  162. if err != nil {
  163. t.Fatal(err)
  164. }
  165. _, err = stmt.Exec()
  166. if err != nil {
  167. t.Fatal(err)
  168. }
  169. err = stmt.Close()
  170. if err != nil {
  171. t.Fatal(err)
  172. }
  173. var num int
  174. var text string
  175. var blob []byte
  176. var nothing sql.NullString
  177. err = txn.QueryRow("SELECT * FROM temp").Scan(&num, &text, &blob, &nothing)
  178. if err != nil {
  179. t.Fatal(err)
  180. }
  181. if num != 1234567890 {
  182. t.Fatal("unexpected result", num)
  183. }
  184. if text != "Héllö\n ☃!\r\t\\" {
  185. t.Fatal("unexpected result", text)
  186. }
  187. if !bytes.Equal(blob, []byte{0, 255, 9, 10, 13}) {
  188. t.Fatal("unexpected result", blob)
  189. }
  190. if nothing.Valid {
  191. t.Fatal("unexpected result", nothing.String)
  192. }
  193. }
  194. func TestCopyInWrongType(t *testing.T) {
  195. db := openTestConn(t)
  196. defer db.Close()
  197. txn, err := db.Begin()
  198. if err != nil {
  199. t.Fatal(err)
  200. }
  201. defer txn.Rollback()
  202. _, err = txn.Exec("CREATE TEMP TABLE temp (num INTEGER)")
  203. if err != nil {
  204. t.Fatal(err)
  205. }
  206. stmt, err := txn.Prepare(CopyIn("temp", "num"))
  207. if err != nil {
  208. t.Fatal(err)
  209. }
  210. defer stmt.Close()
  211. _, err = stmt.Exec("Héllö\n ☃!\r\t\\")
  212. if err != nil {
  213. t.Fatal(err)
  214. }
  215. _, err = stmt.Exec()
  216. if err == nil {
  217. t.Fatal("expected error")
  218. }
  219. if pge := err.(*Error); pge.Code.Name() != "invalid_text_representation" {
  220. t.Fatalf("expected 'invalid input syntax for integer' error, got %s (%+v)", pge.Code.Name(), pge)
  221. }
  222. }
  223. func TestCopyOutsideOfTxnError(t *testing.T) {
  224. db := openTestConn(t)
  225. defer db.Close()
  226. _, err := db.Prepare(CopyIn("temp", "num"))
  227. if err == nil {
  228. t.Fatal("COPY outside of transaction did not return an error")
  229. }
  230. if err != errCopyNotSupportedOutsideTxn {
  231. t.Fatalf("expected %s, got %s", err, err.Error())
  232. }
  233. }
  234. func TestCopyInBinaryError(t *testing.T) {
  235. db := openTestConn(t)
  236. defer db.Close()
  237. txn, err := db.Begin()
  238. if err != nil {
  239. t.Fatal(err)
  240. }
  241. defer txn.Rollback()
  242. _, err = txn.Exec("CREATE TEMP TABLE temp (num INTEGER)")
  243. if err != nil {
  244. t.Fatal(err)
  245. }
  246. _, err = txn.Prepare("COPY temp (num) FROM STDIN WITH binary")
  247. if err != errBinaryCopyNotSupported {
  248. t.Fatalf("expected %s, got %+v", errBinaryCopyNotSupported, err)
  249. }
  250. // check that the protocol is in a valid state
  251. err = txn.Rollback()
  252. if err != nil {
  253. t.Fatal(err)
  254. }
  255. }
  256. func TestCopyFromError(t *testing.T) {
  257. db := openTestConn(t)
  258. defer db.Close()
  259. txn, err := db.Begin()
  260. if err != nil {
  261. t.Fatal(err)
  262. }
  263. defer txn.Rollback()
  264. _, err = txn.Exec("CREATE TEMP TABLE temp (num INTEGER)")
  265. if err != nil {
  266. t.Fatal(err)
  267. }
  268. _, err = txn.Prepare("COPY temp (num) TO STDOUT")
  269. if err != errCopyToNotSupported {
  270. t.Fatalf("expected %s, got %+v", errCopyToNotSupported, err)
  271. }
  272. // check that the protocol is in a valid state
  273. err = txn.Rollback()
  274. if err != nil {
  275. t.Fatal(err)
  276. }
  277. }
  278. func TestCopySyntaxError(t *testing.T) {
  279. db := openTestConn(t)
  280. defer db.Close()
  281. txn, err := db.Begin()
  282. if err != nil {
  283. t.Fatal(err)
  284. }
  285. defer txn.Rollback()
  286. _, err = txn.Prepare("COPY ")
  287. if err == nil {
  288. t.Fatal("expected error")
  289. }
  290. if pge := err.(*Error); pge.Code.Name() != "syntax_error" {
  291. t.Fatalf("expected syntax error, got %s (%+v)", pge.Code.Name(), pge)
  292. }
  293. // check that the protocol is in a valid state
  294. err = txn.Rollback()
  295. if err != nil {
  296. t.Fatal(err)
  297. }
  298. }
  299. // Tests for connection errors in copyin.resploop()
  300. func TestCopyRespLoopConnectionError(t *testing.T) {
  301. db := openTestConn(t)
  302. defer db.Close()
  303. txn, err := db.Begin()
  304. if err != nil {
  305. t.Fatal(err)
  306. }
  307. defer txn.Rollback()
  308. var pid int
  309. err = txn.QueryRow("SELECT pg_backend_pid()").Scan(&pid)
  310. if err != nil {
  311. t.Fatal(err)
  312. }
  313. _, err = txn.Exec("CREATE TEMP TABLE temp (a int)")
  314. if err != nil {
  315. t.Fatal(err)
  316. }
  317. stmt, err := txn.Prepare(CopyIn("temp", "a"))
  318. if err != nil {
  319. t.Fatal(err)
  320. }
  321. defer stmt.Close()
  322. _, err = db.Exec("SELECT pg_terminate_backend($1)", pid)
  323. if err != nil {
  324. t.Fatal(err)
  325. }
  326. if getServerVersion(t, db) < 90500 {
  327. // We have to try and send something over, since postgres before
  328. // version 9.5 won't process SIGTERMs while it's waiting for
  329. // CopyData/CopyEnd messages; see tcop/postgres.c.
  330. _, err = stmt.Exec(1)
  331. if err != nil {
  332. t.Fatal(err)
  333. }
  334. }
  335. _, err = stmt.Exec()
  336. if err == nil {
  337. t.Fatalf("expected error")
  338. }
  339. pge, ok := err.(*Error)
  340. if !ok {
  341. if err == driver.ErrBadConn {
  342. // likely an EPIPE
  343. } else {
  344. t.Fatalf("expected *pq.Error or driver.ErrBadConn, got %+#v", err)
  345. }
  346. } else if pge.Code.Name() != "admin_shutdown" {
  347. t.Fatalf("expected admin_shutdown, got %s", pge.Code.Name())
  348. }
  349. _ = stmt.Close()
  350. }
  351. func BenchmarkCopyIn(b *testing.B) {
  352. db := openTestConn(b)
  353. defer db.Close()
  354. txn, err := db.Begin()
  355. if err != nil {
  356. b.Fatal(err)
  357. }
  358. defer txn.Rollback()
  359. _, err = txn.Exec("CREATE TEMP TABLE temp (a int, b varchar)")
  360. if err != nil {
  361. b.Fatal(err)
  362. }
  363. stmt, err := txn.Prepare(CopyIn("temp", "a", "b"))
  364. if err != nil {
  365. b.Fatal(err)
  366. }
  367. for i := 0; i < b.N; i++ {
  368. _, err = stmt.Exec(int64(i), "hello world!")
  369. if err != nil {
  370. b.Fatal(err)
  371. }
  372. }
  373. _, err = stmt.Exec()
  374. if err != nil {
  375. b.Fatal(err)
  376. }
  377. err = stmt.Close()
  378. if err != nil {
  379. b.Fatal(err)
  380. }
  381. var num int
  382. err = txn.QueryRow("SELECT COUNT(*) FROM temp").Scan(&num)
  383. if err != nil {
  384. b.Fatal(err)
  385. }
  386. if num != b.N {
  387. b.Fatalf("expected %d items, not %d", b.N, num)
  388. }
  389. }