portal_test.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606
  1. package pgserver
  2. import (
  3. "bufio"
  4. "bytes"
  5. "strconv"
  6. "testing"
  7. )
  8. // Bounded Execute is only reachable over the raw protocol. pgx v5.11.0 sends
  9. // maxRows = 0 for every Execute (pgconn.ExecParams, ExecPrepared, and the
  10. // Pipeline equivalents have no row-limit argument), and database/sql drives pgx,
  11. // so neither a pgx nor a database/sql test can request a suspended portal. These
  12. // tests therefore drive the protocol messages directly, one round trip per
  13. // Execute.
  14. // portalSession sends extended-protocol messages to a connection backed by a
  15. // real executor and returns the frames the client would receive for that round
  16. // trip.
  17. type portalSession struct {
  18. t *testing.T
  19. c *Connection
  20. conn *countingConn
  21. }
  22. func newPortalSession(t *testing.T) *portalSession {
  23. t.Helper()
  24. c, _ := newDataConnection(t)
  25. conn := &countingConn{}
  26. c.conn = conn
  27. c.writer = bufio.NewWriter(conn)
  28. return &portalSession{t: t, c: c, conn: conn}
  29. }
  30. // send handles the messages, flushes buffered output the way a client's Flush
  31. // message would, and returns the frames produced since the previous send.
  32. func (s *portalSession) send(messages ...*Message) []*Message {
  33. s.t.Helper()
  34. for _, message := range messages {
  35. if err := s.c.handleMessage(message); err != nil {
  36. s.t.Fatalf("handle %c: %v", message.Type, err)
  37. }
  38. }
  39. if err := s.c.handleMessage(&Message{Type: MsgFlush}); err != nil {
  40. s.t.Fatalf("flush: %v", err)
  41. }
  42. frames := readAllMessages(s.t, s.conn.bytes())
  43. s.conn.reset()
  44. return frames
  45. }
  46. // setup runs SQL outside the protocol; the statements are the test fixture, not
  47. // the behavior under test.
  48. func (s *portalSession) setup(statements ...string) {
  49. s.t.Helper()
  50. for _, sql := range statements {
  51. if _, err := s.c.executor.Execute(parseStmt(s.t, sql)); err != nil {
  52. s.t.Fatalf("%s: %v", sql, err)
  53. }
  54. }
  55. }
  56. // countRows reports how many rows a table holds, so a test can prove a DML
  57. // portal ran once.
  58. func (s *portalSession) countRows(table string) int64 {
  59. s.t.Helper()
  60. result, err := s.c.executor.Execute(parseStmt(s.t, "SELECT COUNT(*) FROM "+table))
  61. if err != nil {
  62. s.t.Fatalf("count %s: %v", table, err)
  63. }
  64. return result.Rows[0][0].(int64)
  65. }
  66. func executeRowsMessage(portal string, maxRows int32) *Message {
  67. mb := NewMessageBuilder()
  68. mb.WriteString(portal)
  69. mb.WriteInt32(maxRows)
  70. return &Message{Type: MsgExecute, Data: mb.Bytes()}
  71. }
  72. func closePortalMessage(portal string) *Message {
  73. return &Message{Type: MsgClose, Data: append([]byte("P"+portal), 0)}
  74. }
  75. func queryMessageData(sql string) *Message {
  76. return &Message{Type: MsgQuery, Data: append([]byte(sql), 0)}
  77. }
  78. func syncMessage() *Message {
  79. return &Message{Type: MsgSync}
  80. }
  81. func commandTagOf(message *Message) string {
  82. if message == nil || message.Type != MsgCommandComplete {
  83. return ""
  84. }
  85. if end := bytes.IndexByte(message.Data, 0); end >= 0 {
  86. return string(message.Data[:end])
  87. }
  88. return string(message.Data)
  89. }
  90. func assertFrameTypes(t *testing.T, frames []*Message, want string) {
  91. t.Helper()
  92. if got := messageTypes(frames); got != want {
  93. t.Fatalf("frames = %q, want %q", got, want)
  94. }
  95. }
  96. func assertRowIDs(t *testing.T, frames []*Message, want ...string) {
  97. t.Helper()
  98. rows := dataRowsRaw(t, frames)
  99. if len(rows) != len(want) {
  100. t.Fatalf("got %d rows, want %d", len(rows), len(want))
  101. }
  102. for i, id := range want {
  103. if len(rows[i]) != 1 || string(rows[i][0]) != id {
  104. t.Fatalf("row %d = %q, want %q", i, rows[i], id)
  105. }
  106. }
  107. }
  108. func readyStatus(t *testing.T, frames []*Message) byte {
  109. t.Helper()
  110. for _, message := range frames {
  111. if message.Type == MsgReadyForQuery {
  112. if len(message.Data) != 1 {
  113. t.Fatalf("ReadyForQuery payload = %q", message.Data)
  114. }
  115. return message.Data[0]
  116. }
  117. }
  118. t.Fatalf("no ReadyForQuery in %q", messageTypes(frames))
  119. return 0
  120. }
  121. const portalSelectSQL = "SELECT id FROM t ORDER BY id"
  122. func createNumberedTable(t *testing.T, s *portalSession, rows int) {
  123. t.Helper()
  124. s.setup("CREATE TABLE t (id INTEGER PRIMARY KEY)")
  125. for i := 1; i <= rows; i++ {
  126. s.setup("INSERT INTO t (id) VALUES (" + strconv.Itoa(i) + ")")
  127. }
  128. }
  129. // A bounded Execute returns at most maxRows rows and suspends; each later
  130. // Execute resumes where the previous one stopped. The RowDescription is sent by
  131. // the first Execute only, and Describe('P') still answers NoData without
  132. // disturbing that sequence.
  133. func TestBoundedExecuteSuspendsAndResumes(t *testing.T) {
  134. s := newPortalSession(t)
  135. createNumberedTable(t, s, 5)
  136. frames := s.send(
  137. parseMessage("stmt", portalSelectSQL),
  138. bindResultMessage("p", "stmt", nil, nil),
  139. &Message{Type: MsgDescribe, Data: []byte("Pp\x00")},
  140. )
  141. assertFrameTypes(t, frames, "12n")
  142. first := s.send(executeRowsMessage("p", 2))
  143. assertFrameTypes(t, first, "TDDs")
  144. assertRowIDs(t, first, "1", "2")
  145. second := s.send(executeRowsMessage("p", 2))
  146. assertFrameTypes(t, second, "DDs")
  147. assertRowIDs(t, second, "3", "4")
  148. third := s.send(executeRowsMessage("p", 2))
  149. assertFrameTypes(t, third, "DC")
  150. assertRowIDs(t, third, "5")
  151. if tag := commandTagOf(third[len(third)-1]); tag != "SELECT 5" {
  152. t.Fatalf("completion tag = %q, want %q", tag, "SELECT 5")
  153. }
  154. // A completed portal repeats its tag without re-running the statement.
  155. fourth := s.send(executeRowsMessage("p", 0))
  156. assertFrameTypes(t, fourth, "C")
  157. if tag := commandTagOf(fourth[0]); tag != "SELECT 5" {
  158. t.Fatalf("repeated completion tag = %q, want %q", tag, "SELECT 5")
  159. }
  160. // Completion drops the row data, so an exhausted portal only keeps the
  161. // metadata a repeated Execute answers with.
  162. if rows := s.c.portals["p"].execution.result.Rows; len(rows) != 0 {
  163. t.Fatalf("completed portal still holds %d rows", len(rows))
  164. }
  165. }
  166. // maxRows = 0 keeps the pre-existing meaning: every remaining row in one
  167. // Execute followed by CommandComplete.
  168. func TestExecuteUnboundedStreamsAllRows(t *testing.T) {
  169. s := newPortalSession(t)
  170. createNumberedTable(t, s, 3)
  171. s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
  172. frames := s.send(executeRowsMessage("p", 0))
  173. assertFrameTypes(t, frames, "TDDDC")
  174. assertRowIDs(t, frames, "1", "2", "3")
  175. if tag := commandTagOf(frames[len(frames)-1]); tag != "SELECT 3" {
  176. t.Fatalf("completion tag = %q, want %q", tag, "SELECT 3")
  177. }
  178. }
  179. // A limit that falls exactly on the last row still suspends, as PostgreSQL
  180. // does; the following Execute reports completion with no further rows.
  181. func TestBoundedExecuteExactLimitSuspends(t *testing.T) {
  182. s := newPortalSession(t)
  183. createNumberedTable(t, s, 3)
  184. s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
  185. frames := s.send(executeRowsMessage("p", 3))
  186. assertFrameTypes(t, frames, "TDDDs")
  187. assertRowIDs(t, frames, "1", "2", "3")
  188. frames = s.send(executeRowsMessage("p", 0))
  189. assertFrameTypes(t, frames, "C")
  190. if tag := commandTagOf(frames[0]); tag != "SELECT 3" {
  191. t.Fatalf("completion tag = %q, want %q", tag, "SELECT 3")
  192. }
  193. }
  194. // Close('P') destroys a suspended portal, and a later Execute reports 34000
  195. // until Sync recovers the cycle.
  196. func TestCloseEndsSuspendedPortal(t *testing.T) {
  197. s := newPortalSession(t)
  198. createNumberedTable(t, s, 3)
  199. s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
  200. assertFrameTypes(t, s.send(executeRowsMessage("p", 1)), "TDs")
  201. assertFrameTypes(t, s.send(closePortalMessage("p")), "3")
  202. frames := s.send(executeRowsMessage("p", 1))
  203. assertFrameTypes(t, frames, "E")
  204. if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
  205. t.Fatalf("SQLSTATE = %q, want %q", code, ErrCodeUndefinedCursor)
  206. }
  207. assertFrameTypes(t, s.send(syncMessage()), "Z")
  208. if s.c.extendedFailed {
  209. t.Fatal("Sync did not clear the extended-failure state")
  210. }
  211. }
  212. // The unnamed portal is replaced by every Bind, which restarts its materialized
  213. // result, row cursor, and row metadata.
  214. func TestUnnamedRebindRestartsPortal(t *testing.T) {
  215. s := newPortalSession(t)
  216. createNumberedTable(t, s, 5)
  217. s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("", "stmt", nil, nil))
  218. assertFrameTypes(t, s.send(executeRowsMessage("", 2)), "TDDs")
  219. frames := s.send(bindResultMessage("", "stmt", nil, nil), executeRowsMessage("", 0))
  220. assertFrameTypes(t, frames, "2TDDDDDC")
  221. assertRowIDs(t, frames, "1", "2", "3", "4", "5")
  222. }
  223. // A named portal cannot be replaced by another Bind: PostgreSQL reports 42P03,
  224. // and the rejected Bind fails the cycle, so an explicit transaction is aborted
  225. // and the portal goes with the ROLLBACK.
  226. func TestNamedRebindIsRejected(t *testing.T) {
  227. s := newPortalSession(t)
  228. createNumberedTable(t, s, 3)
  229. if status := readyStatus(t, s.send(queryMessageData("BEGIN"))); status != TxStatusInBlock {
  230. t.Fatalf("BEGIN status = %c, want %c", status, TxStatusInBlock)
  231. }
  232. s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
  233. assertFrameTypes(t, s.send(executeRowsMessage("p", 1)), "TDs")
  234. frames := s.send(bindResultMessage("p", "stmt", nil, nil))
  235. assertFrameTypes(t, frames, "E")
  236. if code := errorCode(findError(frames)); code != ErrCodeDuplicateCursor {
  237. t.Fatalf("SQLSTATE = %q, want %q", code, ErrCodeDuplicateCursor)
  238. }
  239. if status := readyStatus(t, s.send(syncMessage())); status != TxStatusFailed {
  240. t.Fatalf("status after rejected re-Bind = %c, want %c", status, TxStatusFailed)
  241. }
  242. if status := readyStatus(t, s.send(queryMessageData("ROLLBACK"))); status != TxStatusIdle {
  243. t.Fatalf("ROLLBACK status = %c, want %c", status, TxStatusIdle)
  244. }
  245. frames = s.send(executeRowsMessage("p", 1))
  246. assertFrameTypes(t, frames, "E")
  247. if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
  248. t.Fatalf("SQLSTATE after ROLLBACK = %q, want %q", code, ErrCodeUndefinedCursor)
  249. }
  250. s.send(syncMessage())
  251. }
  252. // Inside an explicit transaction a portal survives Sync and keeps resuming;
  253. // COMMIT and ROLLBACK end the transaction and with it every portal.
  254. func TestPortalLivesUntilTransactionEnd(t *testing.T) {
  255. for _, ending := range []string{"COMMIT", "ROLLBACK"} {
  256. t.Run(ending, func(t *testing.T) {
  257. s := newPortalSession(t)
  258. createNumberedTable(t, s, 3)
  259. if status := readyStatus(t, s.send(queryMessageData("BEGIN"))); status != TxStatusInBlock {
  260. t.Fatalf("BEGIN status = %c, want %c", status, TxStatusInBlock)
  261. }
  262. frames := s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
  263. assertFrameTypes(t, frames, "12")
  264. frames = s.send(executeRowsMessage("p", 2), syncMessage())
  265. assertFrameTypes(t, frames, "TDDsZ")
  266. if status := readyStatus(t, frames); status != TxStatusInBlock {
  267. t.Fatalf("Sync inside the transaction reported status %c, want %c", status, TxStatusInBlock)
  268. }
  269. frames = s.send(executeRowsMessage("p", 0), syncMessage())
  270. assertFrameTypes(t, frames, "DCZ")
  271. assertRowIDs(t, frames, "3")
  272. if status := readyStatus(t, s.send(queryMessageData(ending))); status != TxStatusIdle {
  273. t.Fatalf("%s status = %c, want %c", ending, status, TxStatusIdle)
  274. }
  275. frames = s.send(executeRowsMessage("p", 1))
  276. assertFrameTypes(t, frames, "E")
  277. if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
  278. t.Fatalf("SQLSTATE after %s = %q, want %q", ending, code, ErrCodeUndefinedCursor)
  279. }
  280. s.send(syncMessage())
  281. })
  282. }
  283. }
  284. // Outside an explicit block the portal belongs to the implicit transaction that
  285. // Sync ends, so Execute afterwards reports 34000 (PostgreSQL behaves the same);
  286. // the prepared statement survives and can be bound again.
  287. func TestPortalEndsAtSyncOutsideTransaction(t *testing.T) {
  288. s := newPortalSession(t)
  289. createNumberedTable(t, s, 3)
  290. frames := s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
  291. assertFrameTypes(t, frames, "12")
  292. frames = s.send(executeRowsMessage("p", 1), syncMessage())
  293. assertFrameTypes(t, frames, "TDsZ")
  294. frames = s.send(executeRowsMessage("p", 1))
  295. assertFrameTypes(t, frames, "E")
  296. if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
  297. t.Fatalf("SQLSTATE = %q, want %q", code, ErrCodeUndefinedCursor)
  298. }
  299. s.send(syncMessage())
  300. frames = s.send(bindResultMessage("p", "stmt", nil, nil), executeRowsMessage("p", 0))
  301. assertFrameTypes(t, frames, "2TDDDC")
  302. assertRowIDs(t, frames, "1", "2", "3")
  303. }
  304. // A suspended INSERT ... RETURNING runs the statement exactly once: the resume
  305. // streams the prefetched rows, and a completed portal never runs it again.
  306. func TestBoundedExecuteDMLEvaluatesOnce(t *testing.T) {
  307. s := newPortalSession(t)
  308. s.setup("CREATE TABLE t (id INTEGER PRIMARY KEY)")
  309. frames := s.send(
  310. parseMessage("stmt", "INSERT INTO t (id) VALUES (1), (2), (3) RETURNING id"),
  311. bindResultMessage("p", "stmt", nil, nil),
  312. )
  313. assertFrameTypes(t, frames, "12")
  314. frames = s.send(executeRowsMessage("p", 2))
  315. assertFrameTypes(t, frames, "TDDs")
  316. assertRowIDs(t, frames, "1", "2")
  317. frames = s.send(executeRowsMessage("p", 0))
  318. // The autocommit INSERT publishes its commit LSN with the completion.
  319. assertFrameTypes(t, frames, "DCS")
  320. assertRowIDs(t, frames, "3")
  321. if tag := commandTagOf(frames[len(frames)-2]); tag != "INSERT 0 3" {
  322. t.Fatalf("completion tag = %q, want %q", tag, "INSERT 0 3")
  323. }
  324. frames = s.send(executeRowsMessage("p", 0))
  325. assertFrameTypes(t, frames, "C")
  326. if rows := s.countRows("t"); rows != 3 {
  327. t.Fatalf("table has %d rows after three Executes, want 3", rows)
  328. }
  329. }
  330. // Per-portal result formats keep applying on every Execute of the same portal,
  331. // and only the first Execute carries the RowDescription.
  332. func TestPortalResumeKeepsBinaryResultFormats(t *testing.T) {
  333. s := newPortalSession(t)
  334. s.setup(
  335. "CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT)",
  336. "INSERT INTO t VALUES (1, 'one')",
  337. "INSERT INTO t VALUES (2, 'two')",
  338. )
  339. frames := s.send(
  340. parseMessage("stmt", "SELECT id, name FROM t ORDER BY id"),
  341. bindResultMessage("p", "stmt", nil, []int16{1}),
  342. )
  343. assertFrameTypes(t, frames, "12")
  344. frames = s.send(executeRowsMessage("p", 1))
  345. assertFrameTypes(t, frames, "TDs")
  346. descriptions := 0
  347. for _, message := range frames {
  348. if message.Type == MsgRowDescription {
  349. descriptions++
  350. formats := rowDescriptionFormats(t, message.Data)
  351. if len(formats) != 2 || formats[0] != 1 || formats[1] != 1 {
  352. t.Fatalf("RowDescription formats = %v, want binary for both columns", formats)
  353. }
  354. }
  355. }
  356. if descriptions != 1 {
  357. t.Fatalf("got %d RowDescriptions, want 1", descriptions)
  358. }
  359. rows := dataRowsRaw(t, frames)
  360. assertBytes(t, "first id", rows[0][0], []byte{0, 0, 0, 1})
  361. frames = s.send(executeRowsMessage("p", 1))
  362. assertFrameTypes(t, frames, "Ds")
  363. rows = dataRowsRaw(t, frames)
  364. assertBytes(t, "second id", rows[0][0], []byte{0, 0, 0, 2})
  365. frames = s.send(executeRowsMessage("p", 1))
  366. assertFrameTypes(t, frames, "C")
  367. if tag := commandTagOf(frames[0]); tag != "SELECT 2" {
  368. t.Fatalf("completion tag = %q, want %q", tag, "SELECT 2")
  369. }
  370. }
  371. // Execute of a portal that was never bound is 34000, like PostgreSQL's
  372. // "portal does not exist".
  373. func TestExecuteUnknownPortalIsUndefinedCursor(t *testing.T) {
  374. s := newPortalSession(t)
  375. frames := s.send(executeRowsMessage("missing", 0))
  376. assertFrameTypes(t, frames, "E")
  377. if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
  378. t.Fatalf("SQLSTATE = %q, want %q", code, ErrCodeUndefinedCursor)
  379. }
  380. s.send(syncMessage())
  381. }
  382. // A failed Execute leaves the transaction aborted: a later Execute yields 25P02
  383. // until ROLLBACK, which also destroys the portals.
  384. func TestFailedExecuteLeavesTransactionConsistent(t *testing.T) {
  385. s := newPortalSession(t)
  386. s.setup("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT NOT NULL)")
  387. if status := readyStatus(t, s.send(queryMessageData("BEGIN"))); status != TxStatusInBlock {
  388. t.Fatalf("BEGIN status = %c, want %c", status, TxStatusInBlock)
  389. }
  390. s.send(
  391. parseMessage("failing", "INSERT INTO t (id, name) VALUES (1, NULL)"),
  392. bindResultMessage("bad", "failing", nil, nil),
  393. )
  394. frames := s.send(executeRowsMessage("bad", 0))
  395. assertFrameTypes(t, frames, "E")
  396. if code := errorCode(findError(frames)); code != "23502" {
  397. t.Fatalf("SQLSTATE = %q, want 23502", code)
  398. }
  399. if status := readyStatus(t, s.send(syncMessage())); status != TxStatusFailed {
  400. t.Fatalf("status after failure = %c, want %c", status, TxStatusFailed)
  401. }
  402. s.send(
  403. parseMessage("reading", portalSelectSQL),
  404. bindResultMessage("ok", "reading", nil, nil),
  405. )
  406. frames = s.send(executeRowsMessage("ok", 0))
  407. assertFrameTypes(t, frames, "E")
  408. if code := errorCode(findError(frames)); code != ErrCodeTransactionAborted {
  409. t.Fatalf("SQLSTATE in the aborted transaction = %q, want %q", code, ErrCodeTransactionAborted)
  410. }
  411. if status := readyStatus(t, s.send(syncMessage())); status != TxStatusFailed {
  412. t.Fatalf("status after aborted Execute = %c, want %c", status, TxStatusFailed)
  413. }
  414. if status := readyStatus(t, s.send(queryMessageData("ROLLBACK"))); status != TxStatusIdle {
  415. t.Fatalf("ROLLBACK status = %c, want %c", status, TxStatusIdle)
  416. }
  417. frames = s.send(executeRowsMessage("ok", 0))
  418. assertFrameTypes(t, frames, "E")
  419. if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
  420. t.Fatalf("SQLSTATE after ROLLBACK = %q, want %q", code, ErrCodeUndefinedCursor)
  421. }
  422. s.send(syncMessage())
  423. }
  424. // A ROLLBACK TO SAVEPOINT can undo writes a suspended portal already
  425. // materialized, so it drops every portal. Resuming one would stream rows and a
  426. // completion tag from a removed state; the client gets 34000 instead. Clearing
  427. // on ROLLBACK TO SAVEPOINT is stricter than PostgreSQL, which keeps portals
  428. // created before the savepoint.
  429. func TestRollbackToSavepointClearsSuspendedPortals(t *testing.T) {
  430. s := newPortalSession(t)
  431. s.setup("CREATE TABLE t (id INTEGER PRIMARY KEY)")
  432. if status := readyStatus(t, s.send(queryMessageData("BEGIN"))); status != TxStatusInBlock {
  433. t.Fatalf("BEGIN status = %c, want %c", status, TxStatusInBlock)
  434. }
  435. if status := readyStatus(t, s.send(queryMessageData("SAVEPOINT s"))); status != TxStatusInBlock {
  436. t.Fatalf("SAVEPOINT status = %c, want %c", status, TxStatusInBlock)
  437. }
  438. s.send(
  439. parseMessage("stmt", "INSERT INTO t (id) VALUES (1), (2) RETURNING id"),
  440. bindResultMessage("p", "stmt", nil, nil),
  441. )
  442. frames := s.send(executeRowsMessage("p", 1))
  443. assertFrameTypes(t, frames, "TDs")
  444. assertRowIDs(t, frames, "1")
  445. if status := readyStatus(t, s.send(queryMessageData("ROLLBACK TO SAVEPOINT s"))); status != TxStatusInBlock {
  446. t.Fatalf("ROLLBACK TO SAVEPOINT status = %c, want %c", status, TxStatusInBlock)
  447. }
  448. frames = s.send(executeRowsMessage("p", 1))
  449. assertFrameTypes(t, frames, "E")
  450. if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
  451. t.Fatalf("SQLSTATE = %q, want %q", code, ErrCodeUndefinedCursor)
  452. }
  453. s.send(syncMessage())
  454. if rows := s.countRows("t"); rows != 0 {
  455. t.Fatalf("table has %d rows after ROLLBACK TO SAVEPOINT, want 0", rows)
  456. }
  457. if status := readyStatus(t, s.send(queryMessageData("ROLLBACK"))); status != TxStatusIdle {
  458. t.Fatalf("ROLLBACK status = %c, want %c", status, TxStatusIdle)
  459. }
  460. }
  461. // Transaction commands executed through a portal emit their tag once and drop
  462. // every open portal, so a suspended sibling can never replay rows or a
  463. // completion tag from a transaction that ended or was rewound.
  464. func TestTransactionCommandPortalClearsSiblings(t *testing.T) {
  465. for _, tc := range []struct {
  466. name string
  467. command string
  468. wantTag string
  469. wantStatus byte
  470. }{
  471. {name: "commit", command: "COMMIT", wantTag: "COMMIT", wantStatus: TxStatusIdle},
  472. {name: "rollback", command: "ROLLBACK", wantTag: "ROLLBACK", wantStatus: TxStatusIdle},
  473. {name: "rollback_to_savepoint", command: "ROLLBACK TO SAVEPOINT s", wantTag: "ROLLBACK", wantStatus: TxStatusInBlock},
  474. } {
  475. t.Run(tc.name, func(t *testing.T) {
  476. s := newPortalSession(t)
  477. createNumberedTable(t, s, 3)
  478. if status := readyStatus(t, s.send(queryMessageData("BEGIN"))); status != TxStatusInBlock {
  479. t.Fatalf("BEGIN status = %c, want %c", status, TxStatusInBlock)
  480. }
  481. if status := readyStatus(t, s.send(queryMessageData("SAVEPOINT s"))); status != TxStatusInBlock {
  482. t.Fatalf("SAVEPOINT status = %c, want %c", status, TxStatusInBlock)
  483. }
  484. s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("first", "stmt", nil, nil))
  485. assertFrameTypes(t, s.send(executeRowsMessage("first", 1)), "TDs")
  486. s.send(bindResultMessage("second", "stmt", nil, nil))
  487. assertFrameTypes(t, s.send(executeRowsMessage("second", 1)), "TDs")
  488. s.send(parseMessage("txn", tc.command), bindResultMessage("txn_portal", "txn", nil, nil))
  489. frames := s.send(executeRowsMessage("txn_portal", 0))
  490. assertFrameTypes(t, frames, "C")
  491. if tag := commandTagOf(frames[0]); tag != tc.wantTag {
  492. t.Fatalf("tag = %q, want %q", tag, tc.wantTag)
  493. }
  494. if status := readyStatus(t, s.send(syncMessage())); status != tc.wantStatus {
  495. t.Fatalf("status after %s = %c, want %c", tc.command, status, tc.wantStatus)
  496. }
  497. for _, name := range []string{"first", "second", "txn_portal"} {
  498. frames = s.send(executeRowsMessage(name, 1))
  499. assertFrameTypes(t, frames, "E")
  500. if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
  501. t.Fatalf("Execute %q after %s: SQLSTATE = %q, want %q", name, tc.command, code, ErrCodeUndefinedCursor)
  502. }
  503. s.send(syncMessage())
  504. }
  505. })
  506. }
  507. }
  508. // BEGIN through a portal starts the block and emits its tag. It does not end a
  509. // transaction, so sibling portals are kept, and a repeated Execute repeats the
  510. // stored tag without replaying BEGIN (the engine rejects a second BEGIN, so a
  511. // replay would surface as an error).
  512. func TestBeginPortalStartsTransaction(t *testing.T) {
  513. s := newPortalSession(t)
  514. createNumberedTable(t, s, 3)
  515. s.send(parseMessage("begin_stmt", "BEGIN"), bindResultMessage("begin_portal", "begin_stmt", nil, nil))
  516. frames := s.send(executeRowsMessage("begin_portal", 0))
  517. assertFrameTypes(t, frames, "C")
  518. if tag := commandTagOf(frames[0]); tag != "BEGIN" {
  519. t.Fatalf("tag = %q, want %q", tag, "BEGIN")
  520. }
  521. if status := readyStatus(t, s.send(syncMessage())); status != TxStatusInBlock {
  522. t.Fatalf("status after BEGIN = %c, want %c", status, TxStatusInBlock)
  523. }
  524. s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
  525. assertFrameTypes(t, s.send(executeRowsMessage("p", 1)), "TDs")
  526. frames = s.send(executeRowsMessage("begin_portal", 0))
  527. assertFrameTypes(t, frames, "C")
  528. if tag := commandTagOf(frames[0]); tag != "BEGIN" {
  529. t.Fatalf("repeated tag = %q, want %q", tag, "BEGIN")
  530. }
  531. frames = s.send(executeRowsMessage("p", 0))
  532. assertFrameTypes(t, frames, "DDC")
  533. assertRowIDs(t, frames, "2", "3")
  534. if status := readyStatus(t, s.send(queryMessageData("ROLLBACK"))); status != TxStatusIdle {
  535. t.Fatalf("ROLLBACK status = %c, want %c", status, TxStatusIdle)
  536. }
  537. }