plpgsql: add support for COMMIT and ROLLBACK statements

[DrewKimball]

plpgsql: add parser support for COMMIT and ROLLBACK

This commit adds support in the PL/pgSQL parser for transaction control
statements. It also performs some minor refactoring of the AST statements.

Informs #115294

Release note: None

sql: add support for transaction control during SQL execution

This commit adds a method to commit or abort a transaction from within
a SQL execution plan, through a TxnControlExpr. A TxnControlExpr wraps
a "continuation" routine. Upon evaluation, it directs the connExecutor to
commit/abort the transaction, and stores a plan built for the continuation.
The connExecutor will finish the current transaction, and then use the
stored plan to resume execution in the new transaction. This will be used
in the following commit to implement the PL/pgSQL COMMIT and ROLLBACK
statements.

Informs #115294

Release note: None

plpgsql: add support for COMMIT and ROLLBACK

This commit adds support for the PL/pgSQL COMMIT and ROLLBACK statements,
which allow a stored procedure to finish the current transaction and
resume execution from the new transaction. This behavior is implemented
through the TxnControl expression from the previous commit, which
passes the necessary information (commit or rollback, continuation) to
the connExecutor before exiting. Note that these statements can only
be used from a stored procedure.

Fixes #115294

Release note (sql change): Added support for the PL/pgSQL COMMIT and
ROLLBACK statements.

[blathers-crl]

Your pull request contains more than 1000 changes. It is strongly encouraged to split big PRs into smaller chunks.

_{🦉 Hoot! I am a Blathers, a bot for CockroachDB. My owner is dev-inf.}

[cockroach-teamcity]

This change is 

[yuzefovich]

Great work! I have some nits and questions.

I don't think I saw any special logic / tests for these parts of the postgres docs:

Transaction control is only possible in CALL or DO invocations from the top level or nested CALL or DO invocations without any other intervening command. For example, if the call stack is CALL proc1() → CALL proc2() → CALL proc3(), then the second and third procedures can perform transaction control actions. But if the call stack is CALL proc1() → SELECT func2() → CALL proc3(), then the last procedure cannot do transaction control, because of the SELECT in between.

Transaction commands are not allowed in cursor loops driven by commands that are not read-only (for example UPDATE ... RETURNING).

Are these not implemented yet?

Reviewed 5 of 5 files at r1, 15 of 15 files at r2, 30 of 30 files at r3, all commit messages.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @DrewKimball and @mgartner)

---

pkg/sql/conn_executor.go line 2622 at r2 (raw file):

		// Do not reset the "resume" plan. This will allow a sub-transaction
		// within a stored procedure to be retried individually from any previous or
		// following transactions within the stored procedure..

nit: double period.

---

pkg/sql/conn_fsm.go line 577 at r2 (raw file):

// given event.
func (ts *txnState) finishTxn(ev txnEventType) error {
	finishedTxnID, commitTimestamp := ts.finishSQLTxn()

nit: it might be nice to extract the shared helper from finishTxn and finishTxnNoAdvance in case in the future we need to update the common logic.

---

pkg/sql/opt/optbuilder/plpgsql.go line 236 at r3 (raw file):

// buildRootBlock builds a PL/pgSQL routine starting with the root block.
func (b *plpgsqlBuilder) buildRootBlock(astBlock *ast.Block, s *scope) *scope {
	// Push the scope, since otherwise the routine parameters could be considered

nit: some of the changes here are included in #119608 - I just approved that PR for ease of rebase.

---

pkg/sql/opt/optbuilder/plpgsql.go line 251 at r3 (raw file):

		}
	}
	return b.buildBlock(astBlock, s.push())

nit: IIUC we don't need to push the scope here since we've already pushed above, right?

---

pkg/sql/opt/optbuilder/plpgsql.go line 1781 at r3 (raw file):

		return stmt, false
	}
	return stmt, true

nit: perhaps we can speed up the traversal a bit by doing:

return stmt, !tc.foundTxnControlStatement

---

pkg/sql/sem/tree/routine.go line 284 at r2 (raw file):

	Typ  *types.T
}

nit: add something like var _ tree.Expr = &TxnControlExpr{}.

---

pkg/sql/sem/tree/routine.go line 312 at r2 (raw file):

