Merge #156959

156959: copy: step txn unconditionally r=yuzefovich a=yuzefovich

Currently, we have an issue with COPY FROM command if it runs within the txn which had just been rolled back to a savepoint - we'd hit an assertion error in the KV layer. This commit fixes the oversight by stepping the txn, similar to what we do on the main query path.

Additionally, this commit explicitly prohibits running COPY FROM via the internal executor (previously it could lead to undefined behavior or internal errors) and adds a test with savepoints for COPY TO (which works because it uses different execution machinery than COPY FROM).

Fixes: #144383.

Release note (bug fix): Previously, CockroachDB could encounter an internal error when evaluating COPY FROM command in a transaction after it's been rolled back to a savepoint. The bug has been present since before 23.2 version and is now fixed.

Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: craig[bot]

pkg/sql/conn_executor.go

@@ -3227,6 +3227,12 @@ func (ex *connExecutor) execCopyIn(
 	// Disable the buffered writes for COPY since there is no benefit in this
 	// ability here.
 	ex.state.mu.txn.SetBufferedWritesEnabled(false /* enabled */)
+	// Step the txn in case it had just been rolled back to a savepoint (if it
+	// wasn't, this is harmless). This also matches what we do unconditionally
+	// on the main query path.
+	if err := ex.state.mu.txn.Step(ctx, false /* allowReadTimestampStep */); err != nil {
+		return ex.makeErrEvent(err, cmd.ParsedStmt.AST)
+	}
 	txnOpt := copyTxnOpt{
 		txn:           ex.state.mu.txn,
 		txnTimestamp:  ex.state.sqlTimestamp,

pkg/sql/copy/copy_in_test.go

@@ -646,49 +646,51 @@ func TestCopyTransaction(t *testing.T) {
 	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{})
 	defer s.Stopper().Stop(context.Background())
 
-	if _, err := db.Exec(`
+	_, err := db.Exec(`
 		CREATE DATABASE d;
 		SET DATABASE = d;
 		CREATE TABLE t (
 			i INT PRIMARY KEY
 		);
-	`); err != nil {
-		t.Fatal(err)
-	}
+	`)
+	require.NoError(t, err)
 
 	txn, err := db.Begin()
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	// Note that, at least with lib/pq, this doesn't actually send a Parse msg
-	// (which we wouldn't support, as we don't support Copy-in in extended
-	// protocol mode). lib/pq has magic for recognizing a Copy.
-	stmt, err := txn.Prepare(pq.CopyIn("t", "i"))
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 
-	const val = 2
+	// Run COPY twice with the first one being rolled back via savepoints.
+	for val, doSavepoint := range []bool{true, false} {
+		func() {
+			if doSavepoint {
+				_, err = txn.Exec("SAVEPOINT s")
+				require.NoError(t, err)
+				defer func() {
+					_, err = txn.Exec("ROLLBACK TO SAVEPOINT s")
+					require.NoError(t, err)
+				}()
+			}
+			// Note that, at least with lib/pq, this doesn't actually send a
+			// Parse msg (which we wouldn't support, as we don't support Copy-in
+			// in extended protocol mode). lib/pq has magic for recognizing a
+			// Copy.
+			stmt, err := txn.Prepare(pq.CopyIn("t", "i"))
+			require.NoError(t, err)
 
-	_, err = stmt.Exec(val)
-	if err != nil {
-		t.Fatal(err)
-	}
+			_, err = stmt.Exec(val)
+			require.NoError(t, err)
 
-	if err = stmt.Close(); err != nil {
-		t.Fatal(err)
-	}
+			err = stmt.Close()
+			require.NoError(t, err)
 
-	var i int
-	if err := txn.QueryRow("SELECT i FROM d.t").Scan(&i); err != nil {
-		t.Fatal(err)
-	} else if i != val {
-		t.Fatalf("expected 1, got %d", i)
-	}
-	if err := txn.Commit(); err != nil {
-		t.Fatal(err)
+			var i int
+			err = txn.QueryRow("SELECT i FROM d.t").Scan(&i)
+			require.NoError(t, err)
+			if i != val {
+				t.Fatalf("expected %d, got %d", val, i)
+			}
+		}()
 	}
+	require.NoError(t, txn.Commit())
 }
 
 // TestCopyFromFKCheck verifies that foreign keys are checked during COPY.

pkg/sql/copy/copy_out_test.go

@@ -61,14 +61,26 @@ func TestCopyOutTransaction(t *testing.T) {
 	tx, err := conn.Begin(ctx)
 	require.NoError(t, err)
 
-	_, err = tx.Exec(ctx, "INSERT INTO t VALUES (1)")
-	require.NoError(t, err)
+	for val, doSavepoint := range []bool{true, false} {
+		func() {
+			if doSavepoint {
+				_, err = tx.Exec(ctx, "SAVEPOINT s")
+				require.NoError(t, err)
+				defer func() {
+					_, err = tx.Exec(ctx, "ROLLBACK TO SAVEPOINT s")
+					require.NoError(t, err)
+				}()
+			}
 
-	var buf bytes.Buffer
-	_, err = tx.Conn().PgConn().CopyTo(ctx, &buf, "COPY t TO STDOUT")
-	require.NoError(t, err)
-	require.Equal(t, "1\n", buf.String())
+			_, err = tx.Exec(ctx, "INSERT INTO t VALUES ($1)", val)
+			require.NoError(t, err)
 
+			var buf bytes.Buffer
+			_, err = tx.Conn().PgConn().CopyTo(ctx, &buf, "COPY t TO STDOUT")
+			require.NoError(t, err)
+			require.Equal(t, fmt.Sprintf("%d\n", val), buf.String())
+		}()
+	}
 	require.NoError(t, tx.Rollback(ctx))
 }
 

pkg/sql/internal.go

@@ -1464,15 +1464,17 @@ func (ie *InternalExecutor) commitTxn(ctx context.Context) error {
 	return ex.commitSQLTransactionInternal(ctx)
 }
 
-// checkIfStmtIsAllowed returns an error if the internal executor is not bound
-// with the outer-txn-related info but is used to run DDL statements within an
-// outer txn.
-// TODO (janexing): this will be deprecate soon since it's not a good idea
-// to have `extraTxnState` to store the info from a outer txn.
+// checkIfStmtIsAllowed returns an error if the internal executor cannot execute
+// the given stmt.
 func (ie *InternalExecutor) checkIfStmtIsAllowed(stmt tree.Statement, txn *kv.Txn) error {
 	if stmt == nil {
 		return nil
 	}
+	if _, ok := stmt.(*tree.CopyFrom); ok {
+		// COPY FROM has special handling in the connExecutor, so we can't run
+		// it via the internal executor.
+		return errors.New("COPY cannot be run via the internal executor")
+	}
 	if tree.CanModifySchema(stmt) && txn != nil && ie.extraTxnState == nil {
 		return errors.New("DDL statement is disallowed if internal " +
 			"executor is not bound with txn metadata")

pkg/sql/opt/exec/execbuilder/testdata/execute_internally_builtin

@@ -298,3 +298,6 @@ subtest end
 # Regression test for crashing with SHOW COMMIT TIMESTAMP.
 statement error this statement is disallowed
 SELECT crdb_internal.execute_internally('SHOW COMMIT TIMESTAMP;', true);
+
+statement error COPY cannot be run via the internal executor
+SELECT crdb_internal.execute_internally('COPY t FROM STDIN');