diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel index e724fd3ca416..ff79ee2fb8d8 100644 --- a/pkg/BUILD.bazel +++ b/pkg/BUILD.bazel @@ -663,6 +663,7 @@ ALL_TESTS = [ "//pkg/sql/sqlstats/sslocal:sslocal_test", "//pkg/sql/sqlstats/ssmemstorage:ssmemstorage_test", "//pkg/sql/sqlstats/ssremote:ssremote_test", + "//pkg/sql/sqlstats:sqlstats_test", "//pkg/sql/sqltestutils:sqltestutils_test", "//pkg/sql/stats:stats_test", "//pkg/sql/stmtdiagnostics:stmtdiagnostics_test", @@ -2417,6 +2418,7 @@ GO_TARGETS = [ "//pkg/sql/sqlstats/ssremote:ssremote", "//pkg/sql/sqlstats/ssremote:ssremote_test", "//pkg/sql/sqlstats:sqlstats", + "//pkg/sql/sqlstats:sqlstats_test", "//pkg/sql/sqltelemetry:sqltelemetry", "//pkg/sql/sqltestutils:sqltestutils", "//pkg/sql/sqltestutils:sqltestutils_test", diff --git a/pkg/sql/sqlstats/BUILD.bazel b/pkg/sql/sqlstats/BUILD.bazel index bb58b1850fbe..55b423813130 100644 --- a/pkg/sql/sqlstats/BUILD.bazel +++ b/pkg/sql/sqlstats/BUILD.bazel @@ -1,4 +1,4 @@ -load("@io_bazel_rules_go//go:def.bzl", "go_library") +load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test") go_library( name = "sqlstats", @@ -32,3 +32,26 @@ go_library( "@com_github_cockroachdb_redact//:redact", ], ) + +go_test( + name = "sqlstats_test", + srcs = [ + "main_test.go", + "sqlstats_test.go", + ], + data = glob(["testdata/**"]), + deps = [ + ":sqlstats", + "//pkg/base", + "//pkg/security/securityassets", + "//pkg/security/securitytest", + "//pkg/server", + "//pkg/testutils/datapathutils", + "//pkg/testutils/serverutils", + "//pkg/testutils/sqlutils", + "//pkg/testutils/testcluster", + "//pkg/util/leaktest", + "//pkg/util/log", + "@com_github_cockroachdb_datadriven//:datadriven", + ], +) diff --git a/pkg/sql/sqlstats/main_test.go b/pkg/sql/sqlstats/main_test.go new file mode 100644 index 000000000000..962bb05e2197 --- /dev/null +++ b/pkg/sql/sqlstats/main_test.go @@ -0,0 +1,24 @@ +// Copyright 2025 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package sqlstats_test + +import ( + "os" + "testing" + + "github.com/cockroachdb/cockroach/pkg/security/securityassets" + "github.com/cockroachdb/cockroach/pkg/security/securitytest" + "github.com/cockroachdb/cockroach/pkg/server" + "github.com/cockroachdb/cockroach/pkg/testutils/serverutils" + "github.com/cockroachdb/cockroach/pkg/testutils/testcluster" +) + +func TestMain(m *testing.M) { + securityassets.SetLoader(securitytest.EmbeddedAssets) + serverutils.InitTestServerFactory(server.TestServerFactory) + serverutils.InitTestClusterFactory(testcluster.TestClusterFactory) + os.Exit(m.Run()) +} diff --git a/pkg/sql/sqlstats/sqlstats_test.go b/pkg/sql/sqlstats/sqlstats_test.go new file mode 100644 index 000000000000..f1f5776cbb05 --- /dev/null +++ b/pkg/sql/sqlstats/sqlstats_test.go @@ -0,0 +1,140 @@ +// Copyright 2025 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package sqlstats_test + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/cockroachdb/cockroach/pkg/base" + "github.com/cockroachdb/cockroach/pkg/sql/sqlstats" + "github.com/cockroachdb/cockroach/pkg/testutils/datapathutils" + "github.com/cockroachdb/cockroach/pkg/testutils/serverutils" + "github.com/cockroachdb/cockroach/pkg/testutils/sqlutils" + "github.com/cockroachdb/cockroach/pkg/util/leaktest" + "github.com/cockroachdb/cockroach/pkg/util/log" + "github.com/cockroachdb/datadriven" +) + +func TestDataDrivenTest(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + datadriven.Walk(t, datapathutils.TestDataPath(t), func(t *testing.T, path string) { + s, conn, _ := serverutils.StartServer(t, base.TestServerArgs{ + Knobs: base.TestingKnobs{ + SQLStatsKnobs: &sqlstats.TestingKnobs{ + SynchronousSQLStats: true, + }, + }, + }) + + defer s.Stopper().Stop(context.Background()) + statsRunner := sqlutils.MakeSQLRunner(conn) + datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string { + switch d.Cmd { + case "exec-sql": + sqlConn := s.SQLConn(t) + defer sqlConn.Close() + execSqlRunner := sqlutils.MakeSQLRunner(sqlConn) + var appName, dbName string + if d.HasArg("db") { + d.ScanArgs(t, "db", &dbName) + execSqlRunner.Exec(t, fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", dbName)) + execSqlRunner.Exec(t, fmt.Sprintf("USE %s", dbName)) + } else { + t.Fatalf("db arg is required") + } + if d.HasArg("app-name") { + d.ScanArgs(t, "app-name", &appName) + execSqlRunner.Exec(t, "SET APPLICATION_NAME = $1", appName) + } + execSqlRunner.Exec(t, d.Input) + return "" + case "show-stats": + var appName, dbName string + if d.HasArg("app-name") { + d.ScanArgs(t, "app-name", &appName) + } + if d.HasArg("db") { + d.ScanArgs(t, "db", &dbName) + } + return GetStats(t, statsRunner, appName, dbName) + case "show-txn-stats": + var appName, dbName string + if d.HasArg("app-name") { + d.ScanArgs(t, "app-name", &appName) + } + if d.HasArg("db") { + d.ScanArgs(t, "db", &dbName) + } + return GetTxnStats(t, statsRunner, appName, dbName) + default: + t.Fatalf("unexpected cmd %s", d.Cmd) + } + return "" + }) + }) +} + +func GetStats(t *testing.T, conn *sqlutils.SQLRunner, appName string, dbName string) string { + t.Helper() + query := ` SELECT row_to_json(s) FROM ( +SELECT + encode(fingerprint_id, 'hex') AS fingerprint_id, + encode(transaction_fingerprint_id, 'hex') AS transaction_fingerprint_id, + plan_hash != '\x0000000000000000' AS plan_hash_set, + metadata ->> 'query' AS query, + metadata ->> 'querySummary' AS summary, + metadata -> 'distsql' AS plan_distributed, + metadata -> 'implicitTxn' AS plan_implicit_txn, + metadata -> 'vec' AS plan_vectorized, + metadata -> 'fullScan' AS plan_full_scan, + statistics -> 'statistics'->> 'cnt' AS count, + statistics -> 'statistics'-> 'nodes' AS nodes, + statistics -> 'statistics' -> 'sqlType' AS sql_type, + (statistics -> 'statistics' -> 'parseLat' -> 'mean')::FLOAT > 0 AS parse_lat_not_zero, + (statistics -> 'statistics' -> 'planLat' -> 'mean')::FLOAT > 0 AS plan_lat_not_zero, + (statistics -> 'statistics' -> 'runLat' -> 'mean')::FLOAT > 0 AS run_lat_not_zero, + (statistics -> 'statistics' -> 'svcLat' -> 'mean')::FLOAT > 0 AS svc_lat_not_zero +FROM crdb_internal.statement_statistics +WHERE app_name = $1 +AND metadata->>'db' = $2 +ORDER BY fingerprint_id) s +` + res := conn.QueryStr(t, query, appName, dbName) + rows := make([]string, len(res)) + for rowIdx := range res { + rows[rowIdx] = strings.Join(res[rowIdx], ",") + } + return strings.Join(rows, "\n") +} + +func GetTxnStats(t *testing.T, conn *sqlutils.SQLRunner, appName string, dbName string) string { + query := ` SELECT row_to_json(s) FROM ( +SELECT + encode(fingerprint_id, 'hex') AS txn_fingerprint_id, + metadata->'stmtFingerprintIDs' AS statement_fingerprint_ids, + (statistics -> 'statistics' -> 'commitLat' -> 'mean')::FLOAT > 0 AS commit_lat_not_zero, + (statistics -> 'statistics' -> 'svcLat' -> 'mean')::FLOAT > 0 AS svc_lat_not_zero +FROM crdb_internal.transaction_statistics ts +WHERE ts.fingerprint_id in ( + SELECT DISTINCT ss.transaction_fingerprint_id + FROM crdb_internal.statement_statistics ss + WHERE app_name = $1 + AND metadata->>'db' = $2 +) +ORDER BY fingerprint_id) s +` + res := conn.QueryStr(t, query, appName, dbName) + rows := make([]string, len(res)) + for rowIdx := range res { + rows[rowIdx] = strings.Join(res[rowIdx], ",") + } + return strings.Join(rows, "\n") +} diff --git a/pkg/sql/sqlstats/testdata/plpgsql b/pkg/sql/sqlstats/testdata/plpgsql new file mode 100644 index 000000000000..f8696fc76114 --- /dev/null +++ b/pkg/sql/sqlstats/testdata/plpgsql @@ -0,0 +1,24 @@ +exec-sql db=udf_test app-name=plpgsql +CREATE TABLE IF NOT EXISTS test(a int); + CREATE OR REPLACE PROCEDURE insert_incremental(n INT) + LANGUAGE plpgsql + AS $$ + DECLARE + i INT := 0; + next_val INT; + BEGIN + WHILE i < n LOOP + SELECT COALESCE(MAX(a), 0) + 1 INTO next_val FROM test; + INSERT INTO test(a) VALUES (next_val); + i := i + 1; + END LOOP; + END; + $$; +CALL insert_incremental(5); +---- + +show-stats db=udf_test app-name=plpgsql +---- +{"count": "1", "fingerprint_id": "11c4e5400aa05722", "nodes": [], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "CREATE OR REPLACE PROCEDURE insert_incremental(n INT8)\n\tLANGUAGE plpgsql\n\tAS $$_$$", "run_lat_not_zero": true, "sql_type": "TypeDDL", "summary": "CREATE OR REPLACE PROCEDURE insert_incremental(n INT8)\n\tLANGUAGE plpgsql\n\tAS $$_$$", "svc_lat_not_zero": true, "transaction_fingerprint_id": "3013cc07fc915da5"} +{"count": "1", "fingerprint_id": "a82856550b801042", "nodes": [], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "CALL insert_incremental(_)", "run_lat_not_zero": true, "sql_type": "TypeTCL", "summary": "CALL insert_incremental(_)", "svc_lat_not_zero": true, "transaction_fingerprint_id": "3013cc07fc915da5"} +{"count": "1", "fingerprint_id": "f592c1a0850da0bf", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "CREATE TABLE IF NOT EXISTS test (a INT8)", "run_lat_not_zero": true, "sql_type": "TypeDDL", "summary": "CREATE TABLE IF NOT EXISTS test (a INT8)", "svc_lat_not_zero": true, "transaction_fingerprint_id": "5af17cec030c1760"} diff --git a/pkg/sql/sqlstats/testdata/query b/pkg/sql/sqlstats/testdata/query new file mode 100644 index 000000000000..2a4dddffac52 --- /dev/null +++ b/pkg/sql/sqlstats/testdata/query @@ -0,0 +1,38 @@ +exec-sql db=random_db +CREATE TABLE users (id INT PRIMARY KEY, name STRING); +INSERT INTO users VALUES (1, 'Alice'), (2, 'Bob'); +SELECT * FROM users; +---- + +# Sanity check shouldn't return anything +show-stats app-name=random db=random_db +---- + +show-stats db=random_db +---- +{"count": "1", "fingerprint_id": "269c5d2c33dec958", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "INSERT INTO users VALUES (_, __more__), (__more__)", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "INSERT INTO users", "svc_lat_not_zero": true, "transaction_fingerprint_id": "205354584637cfcb"} +{"count": "1", "fingerprint_id": "2cdc86177e8b242b", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SET database = random_db", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SET database = random_db", "svc_lat_not_zero": true, "transaction_fingerprint_id": "83bf3b5bf88a93f4"} +{"count": "1", "fingerprint_id": "e6733d9c37147d22", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": true, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SELECT * FROM users", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SELECT * FROM users", "svc_lat_not_zero": true, "transaction_fingerprint_id": "205354584637cfcb"} +{"count": "1", "fingerprint_id": "ec1bc23e5d2925f6", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "CREATE TABLE users (id INT8 PRIMARY KEY, name STRING)", "run_lat_not_zero": true, "sql_type": "TypeDDL", "summary": "CREATE TABLE users (id INT8 PRIMARY KEY, name STRING)", "svc_lat_not_zero": true, "transaction_fingerprint_id": "205354584637cfcb"} + +exec-sql db=random_db app-name=transactions +BEGIN; +INSERT INTO users VALUES (3, 'Charlie'); +INSERT INTO users VALUES (4, 'Diana'); +COMMIT; +---- + +exec-sql db=random_db app-name=transactions +BEGIN; +INSERT INTO users VALUES (5, 'Eve'); +INSERT INTO users VALUES (6, 'Frank'); +COMMIT; +---- + +show-stats db=random_db app-name=transactions +---- +{"count": "4", "fingerprint_id": "5fe99858726d3688", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": false, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "INSERT INTO users VALUES (_, __more__)", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "INSERT INTO users", "svc_lat_not_zero": true, "transaction_fingerprint_id": "78d7c1c32632f05d"} + +show-txn-stats db=random_db app-name=transactions +---- +{"commit_lat_not_zero": true, "statement_fingerprint_ids": ["5fe99858726d3688", "5fe99858726d3688"], "svc_lat_not_zero": true, "txn_fingerprint_id": "78d7c1c32632f05d"} diff --git a/pkg/sql/sqlstats/testdata/udf b/pkg/sql/sqlstats/testdata/udf new file mode 100644 index 000000000000..83efccc787c7 --- /dev/null +++ b/pkg/sql/sqlstats/testdata/udf @@ -0,0 +1,39 @@ +exec-sql db=udf_test app-name=udf +CREATE TABLE IF NOT EXISTS test(a int); +CREATE FUNCTION random_int() +RETURNS INT +LANGUAGE SQL +AS $$ + SELECT floor(random() * 100)::INT +$$ +VOLATILE; +SELECT random_int(); +CREATE FUNCTION do_something() +RETURNS INT +LANGUAGE SQL +AS $$ + INSERT INTO test(select random_int()); + SELECT max(a) FROM test; +$$ +VOLATILE; +select do_something(); +---- + +show-stats db=udf_test app-name=udf +---- +{"count": "1", "fingerprint_id": "8be77c892ac9ad38", "nodes": [], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "CREATE FUNCTION do_something()\n\tRETURNS INT8\n\tLANGUAGE SQL\n\tVOLATILE\n\tAS $$_$$", "run_lat_not_zero": true, "sql_type": "TypeDDL", "summary": "CREATE FUNCTION do_something()\n\tRETURNS INT8\n\tLANGUAGE SQL\n\tVOLATILE\n\tAS $$_$$", "svc_lat_not_zero": true, "transaction_fingerprint_id": "7a2be5f0c7083376"} +{"count": "1", "fingerprint_id": "af9bcc145f0d85f3", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SELECT do_something()", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SELECT do_something()", "svc_lat_not_zero": true, "transaction_fingerprint_id": "7a2be5f0c7083376"} +{"count": "1", "fingerprint_id": "bf1896e478760878", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SELECT random_int()", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SELECT random_int()", "svc_lat_not_zero": true, "transaction_fingerprint_id": "3af8b4d3f0ded8e0"} +{"count": "1", "fingerprint_id": "f56262fdc29506d7", "nodes": [], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "CREATE FUNCTION random_int()\n\tRETURNS INT8\n\tLANGUAGE SQL\n\tVOLATILE\n\tAS $$_$$", "run_lat_not_zero": true, "sql_type": "TypeDDL", "summary": "CREATE FUNCTION random_int()\n\tRETURNS INT8\n\tLANGUAGE SQL\n\tVOLATILE\n\tAS $$_$$", "svc_lat_not_zero": true, "transaction_fingerprint_id": "3af8b4d3f0ded8e0"} +{"count": "1", "fingerprint_id": "f592c1a0850da0bf", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": true, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "CREATE TABLE IF NOT EXISTS test (a INT8)", "run_lat_not_zero": true, "sql_type": "TypeDDL", "summary": "CREATE TABLE IF NOT EXISTS test (a INT8)", "svc_lat_not_zero": true, "transaction_fingerprint_id": "5af17cec030c1760"} + +exec-sql db=udf_test app-name=udf_transactions +BEGIN; +select do_something(); +select do_something(); +COMMIT; +---- + +show-stats db=udf_test app-name=udf_transactions +---- +{"count": "2", "fingerprint_id": "af9bcc145f0d85ff", "nodes": [1], "parse_lat_not_zero": true, "plan_distributed": false, "plan_full_scan": false, "plan_hash_set": true, "plan_implicit_txn": false, "plan_lat_not_zero": true, "plan_vectorized": true, "query": "SELECT do_something()", "run_lat_not_zero": true, "sql_type": "TypeDML", "summary": "SELECT do_something()", "svc_lat_not_zero": true, "transaction_fingerprint_id": "1dc175ec90b4a99f"}