// Format is part of the Expr interface.
func (node *TxnControlExpr) Format(ctx *FmtCtx) {
	if node.Op == StoredProcTxnCommit {

nit: should we add an test-build assertion that node.Op is not a no-op?

---

pkg/sql/sem/tree/routine.go line 317 at r2 (raw file):

		ctx.Printf("ROLLBACK; CALL ")
	}
	ctx.Printf("%s(", node.Name)

nit: perhaps include CALL string here rather than above in two places (it seems a bit nicer this way).

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 167 at r3 (raw file):
Is this part of postgres docs implemented?

However, a cursor created as part of a loop like this is automatically converted to a holdable cursor by the first COMMIT or ROLLBACK. That means that the cursor is fully evaluated at the first COMMIT or ROLLBACK rather than row by row.

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 478 at r3 (raw file):

subtest function

nit: perhaps add a comment here saying that postgres creates the function but then runs into an error during the execution time since it defers the checks till execution.

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 601 at r3 (raw file):

statement ok
SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL READ COMMITTED;

nit: do you think that we can somehow verify that txn properties of the session txn are propagated to all nested txns? Or is this only applicable with ON CHAIN?

---

pkg/sql/plpgsql/parser/testdata/stmt_commit_rollback line 1 at r1 (raw file):

error

nit: what is the error that we expect now? I'm surprised that we don't have the expected error in the test file.

[DrewKimball]

TFYR!

Are these not implemented yet?

Yeah, none of these are implemented yet:
UDF/SP calling other UDF/SP: #88198
DO blocks: #107345
Cursor loops: #115296

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @mgartner and @yuzefovich)

---

pkg/sql/conn_executor.go line 2622 at r2 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: double period.

Done.

---

pkg/sql/conn_fsm.go line 577 at r2 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: it might be nice to extract the shared helper from finishTxn and finishTxnNoAdvance in case in the future we need to update the common logic.

Good idea. I removed finishTxnNoAdvance, and added a new parameter to finishTxn to indicate the advance behavior.

---

pkg/sql/opt/optbuilder/plpgsql.go line 236 at r3 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: some of the changes here are included in #119608 - I just approved that PR for ease of rebase.

Thanks! I'll make sure to rebase this once that gets merged.

---

pkg/sql/opt/optbuilder/plpgsql.go line 251 at r3 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: IIUC we don't need to push the scope here since we've already pushed above, right?

Good catch. It's harmless, but unnecessary.

---

pkg/sql/opt/optbuilder/plpgsql.go line 1781 at r3 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: perhaps we can speed up the traversal a bit by doing:

return stmt, !tc.foundTxnControlStatement

Nice idea, Done.

---

pkg/sql/sem/tree/routine.go line 284 at r2 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: add something like var _ tree.Expr = &TxnControlExpr{}.

Done.

---

pkg/sql/sem/tree/routine.go line 312 at r2 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: should we add an test-build assertion that node.Op is not a no-op?

Good idea, done.

---

pkg/sql/sem/tree/routine.go line 317 at r2 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: perhaps include CALL string here rather than above in two places (it seems a bit nicer this way).

Done. I also modified it to use the StoredProcTxnOp.String method.

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 167 at r3 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

Is this part of postgres docs implemented?

However, a cursor created as part of a loop like this is automatically converted to a holdable cursor by the first COMMIT or ROLLBACK. That means that the cursor is fully evaluated at the first COMMIT or ROLLBACK rather than row by row.

Cursor loops aren't implemented yet. Issue here: #115296

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 478 at r3 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: perhaps add a comment here saying that postgres creates the function but then runs into an error during the execution time since it defers the checks till execution.

Done.

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 601 at r3 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: do you think that we can somehow verify that txn properties of the session txn are propagated to all nested txns? Or is this only applicable with ON CHAIN?

Hm, I don't think there's any way to check that within a stored proc right now. But I did add checks after the proc executes, and you inspired me to also add some (buffered) logging to verify that variable assignment isn't duplicated by internal retries.

---

pkg/sql/plpgsql/parser/testdata/stmt_commit_rollback line 1 at r1 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: what is the error that we expect now? I'm surprised that we don't have the expected error in the test file.

Ah, I forgot to replace those headers with parse. Fixed it now.

[yuzefovich]

Reviewed 46 of 46 files at r4, 16 of 16 files at r5, 30 of 30 files at r6, all commit messages.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @DrewKimball and @mgartner)

---

