From 37eec07e4a38a04966219c6f93b8d9790ae35dbc Mon Sep 17 00:00:00 2001 From: Jane Xing Date: Mon, 27 Jun 2022 17:39:18 -0400 Subject: [PATCH] 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 https://github.com/cockroachdb/cockroach/issues/69495 Release note: None --- pkg/sql/planner.go | 95 ++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 88 insertions(+), 7 deletions(-) diff --git a/pkg/sql/planner.go b/pkg/sql/planner.go index e061b25816dd..08591c305c6f 100644 --- a/pkg/sql/planner.go +++ b/pkg/sql/planner.go @@ -20,6 +20,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" @@ -37,6 +38,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" @@ -291,7 +293,7 @@ func NewInternalPlanner( // Returns a cleanup function that must be called once the caller is done with // the planner. func newInternalPlanner( -// TODO(yuzefovich): make this redact.RedactableString. + // TODO(yuzefovich): make this redact.RedactableString. opName string, txn *kv.Txn, user username.SQLUsername, @@ -841,6 +843,20 @@ func validateDescriptor(ctx context.Context, p *planner, descriptor catalog.Desc ) } +// 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(ctx, p.SessionData()) + ie.(*InternalExecutor).extraTxnState = &extraTxnState{ + txn: p.Txn(), + descCollection: p.Descriptors(), + jobs: p.extendedEvalCtx.Jobs, + schemaChangeJobRecords: p.extendedEvalCtx.SchemaChangeJobRecords, + } + 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. // @@ -853,8 +869,21 @@ func (p *planner) QueryRowEx( stmt string, qargs ...interface{}, ) (tree.Datums, error) { - ie := p.ExecCfg().InternalExecutorFactory(ctx, p.SessionData()) - return ie.QueryRowEx(ctx, opName, p.Txn(), override, stmt, qargs...) + ie := initInternalExecutor(ctx, p) + return ie.QueryRowExUpdated(ctx, opName, 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.ExecExUpdated(ctx, opName, override, stmt, qargs...) } // QueryIteratorEx executes the query, returning an iterator that can be used @@ -870,14 +899,66 @@ func (p *planner) QueryIteratorEx( stmt string, qargs ...interface{}, ) (eval.InternalRows, error) { - ie := p.ExecCfg().InternalExecutorFactory(ctx, p.SessionData()) - rows, err := ie.QueryIteratorEx(ctx, opName, p.Txn(), override, stmt, qargs...) + ie := initInternalExecutor(ctx, p) + rows, err := ie.QueryIteratorExUpdated(ctx, opName, override, stmt, qargs...) return rows.(eval.InternalRows), err } +// 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.QueryBufferedExUpdated(ctx, opName, 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.QueryRowExWithColsUpdated(ctx, opName, 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.QueryBufferedExWithColsUpdated(ctx, opName, 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) +} + // extraTxnState is to store extra transaction state info that -// will be passed to an internal executor when it's used under a planner -// context. It should not be exported from the sql package. +// will be passed to an internal executor when it's used under a txn context. +// It should not be exported from the sql package. type extraTxnState struct { txn *kv.Txn descCollection *descs.Collection