sql: allow transaction diagnostics to work with non-fingerprinted statements

Certain statements can be executed within a transaction, but are not recorded
as part of the transaction fingerprint. In these cases, the transaction
diagnostics recorded should record these statements and continue with
with the diagnostics collection instead of resetting the instrumentation
helper.

As part of this fix, the transaction diagnostics collector now has more
states it could be in, including Reset, NotStarted, InProgress, and
ReadyToFinalize. These states make it easier to allow or dissalllow certain
actions, such as continuing diagnostics, adding statement bundles,
finalizing a diagnostics collection, and starting a diagnostics collection.

Epic: CRDB-53541
Release note: None

Author: kyle-a-wong

pkg/sql/conn_executor.go

@@ -3499,12 +3499,6 @@ func isCommit(stmt tree.Statement) bool {
 	return ok
 }
 
-// isRollback returns true if stmt is a "ROLLBACK" statement.
-func isRollback(stmt tree.Statement) bool {
-	_, ok := stmt.(*tree.RollbackTransaction)
-	return ok
-}
-
 var retryableMinTimestampBoundUnsatisfiableError = errors.Newf(
 	"retryable MinTimestampBoundUnsatisfiableError",
 )

pkg/sql/instrumentation.go

@@ -736,7 +736,7 @@ func (ih *instrumentationHelper) Finish(
 				payloadErr, retErr, &p.extendedEvalCtx.Settings.SV, ih.inFlightTraceCollector,
 			)
 