pkg/sql/opt/exec/execbuilder/scalar.go line 1025 at r6 (raw file):

	ctx *buildScalarCtx, routineArgs memo.ScalarListExpr,
) (args tree.TypedExprs, err error) {
	if len(routineArgs) > 0 {

nit: might be cleaner to negate the condition and return early on zero args.

---

pkg/sql/opt/exec/execbuilder/scalar.go line 1296 at r6 (raw file):

			memoArgs[i] = f.ConstructConstVal(evalArgs[i], evalArgs[i].ResolvedType())
		}
		continuationProc := f.ConstructUDFCall(memoArgs, &memo.UDFCallPrivate{Def: txnExpr.Def})

nit: can we ever set TailCall: true?

---

pkg/sql/opt/optbuilder/plpgsql.go line 774 at r6 (raw file):

				panic(txnControlWithChainErr)
			}
			// NOTE: postgres doesn't make the following checks until runtime.

How important is this difference do you think? Seems reasonable to me - I'm assuming that "the following checks" only refers to two panics below, right? In other words, are there cases where, perhaps depending on the arguments, we might return an error (because we eagerly perform some checks) whereas postgres wouldn't return an error (because the error path is not actually hit during execution)?

Update: indeed, there are some cases like this - using your example:

CREATE PROCEDURE p(a INT) LANGUAGE PLpgSQL AS $$
  BEGIN
    IF a > 0 THEN
      COMMIT;
    END IF;
  EXCEPTION WHEN division_by_zero THEN
    RAISE NOTICE 'foo';
  END;
$$;

Calling this procedure with non-positive argument doesn't trigger an error in postgres.

We probably should document this difference. Would it be difficult to mirror the behavior of postgres here?

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 368 at r6 (raw file):

DROP PROCEDURE p;

# COMMIT is not valid in a block with an exception handler.

nit: perhaps add a comment here as well about the difference with postgres.

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 590 at r6 (raw file):

3  3

subtest read_committed_retry

nit: since we now have read committed logic test config as part of the default configs, we can probably remove some of the duplication.

[DrewKimball]

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @mgartner and @yuzefovich)

---

pkg/sql/opt/exec/execbuilder/scalar.go line 1025 at r6 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: might be cleaner to negate the condition and return early on zero args.

Good idea, Done.

---

pkg/sql/opt/exec/execbuilder/scalar.go line 1296 at r6 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: can we ever set TailCall: true?

It wouldn't do any harm, but it wouldn't have any effect either. TailCall describes a routine's position inside of a parent routine, so what does it mean to set it for a routine that isn't nested inside another?

---

pkg/sql/opt/optbuilder/plpgsql.go line 774 at r6 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

How important is this difference do you think? Seems reasonable to me - I'm assuming that "the following checks" only refers to two panics below, right? In other words, are there cases where, perhaps depending on the arguments, we might return an error (because we eagerly perform some checks) whereas postgres wouldn't return an error (because the error path is not actually hit during execution)?

Update: indeed, there are some cases like this - using your example:

CREATE PROCEDURE p(a INT) LANGUAGE PLpgSQL AS $$
  BEGIN
    IF a > 0 THEN
      COMMIT;
    END IF;
  EXCEPTION WHEN division_by_zero THEN
    RAISE NOTICE 'foo';
  END;
$$;

Calling this procedure with non-positive argument doesn't trigger an error in postgres.

We probably should document this difference. What it be difficult to mirror the behavior of postgres here?

It wouldn't be too hard to make the error lazy with a RAISE statement similar to end-of-function errors, but it would add some complication, and doesn't seem worth it when a user will probably never want to do this. I've opened a tracking issue here: #119750

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 368 at r6 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: perhaps add a comment here as well about the difference with postgres.

Done.

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 590 at r6 (raw file):

Previously, yuzefovich (Yahor Yuzefovich) wrote…

nit: since we now have read committed logic test config as part of the default configs, we can probably remove some of the duplication.

Oh nice, that's convenient. I've removed the read-committed subtest.

[yuzefovich]

Great work!

Reviewed 44 of 44 files at r7, 16 of 16 files at r8, 28 of 28 files at r9, all commit messages.
Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @DrewKimball and @mgartner)

---

pkg/sql/opt/optbuilder/plpgsql.go line 774 at r6 (raw file):

Previously, DrewKimball (Drew Kimball) wrote…

It wouldn't be too hard to make the error lazy with a RAISE statement similar to end-of-function errors, but it would add some complication, and doesn't seem worth it when a user will probably never want to do this. I've opened a tracking issue here: #119750

Sounds good to me.

[mgartner]

First commit looks great, still working through the other two.

Reviewed 4 of 5 files at r1, 1 of 46 files at r4, 44 of 44 files at r7.
Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @DrewKimball)

---

