sql: add optional message to CANCEL commands

docs/generated/http/full.md

@@ -2463,6 +2463,7 @@ Request object for issuing a query cancel request.
 | node_id | [string](#cockroach.server.serverpb.CancelQueryRequest-string) |  | ID of gateway node for the query to be canceled.<br><br>TODO(itsbilal): use [(gogoproto.customname) = "NodeID"] below. Need to figure out how to teach grpc-gateway about custom names.<br><br>node_id is a string so that "local" can be used to specify that no forwarding is necessary. | [reserved](#support-status) |
 | query_id | [string](#cockroach.server.serverpb.CancelQueryRequest-string) |  | ID of query to be canceled (converted to string). | [reserved](#support-status) |
 | username | [string](#cockroach.server.serverpb.CancelQueryRequest-string) |  | Username of the user making this cancellation request. This may be omitted if the user is the same as the one issuing the CancelQueryRequest. The caller is responsible for case-folding and NFC normalization. | [reserved](#support-status) |
+| message | [string](#cockroach.server.serverpb.CancelQueryRequest-string) |  | The message to be sent to the clients of the canceled query. | [reserved](#support-status) |
 
 
 

docs/generated/sql/bnf/cancel_query.bnf

@@ -1,5 +1,5 @@
 cancel_queries_stmt ::=
-	'CANCEL' 'QUERY' query_id
-	| 'CANCEL' 'QUERY' 'IF' 'EXISTS' query_id
-	| 'CANCEL' 'QUERIES' select_stmt
-	| 'CANCEL' 'QUERIES' 'IF' 'EXISTS' select_stmt
+	'CANCEL' 'QUERY' query_id opt_cancel_message
+	| 'CANCEL' 'QUERY' 'IF' 'EXISTS' query_id opt_cancel_message
+	| 'CANCEL' 'QUERIES' select_stmt opt_cancel_message
+	| 'CANCEL' 'QUERIES' 'IF' 'EXISTS' select_stmt opt_cancel_message

docs/generated/sql/bnf/stmt_block.bnf

@@ -569,10 +569,10 @@ cancel_jobs_stmt ::=
 	| 'CANCEL' 'JOBS' for_schedules_clause
 
 cancel_queries_stmt ::=
-	'CANCEL' 'QUERY' a_expr
-	| 'CANCEL' 'QUERY' 'IF' 'EXISTS' a_expr
-	| 'CANCEL' 'QUERIES' select_stmt
-	| 'CANCEL' 'QUERIES' 'IF' 'EXISTS' select_stmt
+	'CANCEL' 'QUERY' a_expr opt_cancel_message
+	| 'CANCEL' 'QUERY' 'IF' 'EXISTS' a_expr opt_cancel_message
+	| 'CANCEL' 'QUERIES' select_stmt opt_cancel_message
+	| 'CANCEL' 'QUERIES' 'IF' 'EXISTS' select_stmt opt_cancel_message
 
 cancel_sessions_stmt ::=
 	'CANCEL' 'SESSION' a_expr
@@ -1255,6 +1255,7 @@ unreserved_keyword ::=
 	| 'MATERIALIZED'
 	| 'MAXVALUE'
 	| 'MERGE'
+	| 'MESSAGE'
 	| 'METHOD'
 	| 'MINUTE'
 	| 'MINVALUE'
@@ -1771,6 +1772,10 @@ for_schedules_clause ::=
 	'FOR' 'SCHEDULES' select_stmt
 	| 'FOR' 'SCHEDULE' a_expr
 
+opt_cancel_message ::=
+	'WITH' 'MESSAGE' cancel_message
+	| 
+
 create_database_stmt ::=
 	'CREATE' 'DATABASE' database_name opt_with opt_template_clause opt_encoding_clause opt_lc_collate_clause opt_lc_ctype_clause opt_connection_limit opt_primary_region_clause opt_regions_list opt_survival_goal_clause opt_placement_clause opt_owner_clause opt_super_region_clause opt_secondary_region_clause
 	| 'CREATE' 'DATABASE' 'IF' 'NOT' 'EXISTS' database_name opt_with opt_template_clause opt_encoding_clause opt_lc_collate_clause opt_lc_ctype_clause opt_connection_limit opt_primary_region_clause opt_regions_list opt_survival_goal_clause opt_placement_clause opt_owner_clause opt_super_region_clause opt_secondary_region_clause
@@ -2446,6 +2451,9 @@ sub_type ::=
 	| 'SOME'
 	| 'ALL'
 
+cancel_message ::=
+	'SCONST'
+
 opt_template_clause ::=
 	'TEMPLATE' opt_equal non_reserved_word_or_sconst
 	| 
@@ -3958,6 +3966,7 @@ bare_label_keywords ::=
 	| 'MATERIALIZED'
 	| 'MAXVALUE'
 	| 'MERGE'
+	| 'MESSAGE'
 	| 'METHOD'
 	| 'MINVALUE'
 	| 'MODE'

pkg/server/serverpb/status.proto

@@ -1139,6 +1139,8 @@ message CancelQueryRequest {
   // if the user is the same as the one issuing the CancelQueryRequest.
   // The caller is responsible for case-folding and NFC normalization.
   string username = 3;
+  // The message to be sent to the clients of the canceled query.
+  string message = 4;
 }
 
 // Response returned by target query's gateway node.

pkg/server/status.go

@@ -3515,7 +3515,7 @@ func (s *statusServer) CancelQuery(
 		return nil, err
 	}
 
-	isCanceled := session.CancelQuery(queryID)
+	isCanceled := session.CancelQuery(queryID, req.Message)
 	return &serverpb.CancelQueryResponse{
 		Canceled: isCanceled,
 	}, nil

pkg/sql/cancel_queries.go

@@ -25,6 +25,7 @@ import (
 type cancelQueriesNode struct {
 	rows     planNode
 	ifExists bool
+	message  string
 }
 
 func (n *cancelQueriesNode) startExec(runParams) error {
@@ -62,6 +63,7 @@ func (n *cancelQueriesNode) Next(params runParams) (bool, error) {
 		NodeId:   fmt.Sprintf("%d", nodeID),
 		QueryID:  string(queryIDString),
 		Username: params.SessionData().User().Normalized(),
+		Message:  n.message,
 	}
 
 	response, err := params.extendedEvalCtx.SQLStatusServer.CancelQuery(params.ctx, request)
@@ -81,3 +83,17 @@ func (*cancelQueriesNode) Values() tree.Datums { return nil }
 func (n *cancelQueriesNode) Close(ctx context.Context) {
 	n.rows.Close(ctx)
 }
+
+// get the cancel message for the given query ID. If no message is found, return
+// an empty string.
+func getCancelMessage(s *Server, queryID clusterunique.ID) (message string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	message, ok := s.mu.cancels[queryID]
+	if ok && message != "" {
+		delete(s.mu.cancels, queryID)
+		return message
+	}
+	return ""
+}

pkg/sql/conn_executor.go

@@ -372,6 +372,10 @@ type Server struct {
 		syncutil.Mutex
 		connectionCount     int64
 		rootConnectionCount int64
+
+		// The key is the ID of the query that is being canceled, and the value is
+		// the cancel message sent from the canceler.
+		cancels map[clusterunique.ID]string
 	}
 }
 
@@ -466,6 +470,8 @@ func NewServer(cfg *ExecutorConfig, pool *mon.BytesMonitor) *Server {
 		idxRecommendationsCache: idxrecommendations.NewIndexRecommendationsCache(cfg.Settings),
 	}
 
+	s.mu.cancels = make(map[clusterunique.ID]string)
+
 	telemetryLoggingMetrics := newTelemetryLoggingMetrics(cfg.TelemetryLoggingTestingKnobs, cfg.Settings)
 	s.TelemetryLoggingMetrics = telemetryLoggingMetrics
 
@@ -4141,12 +4147,15 @@ func (ex *connExecutor) hasQuery(queryID clusterunique.ID) bool {
 }
 
 // CancelQuery is part of the RegistrySession interface.
-func (ex *connExecutor) CancelQuery(queryID clusterunique.ID) bool {
+func (ex *connExecutor) CancelQuery(queryID clusterunique.ID, message string) bool {
 	// RLock can be used because map deletion happens in
 	// connExecutor.removeActiveQuery.
 	ex.mu.RLock()
 	defer ex.mu.RUnlock()
 	if queryMeta, exists := ex.mu.ActiveQueries[queryID]; exists {
+		if message != "" {
+			ex.server.mu.cancels[queryID] = message
+		}
 		queryMeta.cancelQuery()
 		return true
 	}

pkg/sql/conn_executor_exec.go

@@ -851,7 +851,15 @@ func (ex *connExecutor) execStmtInOpenState(
 				retEv = eventNonRetriableErr{
 					IsCommit: fsm.FromBool(isCommit(ast)),
 				}
+
 				errToPush := cancelchecker.QueryCanceledError
+
+				// Provide a more detailed error message if the canceler provides it.
+				message := getCancelMessage(ex.server, queryID)
+				if message != "" {
+					errToPush = cancelchecker.QueryCanceledErrorWithMessage(message)
+				}
+
 				// For pausable portal, we can arrive here after encountering a timeout
 				// error and then perform a query-cleanup step. In this case, we don't
 				// want to override the original timeout error with the query-cancelled

pkg/sql/distsql_spec_exec_factory.go

@@ -1164,7 +1164,7 @@ func (e *distSQLSpecExecFactory) ConstructControlSchedules(
 }
 
 func (e *distSQLSpecExecFactory) ConstructCancelQueries(
-	input exec.Node, ifExists bool,
+	input exec.Node, ifExists bool, message string,
 ) (exec.Node, error) {
 	return nil, unimplemented.NewWithIssue(47473, "experimental opt-driven distsql planning: cancel queries")
 }

pkg/sql/exec_util.go

@@ -2365,7 +2365,7 @@ type RegistrySession interface {
 	SessionUser() username.SQLUsername
 	hasQuery(queryID clusterunique.ID) bool
 	// CancelQuery cancels the query specified by queryID if it exists.
-	CancelQuery(queryID clusterunique.ID) bool
+	CancelQuery(queryID clusterunique.ID, message string) bool
 	// CancelActiveQueries cancels all currently active queries.
 	CancelActiveQueries() bool
 	// CancelSession cancels the session.

pkg/sql/logictest/testdata/logic_test/cancel

@@ -0,0 +1,82 @@
+statement ok
+CREATE TABLE t (a INT PRIMARY KEY)
+
+statement ok
+GRANT ALL ON t TO testuser
+
+# insert a row so the following transactions can request x-lock on it, no
+# matter the transaction isolation level
+statement ok
+INSERT INTO t VALUES (1)
+
+subtest cancel
+
+# start a transaction and hold a x-lock on the row
+statement ok
+BEGIN; SELECT * FROM t WHERE a = 1 FOR UPDATE
+
+# issue the CANCEL QUERY with a 1-second delay
+statement async cancelstmt
+CANCEL QUERY (
+  WITH delay AS (
+    SELECT pg_sleep(1)
+  ),
+  stmts AS (
+    SHOW CLUSTER STATEMENTS
+  )
+  SELECT query_id
+  FROM stmts
+  WHERE query = 'SELECT * FROM t WHERE a = 1 FOR UPDATE'
+)
+
+user testuser
+
+# start a second transaction that blocks on the first transaction that holding the
+# x-lock, then this will be killed by the delayed CANCEL QUERY above
+query error pgcode 57014 query execution canceled
+SELECT * FROM t WHERE a = 1 FOR UPDATE
+
+user root
+
+awaitstatement cancelstmt
+
+statement ok
+COMMIT
+
+subtest end
+
+subtest cancel-with-message
+
+# start a transaction and hold a x-lock on the row
+statement ok
+BEGIN; SELECT * FROM t WHERE a = 1 FOR UPDATE
+
+# issue the CANCEL QUERY with a 1-second delay
+statement async cancelstmt
+CANCEL QUERY (
+  WITH delay AS (
+    SELECT pg_sleep(1)
+  ),
+  stmts AS (
+    SHOW CLUSTER STATEMENTS
+  )
+  SELECT query_id
+  FROM stmts
+  WHERE query = 'SELECT * FROM t WHERE a = 1 FOR UPDATE'
+) WITH MESSAGE 'cancel due to cluster maintenance'
+
+user testuser
+
+# start a second transaction that blocks on the first transaction that holding the
+# x-lock, then this will be killed by the delayed CANCEL QUERY above
+query error pgcode 57014 query execution canceled, message from the canceler: "cancel due to cluster maintenance"
+SELECT * FROM t WHERE a = 1 FOR UPDATE
+
+user root
+
+awaitstatement cancelstmt
+
+statement ok
+COMMIT
+
+subtest end