server: proper transaction state management in sql-over-http

We need to construct the internal executor in the context of the transaction so that we can make sure that its side-effects are properly managed. Without this change, we'd be throwing away all of the extraTxnState between each statement. We'd fail to create the jobs (which we defer to the end of the transaction), and we'd fail to run those jobs and check for errors. We'd also fail to validate the two-version invariant or wait for one version. Fixes #86332 Release justification: Fixes critical bugs in new functionality. Release note: None
cockroachdb · Aug 20, 2022 · 4812378 · 4812378
1 parent f266aa2
commit 4812378
Show file tree

Hide file tree

Showing 2 changed files with 203 additions and 5 deletions.
diff --git a/pkg/server/api_v2_sql.go b/pkg/server/api_v2_sql.go
@@ -21,6 +21,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -360,13 +361,27 @@ func (a *apiV2Server) execSQL(w http.ResponseWriter, r *http.Request) {
 	// runner is the function that will execute all the statements as a group.
 	// If there's just one statement, we execute them with an implicit,
 	// auto-commit transaction.
-	runner := func(ctx context.Context, fn func(context.Context, *kv.Txn) error) error { return fn(ctx, nil) }
+
+	type (
+		txnFunc    = func(context.Context, *kv.Txn, sqlutil.InternalExecutor) error
+		runnerFunc = func(ctx context.Context, fn txnFunc) error
+	)
+	var runner runnerFunc
 	if len(requestPayload.Statements) > 1 {
 		// We need a transaction to group the statements together.
 		// We use TxnWithSteppingEnabled here even though we don't
 		// use stepping below, because that buys us admission control.
-		runner = func(ctx context.Context, fn func(context.Context, *kv.Txn) error) error {
-			return a.admin.server.db.TxnWithSteppingEnabled(ctx, sessiondatapb.Normal, fn)
+		cf := a.admin.server.sqlServer.execCfg.CollectionFactory
+		runner = func(ctx context.Context, fn txnFunc) error {
+			return cf.TxnWithExecutor(ctx, a.admin.server.db, nil, func(
+				ctx context.Context, txn *kv.Txn, _ *descs.Collection, ie sqlutil.InternalExecutor,
+			) error {
+				return fn(ctx, txn, ie)
+			}, descs.SteppingEnabled())
+		}
+	} else {
+		runner = func(ctx context.Context, fn func(context.Context, *kv.Txn, sqlutil.InternalExecutor) error) error {
+			return fn(ctx, nil, a.admin.ie)
 		}
 	}
 
@@ -376,7 +391,7 @@ func (a *apiV2Server) execSQL(w http.ResponseWriter, r *http.Request) {
 	err = contextutil.RunWithTimeout(ctx, "run-sql-via-api", timeout, func(ctx context.Context) error {
 		retryNum := 0
 
-		return runner(ctx, func(ctx context.Context, txn *kv.Txn) error {
+		return runner(ctx, func(ctx context.Context, txn *kv.Txn, ie sqlutil.InternalExecutor) error {
 			result.Execution.TxnResults = result.Execution.TxnResults[:0]
 			result.Execution.Retries = retryNum
 			retryNum++
@@ -413,7 +428,7 @@ func (a *apiV2Server) execSQL(w http.ResponseWriter, r *http.Request) {
 						}
 					}()
 
-					it, err := a.admin.ie.QueryIteratorEx(ctx, "run-query-via-api", txn,
+					it, err := ie.QueryIteratorEx(ctx, "run-query-via-api", txn,
 						sessiondata.InternalExecutorOverride{
 							User:            username,
 							Database:        requestPayload.Database,

diff --git a/pkg/server/testdata/api_v2_sql b/pkg/server/testdata/api_v2_sql
@@ -271,3 +271,186 @@ sql admin expect-error
 }
 ----
 XXUUU|parsing statement 1: expected 1 placeholder(s), got 2
+
+sql admin
+{
+  "database": "mydb",
+  "execute": true,
+  "statements": [{"sql": "CREATE TABLE foo (i INT PRIMARY KEY, j INT UNIQUE)"}]
+}
+----
+{
+ "execution": {
+  "txn_results": [
+   {
+    "columns": [
+     {
+      "name": "rows_affected",
+      "oid": 20,
+      "type": "INT8"
+     }
+    ],
+    "end": "1970-01-01T00:00:00Z",
+    "rows_affected": 0,
+    "start": "1970-01-01T00:00:00Z",
+    "statement": 1,
+    "tag": "CREATE TABLE"
+   }
+  ]
+ },
+ "num_statements": 1
+}
+
+sql admin
+{
+  "database": "mydb",
+  "execute": true,
+  "statements": [
+    {"sql": "ALTER TABLE foo RENAME TO bar"},
+    {"sql": "INSERT INTO bar (i) VALUES (1), (2)"},
+    {"sql": "ALTER TABLE bar DROP COLUMN j"},
+    {"sql": "ALTER TABLE bar ADD COLUMN k INT DEFAULT 42"}
+  ]
+}
+----
+{
+ "execution": {
+  "txn_results": [
+   {
+    "columns": [
+     {
+      "name": "rows_affected",
+      "oid": 20,
+      "type": "INT8"
+     }
+    ],
+    "end": "1970-01-01T00:00:00Z",
+    "rows_affected": 0,
+    "start": "1970-01-01T00:00:00Z",
+    "statement": 1,
+    "tag": "ALTER TABLE"
+   },
+   {
+    "columns": [
+     {
+      "name": "rows_affected",
+      "oid": 20,
+      "type": "INT8"
+     }
+    ],
+    "end": "1970-01-01T00:00:00Z",
+    "rows_affected": 2,
+    "start": "1970-01-01T00:00:00Z",
+    "statement": 2,
+    "tag": "INSERT"
+   },
+   {
+    "columns": [
+     {
+      "name": "rows_affected",
+      "oid": 20,
+      "type": "INT8"
+     }
+    ],
+    "end": "1970-01-01T00:00:00Z",
+    "rows_affected": 0,
+    "start": "1970-01-01T00:00:00Z",
+    "statement": 3,
+    "tag": "ALTER TABLE"
+   },
+   {
+    "columns": [
+     {
+      "name": "rows_affected",
+      "oid": 20,
+      "type": "INT8"
+     }
+    ],
+    "end": "1970-01-01T00:00:00Z",
+    "rows_affected": 0,
+    "start": "1970-01-01T00:00:00Z",
+    "statement": 4,
+    "tag": "ALTER TABLE"
+   }
+  ]
+ },
+ "num_statements": 4
+}
+
+sql admin
+{
+  "database": "mydb",
+  "execute": true,
+  "statements": [
+    {"sql": "SELECT * FROM bar"}
+  ]
+}
+----
+{
+ "execution": {
+  "txn_results": [
+   {
+    "columns": [
+     {
+      "name": "i",
+      "oid": 20,
+      "type": "INT8"
+     },
+     {
+      "name": "k",
+      "oid": 20,
+      "type": "INT8"
+     }
+    ],
+    "end": "1970-01-01T00:00:00Z",
+    "rows": [
+     {
+      "i": 1,
+      "k": 42
+     },
+     {
+      "i": 2,
+      "k": 42
+     }
+    ],
+    "rows_affected": 0,
+    "start": "1970-01-01T00:00:00Z",
+    "statement": 1,
+    "tag": "SELECT"
+   }
+  ]
+ },
+ "num_statements": 1
+}
+
+
+sql admin
+{
+  "database": "mydb",
+  "execute": true,
+  "statements": [
+    {"sql": "DROP TABLE bar"}
+  ]
+}
+----
+{
+ "execution": {
+  "txn_results": [
+   {
+    "columns": [
+     {
+      "name": "rows_affected",
+      "oid": 20,
+      "type": "INT8"
+     }
+    ],
+    "end": "1970-01-01T00:00:00Z",
+    "rows_affected": 0,
+    "start": "1970-01-01T00:00:00Z",
+    "statement": 1,
+    "tag": "DROP TABLE"
+   }
+  ]
+ },
+ "num_statements": 1
+}