server: Tenant pods fan out to cancel SQL queries

docs/generated/http/full.md

@@ -2031,6 +2031,53 @@ Response returned by target query's gateway node.
 
 
 
+## CancelLocalQuery
+
+`POST /_status/cancel_local_query`
+
+CancelLocalQuery cancels a SQL query running on this node given its ID.
+
+Support status: [reserved](#support-status)
+
+#### Request Parameters
+
+
+
+
+Request object for issuing a query cancel request.
+
+
+| Field | Type | Label | Description | Support status |
+| ----- | ---- | ----- | ----------- | -------------- |
+| 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) |
+
+
+
+
+
+
+
+#### Response Parameters
+
+
+
+
+Response returned by target query's gateway node.
+
+
+| Field | Type | Label | Description | Support status |
+| ----- | ---- | ----- | ----------- | -------------- |
+| canceled | [bool](#cockroach.server.serverpb.CancelQueryResponse-bool) | | Whether the cancellation request succeeded and the query was canceled. | [reserved](#support-status) |
+| error | [string](#cockroach.server.serverpb.CancelQueryResponse-string) | | Error message (accompanied with canceled = false). | [reserved](#support-status) |
+
+
+
+
+
+
+
 ## ListContentionEvents
 
 `GET /_status/contention_events`

pkg/ccl/serverccl/tenant_status_test.go

@@ -14,6 +14,7 @@ import (
  "encoding/hex"
  "fmt"
  "net/url"
+ "reflect"
  "sort"
  "strconv"
  "strings"
@@ -537,3 +538,87 @@ func TestTenantStatusCancelSession(t *testing.T) {
  require.Equal(t, false, cancelSessionResp.Canceled)
  require.Equal(t, fmt.Sprintf("session ID %s not found", hex.EncodeToString(session.ID)), cancelSessionResp.Error)
 }
+
+func selectClusterQueryIDs(t *testing.T, conn *sqlutils.SQLRunner) []string {
+ var queryIDs []string
+ rows := conn.QueryStr(t, "SELECT query_id FROM crdb_internal.cluster_queries")
+ for _, row := range rows {
+ queryIDs = append(queryIDs, row[0])
+ }
+ return queryIDs
+}
+
+func TestTenantStatusCancelQuery(t *testing.T) {
+ defer leaktest.AfterTest(t)()
+ defer log.Scope(t).Close(t)
+
+ skip.UnderStressRace(t, "expensive tests")
+
+ ctx := context.Background()
+ helper := newTestTenantHelper(t, 3, base.TestingKnobs{})
+ defer helper.cleanup(ctx, t)
+
+ // Open a SQL session on tenant SQL pod 0 and start a long-running query.
+ sqlPod0 := helper.testCluster().tenantConn(0)
+ results := make(chan struct{})
+ errors := make(chan error)
+ defer close(results)
+ defer close(errors)
+ go func() {
+ if _, err := sqlPod0.DB.ExecContext(ctx, "SELECT pg_sleep(60)"); err != nil {
+ errors <- err
+ } else {
+ results <- struct{}{}
+ }
+ }()
+
+ // See the query over HTTP on tenant SQL pod 1.
+ httpPod1, err := helper.testCluster().tenantHTTPJSONClient(1)
+ require.NoError(t, err)
+ defer httpPod1.Close()
+ var listSessionsResp serverpb.ListSessionsResponse
+ var query serverpb.ActiveQuery
+ require.Eventually(t, func() bool {
+ err = httpPod1.GetJSON("/_status/sessions", &listSessionsResp)
+ require.NoError(t, err)
+ for _, s := range listSessionsResp.Sessions {
+ for _, q := range s.ActiveQueries {
+ if q.Sql == "SELECT pg_sleep(60)" {
+ query = q
+ break
+ }
+ }
+ }
+ return query.ID != ""
+ }, 10*time.Second, 100*time.Millisecond, "query not found")
+
+ // See the query over SQL on tenant SQL pod 0.
+ require.Contains(t, selectClusterQueryIDs(t, sqlPod0), query.ID)
+
+ // Cancel the query over HTTP on tenant SQL pod 1.
+ cancelQueryReq := serverpb.CancelQueryRequest{QueryID: query.ID}
+ cancelQueryResp := serverpb.CancelQueryResponse{}
+ err = httpPod1.PostJSON("/_status/cancel_query/0", &cancelQueryReq, &cancelQueryResp)
+ require.NoError(t, err)
+ require.Equal(t, true, cancelQueryResp.Canceled,
+ "expected query to be canceled, but encountered unexpected error %s", cancelQueryResp.Error)
+
+ // No longer see the query over SQL on tenant SQL pod 0.
+ require.Eventually(t, func() bool {
+ return !strings.Contains(reflect.ValueOf(selectClusterQueryIDs(t, sqlPod0)).String(), query.ID)
+ }, 10*time.Second, 100*time.Millisecond,
+ "expected query %s to no longer be visible in crdb_internal.cluster_queries", query.ID)
+
+ select {
+ case <-results:
+ t.Fatalf("Expected long-running query to have been canceled with error.")
+ case err := <-errors:
+ require.Equal(t, "pq: query execution canceled", err.Error())
+ }
+
+ // Attempt to cancel the query again over HTTP from tenant SQL pod 1, so that we can see the error message.
+ err = httpPod1.PostJSON("/_status/cancel_query/0", &cancelQueryReq, &cancelQueryResp)
+ require.NoError(t, err)
+ require.Equal(t, false, cancelQueryResp.Canceled)
+ require.Equal(t, fmt.Sprintf("query ID %s not found", query.ID), cancelQueryResp.Error)
+}

pkg/rpc/auth_tenant.go

@@ -97,6 +97,12 @@ func (a tenantAuthorizer) authorize(
  case "/cockroach.server.serverpb.Status/CancelLocalSession":
  return a.authTenant(tenID)
 
+ case "/cockroach.server.serverpb.Status/CancelQuery":
+ return a.authTenant(tenID)
+
+ case "/cockroach.server.serverpb.Status/CancelLocalQuery":
+ return a.authTenant(tenID)
+
  case "/cockroach.roachpb.Internal/GetSpanConfigs":
  return a.authGetSpanConfigs(tenID, req.(*roachpb.GetSpanConfigsRequest))