github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/conn_executor_savepoints_test.go

// Copyright 2020 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package sql_test

import (
	"context"
	"fmt"
	"strings"
	"testing"

	"github.com/cockroachdb/cockroach/pkg/base"
	"github.com/cockroachdb/cockroach/pkg/sql"
	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
	"github.com/cockroachdb/cockroach/pkg/util"
	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
	"github.com/cockroachdb/datadriven"
)

func TestSavepoints(t *testing.T) {
	defer leaktest.AfterTest(t)()

	ctx := context.Background()
	datadriven.Walk(t, "testdata/savepoints", func(t *testing.T, path string) {

		params := base.TestServerArgs{}
		s, sqlConn, _ := serverutils.StartServer(t, params)
		defer s.Stopper().Stop(ctx)

		if _, err := sqlConn.Exec("CREATE TABLE progress(n INT, marker BOOL)"); err != nil {
			t.Fatal(err)
		}

		datadriven.RunTest(t, path, func(t *testing.T, td *datadriven.TestData) string {
			switch td.Cmd {
			case "sql":
				// Implicitly abort any previously-ongoing txn.
				_, _ = sqlConn.Exec("ABORT")
				// Prepare for the next test.
				if _, err := sqlConn.Exec("DELETE FROM progress"); err != nil {
					td.Fatalf(t, "cleaning up: %v", err)
				}

				// Prepare a buffer to accumulate the results.
				var buf strings.Builder

				// We're going to execute the input line-by-line.
				stmts := strings.Split(td.Input, "\n")

				// progressBar is going to show the cancellation of writes
				// during rollbacks.
				progressBar := make([]byte, len(stmts))
				erase := func(status string) {
					char := byte('.')
					if !isOpenTxn(status) {
						char = 'X'
					}
					for i := range progressBar {
						progressBar[i] = char
					}
				}

				// stepNum is the index of the current statement
				// in the input.
				var stepNum int

				// updateProgress loads the current set of writes
				// into the progress bar.
				updateProgress := func() {
					rows, err := sqlConn.Query("SELECT n FROM progress")
					if err != nil {
						t.Logf("%d: reading progress: %v", stepNum, err)
						// It's OK if we can't read this.
						return
					}
					defer rows.Close()
					for rows.Next() {
						var n int
						if err := rows.Scan(&n); err != nil {
							td.Fatalf(t, "%d: unexpected error while reading progress: %v", stepNum, err)
						}
						if n < 1 || n > len(progressBar) {
							td.Fatalf(t, "%d: unexpected stepnum in progress table: %d", stepNum, n)
						}
						progressBar[n-1] = '#'
					}
				}

				// getTxnStatus retrieves the current txn state.
				// This is guaranteed to always succeed because SHOW TRANSACTION STATUS
				// is an observer statement.
				getTxnStatus := func() string {
					row := sqlConn.QueryRow("SHOW TRANSACTION STATUS")
					var status string
					if err := row.Scan(&status); err != nil {
						td.Fatalf(t, "%d: unable to retrieve txn status: %v", stepNum, err)
					}
					return status
				}
				// showSavepointStatus is like getTxnStatus but retrieves the
				// savepoint stack.
				showSavepointStatus := func() {
					rows, err := sqlConn.Query("SHOW SAVEPOINT STATUS")
					if err != nil {
						td.Fatalf(t, "%d: unable to retrieve savepoint status: %v", stepNum, err)
					}
					defer rows.Close()

					comma := ""
					hasSavepoints := false
					for rows.Next() {
						var name string
						var isRestart bool
						if err := rows.Scan(&name, &isRestart); err != nil {
							td.Fatalf(t, "%d: unexpected error while reading savepoints: %v", stepNum, err)
						}
						if isRestart {
							name += "(r)"
						}
						buf.WriteString(comma)
						buf.WriteString(name)
						hasSavepoints = true
						comma = ">"
					}
					if !hasSavepoints {
						buf.WriteString("(none)")
					}
				}
				// report shows the progress of execution so far after
				// each statement executed.
				report := func(beforeStatus, afterStatus string) {
					erase(afterStatus)
					if isOpenTxn(afterStatus) {
						updateProgress()
					}
					fmt.Fprintf(&buf, "-- %-11s -> %-11s %s ", beforeStatus, afterStatus, string(progressBar))
					buf.WriteByte(' ')
					showSavepointStatus()
					buf.WriteByte('\n')
				}

				// The actual execution of the statements starts here.

				beforeStatus := getTxnStatus()
				for i, stmt := range stmts {
					stepNum = i + 1
					// Before each statement, mark the progress so far with
					// a KV write.
					if isOpenTxn(beforeStatus) {
						_, err := sqlConn.Exec("INSERT INTO progress(n, marker) VALUES ($1, true)", stepNum)
						if err != nil {
							td.Fatalf(t, "%d: before-stmt: %v", stepNum, err)
						}
					}

					// Run the statement and report errors/results.
					fmt.Fprintf(&buf, "%d: %s -- ", stepNum, stmt)
					execRes, err := sqlConn.Exec(stmt)
					if err != nil {
						fmt.Fprintf(&buf, "%v\n", err)
					} else {
						nRows, err := execRes.RowsAffected()
						if err != nil {
							fmt.Fprintf(&buf, "error retrieving rows: %v\n", err)
						} else {
							fmt.Fprintf(&buf, "%d row%s\n", nRows, util.Pluralize(nRows))
						}
					}

					// Report progress on the next line
					afterStatus := getTxnStatus()
					report(beforeStatus, afterStatus)
					beforeStatus = afterStatus
				}

				return buf.String()

			default:
				td.Fatalf(t, "unknown directive: %s", td.Cmd)
			}
			return ""
		})
	})
}

func isOpenTxn(status string) bool {
	return status == sql.OpenStateStr || status == sql.NoTxnStateStr
}