server: add fetch mode option to combined req

This commit adds a new field, fetch_mode, which
is a message containing an enum specifying what
kind of stats to return from the combined stmts
endpoint. The options are
- StmtStatsOnly
- TxnStatsOnly

Leaving this field null will fetch both stmts
and txns.

If `TxnStatsOnly` is specified, we will also include
the stmt stats for the stms in each txn in the returned
txn payload. This is because in the client app we need
the stmt stats to properly build the txn fingerprint pages.

In the future, we will split this API instead of using
this fetch mode flag, but this is currently preferred to
make it easier to backport for perf improvements.

Epic: none
Part of: cockroachdb#97252
Part of: cockroachdb#97875

Release note: None

docs/generated/http/full.md

@@ -4004,12 +4004,26 @@ Support status: [reserved](#support-status)
 | ----- | ---- | ----- | ----------- | -------------- |
 | start | [int64](#cockroach.server.serverpb.CombinedStatementsStatsRequest-int64) |  | Unix time range for aggregated statements. | [reserved](#support-status) |
 | end | [int64](#cockroach.server.serverpb.CombinedStatementsStatsRequest-int64) |  |  | [reserved](#support-status) |
+| fetch_mode | [CombinedStatementsStatsRequest.FetchMode](#cockroach.server.serverpb.CombinedStatementsStatsRequest-cockroach.server.serverpb.CombinedStatementsStatsRequest.FetchMode) |  | Note that if fetch_mode is set to transactions only, we will also include the statement statistics for the stmts in the transactions response. This is more of a hack-y method to get the complete stats for txns, because in the client we need to fill in some txn stats info from its stmt stats, such as the query string.<br><br>We prefer this hackier method right now to reduce surface area for backporting these changes, but in the future we will introduce more endpoints to properly organize these differing requests. TODO (xinhaoz) - Split this API into stmts and txns properly instead of using  this param. | [reserved](#support-status) |
 
 
 
 
 
 
+<a name="cockroach.server.serverpb.CombinedStatementsStatsRequest-cockroach.server.serverpb.CombinedStatementsStatsRequest.FetchMode"></a>
+#### CombinedStatementsStatsRequest.FetchMode
+
+
+
+| Field | Type | Label | Description | Support status |
+| ----- | ---- | ----- | ----------- | -------------- |
+| stats_type | [CombinedStatementsStatsRequest.StatsType](#cockroach.server.serverpb.CombinedStatementsStatsRequest-cockroach.server.serverpb.CombinedStatementsStatsRequest.StatsType) |  |  | [reserved](#support-status) |
+
+
+
+
+
 
 #### Response Parameters
 

pkg/server/combined_statement_stats.go

@@ -73,12 +73,26 @@ func getCombinedStatementStats(
 	showInternal := SQLStatsShowInternal.Get(&settings.SV)
 	whereClause, orderAndLimit, args := getCombinedStatementsQueryClausesAndArgs(
 		startTime, endTime, limit, testingKnobs, showInternal)
-	statements, err := collectCombinedStatements(ctx, ie, whereClause, args, orderAndLimit)
-	if err != nil {
-		return nil, serverError(ctx, err)
+	var statements []serverpb.StatementsResponse_CollectedStatementStatistics
+	var transactions []serverpb.StatementsResponse_ExtendedCollectedTransactionStatistics
+	var err error
+
+	if req.FetchMode == nil || req.FetchMode.StatsType == serverpb.CombinedStatementsStatsRequest_TxnStatsOnly {
+		transactions, err = collectCombinedTransactions(ctx, ie, whereClause, args, orderAndLimit)
+		if err != nil {
+			return nil, serverError(ctx, err)
+		}
 	}
 
-	transactions, err := collectCombinedTransactions(ctx, ie, whereClause, args, orderAndLimit)
+	if req.FetchMode != nil && req.FetchMode.StatsType == serverpb.CombinedStatementsStatsRequest_TxnStatsOnly {
+		// Change the whereClause for the statements to those matching the txn_fingerprint_ids in the
+		// transactions response that are within the desired interval. We also don't need the order and
+		// limit anymore.
+		orderAndLimit = ""
+		whereClause, args = buildWhereClauseForStmtsByTxn(req, transactions, testingKnobs)
+	}
+
+	statements, err = collectCombinedStatements(ctx, ie, whereClause, args, orderAndLimit)
 	if err != nil {
 		return nil, serverError(ctx, err)
 	}
@@ -93,6 +107,44 @@ func getCombinedStatementStats(
 	return response, nil
 }
 
+// buildWhereClauseForStmtsByTxn builds the where clause to get the statement
+// stats based on a list of transactions. The list of transactions provided must
+// contain no duplicate transaction fingerprint ids.
+func buildWhereClauseForStmtsByTxn(
+	req *serverpb.CombinedStatementsStatsRequest,
+	transactions []serverpb.StatementsResponse_ExtendedCollectedTransactionStatistics,
+	testingKnobs *sqlstats.TestingKnobs,
+) (whereClause string, args []interface{}) {
+	var buffer strings.Builder
+	buffer.WriteString(testingKnobs.GetAOSTClause())
+
+	buffer.WriteString(" WHERE true")
+
+	// Add start and end filters from request.
+	startTime := getTimeFromSeconds(req.Start)
+	endTime := getTimeFromSeconds(req.End)
+	if startTime != nil {
+		args = append(args, *startTime)
+		buffer.WriteString(fmt.Sprintf(" AND aggregated_ts >= $%d", len(args)))
+	}
+
+	if endTime != nil {
+		args = append(args, *endTime)
+		buffer.WriteString(fmt.Sprintf(" AND aggregated_ts <= $%d", len(args)))
+	}
+
+	txnFingerprints := make([]string, 0, len(transactions))
+	for i := range transactions {
+		fingerprint := uint64(transactions[i].StatsData.TransactionFingerprintID)
+		txnFingerprints = append(txnFingerprints, fmt.Sprintf("\\x%016x", fingerprint))
+	}
+
+	args = append(args, txnFingerprints)
+	buffer.WriteString(fmt.Sprintf(" AND transaction_fingerprint_id = any $%d", len(args)))
+
+	return buffer.String(), args
+}
+
 // getCombinedStatementsQueryClausesAndArgs returns:
 // - where clause (filtering by name and aggregates_ts when defined)
 // - order and limit clause

pkg/server/serverpb/status.proto

@@ -1455,9 +1455,31 @@ message StatementsResponse {
 }
 
 message CombinedStatementsStatsRequest {
+  enum StatsType {
+    StmtStatsOnly = 0;
+    TxnStatsOnly = 1;
+  }
+
+  message FetchMode {
+    StatsType stats_type = 1;
+  }
+
   // Unix time range for aggregated statements.
   int64 start = 1 [(gogoproto.nullable) = true];
   int64 end = 2 [(gogoproto.nullable) = true];
+
+  // Note that if fetch_mode is set to transactions only, we will also
+  // include the statement statistics for the stmts in the transactions
+  // response. This is more of a hack-y method to get the complete stats
+  // for txns, because in the client we need to fill in some txn stats info
+  // from its stmt stats, such as the query string.
+  //
+  // We prefer this hackier method right now to reduce surface area for backporting
+  // these changes, but in the future we will introduce more endpoints to properly
+  // organize these differing requests.
+  // TODO (xinhaoz) - Split this API into stmts and txns properly instead of using
+  //  this param.
+  FetchMode fetch_mode = 5 [(gogoproto.nullable) = true];
 }
 
 // StatementDetailsRequest requests the details of a Statement, based on its keys.
@@ -2051,6 +2073,7 @@ service Status {
       get: "/_status/combinedstmts"
     };
   }
+
   rpc StatementDetails(StatementDetailsRequest) returns (StatementDetailsResponse) {
     option (google.api.http) = {
       get: "/_status/stmtdetails/{fingerprint_id}"