From 9091c86dfdfb37a6a778430d0aa141d35e8956f3 Mon Sep 17 00:00:00 2001
From: Tommy Reilly <treilly@cockroachlabs.com>
Date: Wed, 30 Nov 2022 17:22:20 -0500
Subject: [PATCH] sql: add profiler labels for sql statement running remote
 flows

Fixes: #82464

Release note (sql change): SQL queries running on remote nodes
now show up in cpu profiles with "distsql.*" labels. Currently include
appname, gateway, txn and stmt.
---
 pkg/sql/distsql/server.go | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/pkg/sql/distsql/server.go b/pkg/sql/distsql/server.go
index 7c02ccd55de3..0e31f0305474 100644
--- a/pkg/sql/distsql/server.go
+++ b/pkg/sql/distsql/server.go
@@ -416,6 +416,16 @@ func (ds *ServerImpl) setupFlow(
 		if req.JobTag != "" {
 			flowCtx.AmbientContext.AddLogTag("job", req.JobTag)
 		}
+		if req.StatementSQL != "" {
+			flowCtx.AmbientContext.AddLogTag("distsql.stmt", req.StatementSQL)
+		}
+		flowCtx.AmbientContext.AddLogTag("distsql.gateway", req.Flow.Gateway)
+		if req.EvalContext.SessionData.ApplicationName != "" {
+			flowCtx.AmbientContext.AddLogTag("distsql.appname", req.EvalContext.SessionData.ApplicationName)
+		}
+		if leafTxn != nil {
+			flowCtx.AmbientContext.AddLogTag("distsql.txn", leafTxn.ID())
+		}
 		ctx = flowCtx.AmbientContext.AnnotateCtx(ctx)
 		telemetry.Inc(sqltelemetry.DistSQLExecCounter)
 	}