pkg/sql/plpgsql/parser/plpgsql.y line 666 at r7 (raw file):

| stmt_while
  {
    $$.val = $1.statement()

Did you mean to add this in this commit?

---

pkg/sql/plpgsql/parser/testdata/stmt_commit_rollback line 1 at r7 (raw file):

parse

Do we still need this test file if we have the commit and rollback ones?

[DrewKimball]

Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @mgartner)

---

pkg/sql/plpgsql/parser/plpgsql.y line 666 at r7 (raw file):

Previously, mgartner (Marcus Gartner) wrote…

Did you mean to add this in this commit?

Yes, it seemed simple enough not to make its own commit, but I'm happy to separate it out if you'd prefer.

---

pkg/sql/plpgsql/parser/testdata/stmt_commit_rollback line 1 at r7 (raw file):

Previously, mgartner (Marcus Gartner) wrote…

Do we still need this test file if we have the commit and rollback ones?

Good point, I hadn't noticed it was a dupe.

[mgartner]

Very nice!

Reviewed 46 of 46 files at r10, 16 of 16 files at r11, 28 of 28 files at r12, all commit messages.
Reviewable status: complete! 1 of 0 LGTMs obtained (and 1 stale) (waiting on @DrewKimball)

---

pkg/sql/conn_executor.go line 2617 at r11 (raw file):

		// Do not reset the "resume" plan. This will allow a sub-transaction
		// within a stored procedure to be retried individually from any previous or
		// following transactions within the stored procedure.

In this case, will we eventually hit the default case below? I want to make sure that we don't leave a pointer to the memo in this struct after the SP executes.

---

pkg/sql/plpgsql/parser/plpgsql.y line 666 at r7 (raw file):

Previously, DrewKimball (Drew Kimball) wrote…

Yes, it seemed simple enough not to make its own commit, but I'm happy to separate it out if you'd prefer.

No need to—we'll still get an unsupported error if someone tries to use a while loop, right?

---

pkg/sql/sem/tree/routine.go line 261 at r11 (raw file):

		return "ROLLBACK"
	default:
		return "NO-OP"

nit: might prevent future confusion when adding a value if you make the default "unknown" and the no-op case explicit.

---

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn line 61 at r12 (raw file):

    RAISE NOTICE 'ROLLBACK;';
    ROLLBACK;
    RAISE NOTICE 'max: %', (SELECT max(x) FROM t);

Cool!!!

[DrewKimball]

Reviewable status: complete! 1 of 0 LGTMs obtained (and 1 stale) (waiting on @mgartner)

---

pkg/sql/conn_executor.go line 2617 at r11 (raw file):

Previously, mgartner (Marcus Gartner) wrote…

In this case, will we eventually hit the default case below? I want to make sure that we don't leave a pointer to the memo in this struct after the SP executes.

Yes, unless we have an "infinite retry" bug, in which case a lot of other things won't be (and shouldn't be) cleaned up either. I added a clarifying note to the comment.

---

pkg/sql/plpgsql/parser/plpgsql.y line 666 at r7 (raw file):

Previously, mgartner (Marcus Gartner) wrote…

No need to—we'll still get an unsupported error if someone tries to use a while loop, right?

WHILE loops actually are supported, it just happened to work before because the interface "union" val was never overwritten. This change won't cause a user-visible change in behavior, but will make sure we don't break anything due to grammar changes down the line (which seems unlikely, but best to be safe I think).

---

pkg/sql/sem/tree/routine.go line 261 at r11 (raw file):

Previously, mgartner (Marcus Gartner) wrote…

nit: might prevent future confusion when adding a value if you make the default "unknown" and the no-op case explicit.

Done.

[blathers-crl]

Your pull request contains more than 1000 changes. It is strongly encouraged to split big PRs into smaller chunks.

Thank you for updating your pull request.

Before a member of our team reviews your PR, I have some potential action items for you:

• We notice you have more than one commit in your PR. We try break logical changes into separate commits, but commits such as "fix typo" or "address review commits" should be squashed into one commit and pushed with --force
• Please ensure your git commit message contains a release note.
• When CI has completed, please ensure no errors have appeared.

_{🦉 Hoot! I am a Blathers, a bot for CockroachDB. My owner is dev-inf.}

[cockroach-teamcity]

All committers have signed the CLA.

[DrewKimball]

TFYRs!

bors r+

[craig]

Build succeeded:

• Bazel Essential CI (Cockroach)
• macos_amd64_build
• macos_arm64_build
• windows_build