github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/server/statements.go (about)

     1  // Copyright 2014 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package server
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    18  	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    20  	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
    21  	"google.golang.org/grpc/codes"
    22  	"google.golang.org/grpc/status"
    23  )
    24  
    25  func (s *statusServer) Statements(
    26  	ctx context.Context, req *serverpb.StatementsRequest,
    27  ) (*serverpb.StatementsResponse, error) {
    28  	if _, err := s.admin.requireAdminUser(ctx); err != nil {
    29  		return nil, err
    30  	}
    31  
    32  	ctx = propagateGatewayMetadata(ctx)
    33  	ctx = s.AnnotateCtx(ctx)
    34  
    35  	response := &serverpb.StatementsResponse{
    36  		Statements:            []serverpb.StatementsResponse_CollectedStatementStatistics{},
    37  		LastReset:             timeutil.Now(),
    38  		InternalAppNamePrefix: sqlbase.InternalAppNamePrefix,
    39  	}
    40  
    41  	localReq := &serverpb.StatementsRequest{
    42  		NodeID: "local",
    43  	}
    44  
    45  	if len(req.NodeID) > 0 {
    46  		requestedNodeID, local, err := s.parseNodeID(req.NodeID)
    47  		if err != nil {
    48  			return nil, status.Errorf(codes.InvalidArgument, err.Error())
    49  		}
    50  		if local {
    51  			return s.StatementsLocal(ctx)
    52  		}
    53  		status, err := s.dialNode(ctx, requestedNodeID)
    54  		if err != nil {
    55  			return nil, err
    56  		}
    57  		return status.Statements(ctx, localReq)
    58  	}
    59  
    60  	dialFn := func(ctx context.Context, nodeID roachpb.NodeID) (interface{}, error) {
    61  		client, err := s.dialNode(ctx, nodeID)
    62  		return client, err
    63  	}
    64  	nodeStatement := func(ctx context.Context, client interface{}, _ roachpb.NodeID) (interface{}, error) {
    65  		status := client.(serverpb.StatusClient)
    66  		return status.Statements(ctx, localReq)
    67  	}
    68  
    69  	if err := s.iterateNodes(ctx, fmt.Sprintf("statement statistics for node %s", req.NodeID),
    70  		dialFn,
    71  		nodeStatement,
    72  		func(nodeID roachpb.NodeID, resp interface{}) {
    73  			statementsResp := resp.(*serverpb.StatementsResponse)
    74  			response.Statements = append(response.Statements, statementsResp.Statements...)
    75  			if response.LastReset.After(statementsResp.LastReset) {
    76  				response.LastReset = statementsResp.LastReset
    77  			}
    78  		},
    79  		func(nodeID roachpb.NodeID, err error) {
    80  			// TODO(couchand): do something here...
    81  		},
    82  	); err != nil {
    83  		return nil, err
    84  	}
    85  
    86  	return response, nil
    87  }
    88  
    89  func (s *statusServer) StatementsLocal(ctx context.Context) (*serverpb.StatementsResponse, error) {
    90  	stmtStats := s.admin.server.sqlServer.pgServer.SQLServer.GetUnscrubbedStmtStats()
    91  	lastReset := s.admin.server.sqlServer.pgServer.SQLServer.GetStmtStatsLastReset()
    92  
    93  	resp := &serverpb.StatementsResponse{
    94  		Statements:            make([]serverpb.StatementsResponse_CollectedStatementStatistics, len(stmtStats)),
    95  		LastReset:             lastReset,
    96  		InternalAppNamePrefix: sqlbase.InternalAppNamePrefix,
    97  	}
    98  
    99  	for i, stmt := range stmtStats {
   100  		resp.Statements[i] = serverpb.StatementsResponse_CollectedStatementStatistics{
   101  			Key: serverpb.StatementsResponse_ExtendedStatementStatisticsKey{
   102  				KeyData: stmt.Key,
   103  				NodeID:  s.gossip.NodeID.Get(),
   104  			},
   105  			Stats: stmt.Stats,
   106  		}
   107  	}
   108  
   109  	return resp, nil
   110  }