release-22.2: sql: fix extra round trips during DISCARD, fix temp discard bug

pkg/bench/rttanalysis/BUILD.bazel

@@ -38,6 +38,7 @@ go_test(
         "alter_table_bench_test.go",
         "bench_test.go",
         "create_alter_role_bench_test.go",
+        "discard_bench_test.go",
         "drop_bench_test.go",
         "grant_revoke_bench_test.go",
         "grant_revoke_role_bench_test.go",

pkg/bench/rttanalysis/discard_bench_test.go

@@ -0,0 +1,53 @@
+// 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 rttanalysis
+
+import "testing"
+
+func BenchmarkDiscard(b *testing.B) { reg.Run(b) }
+func init() {
+	reg.Register("Discard", []RoundTripBenchTestCase{
+		{
+			Name: "DISCARD ALL, no tables",
+			Setup: `DROP DATABASE IF EXISTS db CASCADE; 
+CREATE DATABASE db;`,
+			Stmt: "DISCARD ALL",
+		},
+		{
+			Name: "DISCARD ALL, 1 tables in 1 db",
+			Setup: `
+SET experimental_enable_temp_tables = true;
+DROP DATABASE IF EXISTS db CASCADE; 
+CREATE DATABASE db; 
+USE db; 
+CREATE TEMPORARY TABLE foo(i int);
+`,
+			Stmt: "DISCARD ALL",
+		},
+		{
+			Name: "DISCARD ALL, 2 tables in 2 dbs",
+			Setup: `
+SET experimental_enable_temp_tables = true;
+DROP DATABASE IF EXISTS db CASCADE; 
+CREATE DATABASE db; 
+USE db; 
+CREATE TEMPORARY TABLE foo(i int);
+CREATE TEMPORARY TABLE bar(i int);
+DROP DATABASE IF EXISTS db2 CASCADE; 
+CREATE DATABASE db2; 
+USE db2; 
+CREATE TEMPORARY TABLE foo(i int);
+CREATE TEMPORARY TABLE bar(i int);
+`,
+			Stmt: "DISCARD ALL",
+		},
+	})
+}

pkg/bench/rttanalysis/testdata/benchmark_expectations

@@ -31,6 +31,9 @@ exp,benchmark
 16,CreateRole/create_role_with_2_options
 19,CreateRole/create_role_with_3_options
 14,CreateRole/create_role_with_no_options
+18,"Discard/DISCARD_ALL,_1_tables_in_1_db"
+46,"Discard/DISCARD_ALL,_2_tables_in_2_dbs"
+0,"Discard/DISCARD_ALL,_no_tables"
 14,DropDatabase/drop_database_0_tables
 16,DropDatabase/drop_database_1_table
 18,DropDatabase/drop_database_2_tables

pkg/sql/discard.go

@@ -86,6 +86,11 @@ func (n *discardNode) startExec(params runParams) error {
 }
 
 func deleteTempTables(ctx context.Context, p *planner) error {
+	// If this session has no temp schemas, then there is nothing to do here.
+	// This is the common case.
+	if len(p.SessionData().DatabaseIDToTempSchemaID) == 0 {
+		return nil
+	}
 	return p.WithInternalExecutor(ctx, func(ctx context.Context, txn *kv.Txn, ie sqlutil.InternalExecutor) error {
 		codec := p.execCfg.Codec
 		descCol := p.Descriptors()

pkg/sql/internal.go

@@ -811,11 +811,17 @@ func (ie *InternalExecutor) execInternal(
 		sd.ApplicationName = catconstants.InternalAppNamePrefix + "-" + opName
 	}
 	// If the caller has injected a mapping to temp schemas, install it, and
-	// leave it installed for the rest of the transaction.
+	// leave it installed until the statement concludes.
+	resetTemporarySchemaProvider := func() {}
 	if ie.extraTxnState != nil && sd.DatabaseIDToTempSchemaID != nil {
 		ie.extraTxnState.descCollection.SetTemporaryDescriptors(
 			descs.NewTemporarySchemaProvider(sessiondata.NewStack(sd)),
 		)
+		resetTemporarySchemaProvider = func() {
+			ie.extraTxnState.descCollection.SetTemporaryDescriptors(
+				descs.NewTemporarySchemaProvider(ie.sessionDataStack),
+			)
+		}
 	}
 
 	// The returned span is finished by this function in all error paths, but if
@@ -834,6 +840,7 @@ func (ie *InternalExecutor) execInternal(
 		// TODO(knz): track the callers and check whether opName could be turned
 		// into a type safe for reporting.
 		if retErr != nil || r == nil {
+			resetTemporarySchemaProvider()
 			// Both retErr and r can be nil in case of panic.
 			if retErr != nil && !errIsRetriable(retErr) {
 				retErr = errors.Wrapf(retErr, "%s", opName)
@@ -843,6 +850,7 @@ func (ie *InternalExecutor) execInternal(
 			sp.Finish()
 		} else {
 			r.errCallback = func(err error) error {
+				resetTemporarySchemaProvider()
 				if err != nil && !errIsRetriable(err) {
 					err = errors.Wrapf(err, "%s", opName)
 				}

pkg/sql/logictest/testdata/logic_test/temp_table

@@ -347,4 +347,44 @@ to_drop.pg_temp.testuser_tmp
 to_drop.pg_temp.tempuser_view
 to_drop.pg_temp.root_temp
 
+subtest create_after_discard_and_drop_database
+
+statement ok
+ALTER ROLE testuser WITH CREATEDB;
+
 user testuser
+
+statement ok
+CREATE DATABASE to_drop
+
+statement ok
+USE to_drop
+
+statement ok
+CREATE TEMPORARY TABLE t (i INT PRIMARY KEY);
+
+statement ok
+SELECT * FROM pg_temp.t
+
+statement ok
+DISCARD TEMP
+
+statement error pgcode 42P01 relation "pg_temp.t" does not exist
+SELECT * FROM pg_temp.t
+
+statement ok
+USE defaultdb;
+DROP DATABASE to_drop CASCADE;
+
+statement ok
+CREATE DATABASE to_drop
+
+statement ok
+CREATE TEMPORARY TABLE t (i INT PRIMARY KEY);
+
+statement ok
+SELECT * FROM pg_temp.t
+
+statement ok
+USE defaultdb;
+DROP DATABASE to_drop CASCADE;

pkg/sql/user.go

@@ -563,7 +563,7 @@ func (p *planner) setRole(ctx context.Context, local bool, s username.SQLUsernam
 	sessionUser := p.SessionData().SessionUser()
 	becomeUser := sessionUser
 	// Check the role exists - if so, populate becomeUser.
-	if !s.IsNoneRole() {
+	if !s.IsNoneRole() && s != sessionUser {
 		becomeUser = s
 
 		exists, err := p.RoleExists(ctx, becomeUser)