sql: introduce query functions under sql.planner

Currently, the internal executor always create its own descriptor collections,
txn state, job collection and etc. for its conn executor, even though it's run
underneath a "parent" query. These recreation can unneccesarily reduce the
query efficiency in some use cases, such as when an internal executor is
used under a planner context. In this case, the internal executor is expected
to inherit these info from the planner, rather than creating its own.

To make this rule more explicit, this commit adds a series of query functions
under `sql.planner`. Each of these functions wrap both the init of an internal
executor and the query execution. In this way, the internal executor always
stores the info inherited from the parent planner, and will pass it to its child
conn executor.

fixes cockroachdb#69495

Release note: None

pkg/bench/rttanalysis/testdata/benchmark_expectations

@@ -67,11 +67,11 @@ exp,benchmark
 2,ORMQueries/has_table_privilege_1
 4,ORMQueries/has_table_privilege_3
 6,ORMQueries/has_table_privilege_5
-15,ORMQueries/information_schema._pg_index_position
+3,ORMQueries/information_schema._pg_index_position
 2,ORMQueries/pg_attribute
 2,ORMQueries/pg_class
 11,ORMQueries/pg_is_other_temp_schema
-23,ORMQueries/pg_is_other_temp_schema_multiple_times
+18,ORMQueries/pg_is_other_temp_schema_multiple_times
 4,ORMQueries/pg_my_temp_schema
 4,ORMQueries/pg_my_temp_schema_multiple_times
 4,ORMQueries/pg_namespace

pkg/sql/planner.go

@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/spanconfig"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/lease"
@@ -39,6 +40,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats/persistedsqlstats"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlutil"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/upgrade"
 	"github.com/cockroachdb/cockroach/pkg/util/cancelchecker"
@@ -889,6 +891,26 @@ func validateDescriptor(ctx context.Context, p *planner, descriptor catalog.Desc
 	)
 }
 
+// IsActive implements the Planner interface.
+func (p *planner) IsActive(ctx context.Context, key clusterversion.Key) bool {
+	return p.execCfg.Settings.Version.IsActive(ctx, key)
+}
+
+// initInternalExecutor is to initialize an internal executor with a planner.
+// Note that this function should only be used when using internal executor
+// to run sql statement under the planner context.
+func initInternalExecutor(ctx context.Context, p *planner) sqlutil.InternalExecutor {
+	ie := p.ExecCfg().InternalExecutorFactory.NewInternalExecutor(p.SessionData())
+	ie.(*InternalExecutor).extraTxnState = &extraTxnState{
+		txn:                    p.Txn(),
+		descCollection:         p.Descriptors(),
+		jobs:                   p.extendedEvalCtx.Jobs,
+		schemaChangeJobRecords: p.extendedEvalCtx.SchemaChangeJobRecords,
+		schemaChangerState:     p.extendedEvalCtx.SchemaChangerState,
+	}
+	return ie
+}
+
 // QueryRowEx executes the supplied SQL statement and returns a single row, or
 // nil if no row is found, or an error if more that one row is returned.
 //
@@ -901,10 +923,23 @@ func (p *planner) QueryRowEx(
 	stmt string,
 	qargs ...interface{},
 ) (tree.Datums, error) {
-	ie := p.ExecCfg().InternalExecutorFactory.NewInternalExecutor(p.SessionData())
+	ie := initInternalExecutor(ctx, p)
 	return ie.QueryRowEx(ctx, opName, p.Txn(), override, stmt, qargs...)
 }
 
+// ExecEx is like Exec, but allows the caller to override some session data
+// fields (e.g. the user).
+func (p *planner) ExecEx(
+	ctx context.Context,
+	opName string,
+	override sessiondata.InternalExecutorOverride,
+	stmt string,
+	qargs ...interface{},
+) (int, error) {
+	ie := initInternalExecutor(ctx, p)
+	return ie.ExecEx(ctx, opName, p.Txn(), override, stmt, qargs...)
+}
+
 // QueryIteratorEx executes the query, returning an iterator that can be used
 // to get the results. If the call is successful, the returned iterator
 // *must* be closed.
@@ -918,12 +953,59 @@ func (p *planner) QueryIteratorEx(
 	stmt string,
 	qargs ...interface{},
 ) (eval.InternalRows, error) {
-	ie := p.ExecCfg().InternalExecutorFactory.NewInternalExecutor(p.SessionData())
+	ie := initInternalExecutor(ctx, p)
 	rows, err := ie.QueryIteratorEx(ctx, opName, p.Txn(), override, stmt, qargs...)
 	return rows.(eval.InternalRows), err
 }
 
-// IsActive implements the Planner interface.
-func (p *planner) IsActive(ctx context.Context, key clusterversion.Key) bool {
-	return p.execCfg.Settings.Version.IsActive(ctx, key)
+// QueryBufferedEx executes the supplied SQL statement and returns the resulting
+// rows (meaning all of them are buffered at once).
+// The fields set in session that are set override the respective fields if they
+// have previously been set through SetSessionData().
+func (p *planner) QueryBufferedEx(
+	ctx context.Context,
+	opName string,
+	session sessiondata.InternalExecutorOverride,
+	stmt string,
+	qargs ...interface{},
+) ([]tree.Datums, error) {
+	ie := initInternalExecutor(ctx, p)
+	return ie.QueryBufferedEx(ctx, opName, p.Txn(), session, stmt, qargs...)
+}
+
+// QueryRowExWithCols is like QueryRowEx, additionally returning the computed
+// ResultColumns of the input query.
+func (p *planner) QueryRowExWithCols(
+	ctx context.Context,
+	opName string,
+	session sessiondata.InternalExecutorOverride,
+	stmt string,
+	qargs ...interface{},
+) (tree.Datums, colinfo.ResultColumns, error) {
+	ie := initInternalExecutor(ctx, p)
+	return ie.QueryRowExWithCols(ctx, opName, p.Txn(), session, stmt, qargs...)
+}
+
+// QueryBufferedExWithCols is like QueryBufferedEx, additionally returning the
+// computed ResultColumns of the input query.
+func (p *planner) QueryBufferedExWithCols(
+	ctx context.Context,
+	opName string,
+	session sessiondata.InternalExecutorOverride,
+	stmt string,
+	qargs ...interface{},
+) ([]tree.Datums, colinfo.ResultColumns, error) {
+	ie := initInternalExecutor(ctx, p)
+	return ie.QueryBufferedExWithCols(ctx, opName, p.Txn(), session, stmt, qargs...)
+}
+
+// WithInternalExecutor let user run multiple sql statements within the same
+// internal executor initialized under a planner context. To run single sql
+// statements, please use the query functions above.
+func (p *planner) WithInternalExecutor(
+	ctx context.Context,
+	run func(ctx context.Context, txn *kv.Txn, ie sqlutil.InternalExecutor) error,
+) error {
+	ie := initInternalExecutor(ctx, p)
+	return run(ctx, p.Txn(), ie)
 }