Merge #62617

62617: sql, server: add SQLStatsResetter to distsql and internalExecutor r=asubiotto,yuzefovich a=Azhng Closes #62587. Release note: None Co-authored-by: Azhng <[email protected]>
cockroachdb · Mar 26, 2021 · 8b552c7 · 8b552c7
2 parents fd5d337 + d43090e
commit 8b552c7
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 0 deletions.
diff --git a/pkg/server/server_sql.go b/pkg/server/server_sql.go
@@ -599,6 +599,8 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 		execCfg,
 	)
 
+	distSQLServer.ServerConfig.SQLStatsResetter = pgServer.SQLServer
+
 	// Now that we have a pgwire.Server (which has a sql.Server), we can close a
 	// circular dependency between the rowexec.Server and sql.Server and set
 	// SessionBoundInternalExecutorFactory. The same applies for setting a

diff --git a/pkg/sql/distsql/server.go b/pkg/sql/distsql/server.go
@@ -299,6 +299,7 @@ func (ds *ServerImpl) setupFlow(
 			InternalExecutor:   ie,
 			Txn:                leafTxn,
 			SQLLivenessReader:  ds.ServerConfig.SQLLivenessReader,
+			SQLStatsResetter:   ds.ServerConfig.SQLStatsResetter,
 		}
 		evalCtx.SetStmtTimestamp(timeutil.Unix(0 /* sec */, req.EvalContext.StmtTimestampNanos))
 		evalCtx.SetTxnTimestamp(timeutil.Unix(0 /* sec */, req.EvalContext.TxnTimestampNanos))

diff --git a/pkg/sql/execinfra/server_config.go b/pkg/sql/execinfra/server_config.go
@@ -30,6 +30,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/hydratedtables"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlutil"
 	"github.com/cockroachdb/cockroach/pkg/storage/cloud"
@@ -151,6 +152,10 @@ type ServerConfig struct {
 	// HydratedTables is a node-level cache of table descriptors which utilize
 	// user-defined types.
 	HydratedTables *hydratedtables.Cache
+
+	// SQLStatsResetter is an interface used to reset SQL stats without the need to
+	// introduce dependency on the sql package.
+	SQLStatsResetter tree.SQLStatsResetter
 }
 
 // RuntimeStats is an interface through which the rowexec layer can get

diff --git a/pkg/sql/logictest/testdata/logic_test/distsql_crdb_internal b/pkg/sql/logictest/testdata/logic_test/distsql_crdb_internal
@@ -0,0 +1,48 @@
+# LogicTest: 5node-default-configs
+
+statement ok
+CREATE TABLE data (a INT, b INT, c FLOAT, d DECIMAL, PRIMARY KEY (a, b, c, d))
+
+# Split into ten parts.
+statement ok
+ALTER TABLE data SPLIT AT SELECT i FROM generate_series(1, 9) AS g(i)
+
+# Relocate the ten parts to the five nodes.
+statement ok
+ALTER TABLE data EXPERIMENTAL_RELOCATE
+  SELECT ARRAY[i%5+1], i FROM generate_series(0, 9) AS g(i)
+
+# Generate all combinations of values 1 to 10.
+statement ok
+INSERT INTO data SELECT a, b, c::FLOAT, d::DECIMAL FROM
+   generate_series(1, 10) AS a(a),
+   generate_series(1, 10) AS b(b),
+   generate_series(1, 10) AS c(c),
+   generate_series(1, 10) AS d(d)
+
+# Verify data placement.
+query TTTI colnames,rowsort
+SELECT start_key, end_key, replicas, lease_holder FROM [SHOW RANGES FROM TABLE data]
+----
+start_key  end_key  replicas  lease_holder
+NULL       /1       {1}       1
+/1         /2       {2}       2
+/2         /3       {3}       3
+/3         /4       {4}       4
+/4         /5       {5}       5
+/5         /6       {1}       1
+/6         /7       {2}       2
+/7         /8       {3}       3
+/8         /9       {4}       4
+/9         NULL     {5}       5
+
+
+query II
+SELECT a, count(*) AS cnt FROM data GROUP BY a HAVING crdb_internal.reset_sql_stats() ORDER BY a LIMIT 5
+----
+1 1000
+2 1000
+3 1000
+4 1000
+5 1000
+
diff --git a/pkg/sql/planner.go b/pkg/sql/planner.go
@@ -414,6 +414,11 @@ func internalExtendedEvalCtx(
 ) extendedEvalContext {
 	evalContextTestingKnobs := execCfg.EvalContextTestingKnobs
 
+	var sqlStatsResetter tree.SQLStatsResetter
+	if execCfg.InternalExecutor != nil {
+		sqlStatsResetter = execCfg.InternalExecutor.s
+	}
+
 	return extendedEvalContext{
 		EvalContext: tree.EvalContext{
 			Txn:              txn,
@@ -428,6 +433,7 @@ func internalExtendedEvalCtx(
 			StmtTimestamp:    stmtTimestamp,
 			TxnTimestamp:     txnTimestamp,
 			InternalExecutor: execCfg.InternalExecutor,
+			SQLStatsResetter: sqlStatsResetter,
 		},
 		SessionMutator:    dataMutator,
 		VirtualSchemas:    execCfg.VirtualSchemas,