-			if !txnHelper.DiagnosticsInProgress() {
+			if !txnHelper.diagnosticsCollector.InProgress() {
 				// Include all non-critical errors as warnings. Note that these
 				// error strings might contain PII, but the warnings are only shown
 				// to the current user and aren't included into the bundle.
@@ -749,7 +749,7 @@ func (ih *instrumentationHelper) Finish(
 		}
 	}
 
-	if txnHelper.DiagnosticsInProgress() {
+	if txnHelper.diagnosticsCollector.InProgress() {
 		// If we're collecting a transaction bundle, add the statement to the
 		// current txn diagnostic bundle. These will be persisted at the end
 		// of the transaction if the transaction matches the request.

pkg/sql/stmtdiagnostics/statement_diagnostics_test.go

@@ -845,9 +845,110 @@ LIMIT 1`)
 	}
 
 	registry := tc.Server(0).ExecutorConfig().(sql.ExecutorConfig).TxnDiagnosticsRecorder
+	t.Run("diagnostics", func(t *testing.T) {
+		for _, testCase := range []struct {
+			name       string
+			statements []string
+		}{
+			{
+				name:       "regular request",
+				statements: txnStatements,
+			},
+			{
+				name: "with savepoint cockroach_restart",
+				statements: []string{
+					"SAVEPOINT cockroach_restart",
+					"SELECT 'aaaaaa', 'bbbbbb'",
+					"SELECT 'aaaaaa', 'bbbbbb', 'ccccc' UNION select 'ddddd', 'eeeee', 'fffff'",
+					"RELEASE SAVEPOINT cockroach_restart",
+				},
+			},
+			{
+				name: "with savepoint",
+				statements: []string{
+					"SAVEPOINT my_savepoint",
+					"SELECT 'aaaaaa', 'bbbbbb'",
+					"SELECT 'aaaaaa', 'bbbbbb', 'ccccc' UNION select 'ddddd', 'eeeee', 'fffff'",
+					"RELEASE SAVEPOINT my_savepoint",
+				},
+			},
+			{
+				name: "with savepoint rollback",
+				statements: []string{
+					"SAVEPOINT my_savepoint",
+					"SELECT 'aaaaaa', 'bbbbbb'",
+					"SELECT 'aaaaaa', 'bbbbbb', 'ccccc' UNION select 'ddddd', 'eeeee', 'fffff'",
+					"ROLLBACK TO SAVEPOINT my_savepoint",
+				},
+			},
+			{
+				name: "with cockroach_restart and regular savepoint",
+				statements: []string{
+					"SAVEPOINT cockroach_restart",
+					"SAVEPOINT my_savepoint",
+					"SELECT 'aaaaaa', 'bbbbbb'",
+					"RELEASE SAVEPOINT my_savepoint",
+					"SELECT 'aaaaaa', 'bbbbbb', 'ccccc' UNION select 'ddddd', 'eeeee', 'fffff'",
+					"RELEASE SAVEPOINT cockroach_restart",
+				},
+			},
+			{
+				name: "with cockroach_restart and regular savepoint and rollback",
+				statements: []string{
+					"SAVEPOINT cockroach_restart",
+					"SELECT 'aaaaaa', 'bbbbbb'",
+					"SAVEPOINT my_savepoint",
+					"ROLLBACK TO SAVEPOINT my_savepoint",
+					"SELECT 'aaaaaa', 'bbbbbb', 'ccccc' UNION select 'ddddd', 'eeeee', 'fffff'",
+					"SHOW SAVEPOINT STATUS",
+					"RELEASE SAVEPOINT cockroach_restart",
+				},
+			},
+			{
+				name: "with show commit timestamp",
+				statements: []string{
+					"SELECT 'aaaaaa', 'bbbbbb'",
+					"SELECT 'aaaaaa', 'bbbbbb', 'ccccc' UNION select 'ddddd', 'eeeee', 'fffff'",
+					"SHOW COMMIT TIMESTAMP",
+				},
+			},
+		} {
+			t.Run(testCase.name, func(t *testing.T) {
+				// Insert a request for the transaction fingerprint
+				reqId, err := registry.InsertTxnRequestInternal(
+					ctx,
+					txnFingerprintId,
+					statementFingerprintIds,
+					"testuser",
+					0,
+					0,
+					0,
+					false,
+				)
+				require.NoError(t, err)
 
-	t.Run("txn diagnostic request", func(t *testing.T) {
-		// Insert a request for the transaction fingerprint
+				// Ensure that the request shows up in the system.transaction_diagnostics_requests table
+				var count int
+				sqlConn.QueryRow(t, "SELECT count(*) FROM system.transaction_diagnostics_requests WHERE completed=false and id=$1", reqId).Scan(&count)
+				require.Equal(t, 1, count)
+
+				// Execute the transaction until the request is marked as complete.
+				// We use a connection to a different server than the diagnostics request
+				// was made on to ensure that the request is properly gossiped.
+				runTxnUntilDiagnosticsCollected(t, tc, sqlutils.MakeSQLRunner(tc.ServerConn(1)), testCase.statements, reqId)
+			})
+		}
+	})
+
+	t.Run("diagnostics with prepared statement", func(t *testing.T) {
+		// Since you can't create the same named prepared statement, we have this
+		// in a separate test that is guaranteed to complete on the first try
+		// by using the same node the request was made on.
+		statements := []string{
+			"PREPARE my_statement as SELECT 'aaaaaa', 'bbbbbb'",
+			"EXECUTE my_statement",
+			"SELECT 'aaaaaa', 'bbbbbb', 'ccccc' UNION select 'ddddd', 'eeeee', 'fffff'",
+		}
 		reqId, err := registry.InsertTxnRequestInternal(
 			ctx,
 			txnFingerprintId,
@@ -859,19 +960,8 @@ LIMIT 1`)
 			false,
 		)
 		require.NoError(t, err)
-
-		// Ensure that the request shows up in the system.transaction_diagnostics_requests table
-		var count int
-		sqlConn.QueryRow(t, "SELECT count(*) FROM system.transaction_diagnostics_requests WHERE completed=false and id=$1", reqId).Scan(&count)
-		require.Equal(t, 1, count)
-
-		// Execute the transaction until the request is marked as complete.
-		// We use a connection to a different server than the diagnostics request
-		// was made on to ensure that the request is properly gossiped.
-		runTxnUntilDiagnosticsCollected(t, tc, sqlutils.MakeSQLRunner(tc.ServerConn(1)), txnStatements, reqId)
-
+		runTxnUntilDiagnosticsCollected(t, tc, sqlConn, statements, reqId)
 	})
-
 	t.Run("txn diagnostic with statement diagnostic request", func(t *testing.T) {
 
 		stmtReqId, err := registry.StmtRegistry.InsertRequestInternal(