Merge pull request #129805 from xinhaoz/backport24.1-128856

release-24.1: ui: use status server apis for stmt bundle ops

pkg/ui/workspaces/cluster-ui/src/api/statementDiagnosticsApi.ts

@@ -8,16 +8,17 @@
 // by the Apache License, Version 2.0, included in the file
 // licenses/APL.txt.
 
+import { cockroach } from "@cockroachlabs/crdb-protobuf-client";
+import Long from "long";
 import moment from "moment-timezone";
-import { Duration } from "src/util/format";
-import {
-  createSqlExecutionRequest,
-  executeInternalSql,
-  executeInternalSqlHelper,
-  SqlExecutionRequest,
-  SqlTxnResult,
-  txnResultSetIsEmpty,
-} from "src/api";
+
+import { fetchData } from "src/api";
+
+import { NumberToDuration } from "../util";
+
+const STATEMENT_DIAGNOSTICS_PATH = "_status/stmtdiagreports";
+const CANCEL_STATEMENT_DIAGNOSTICS_PATH =
+  STATEMENT_DIAGNOSTICS_PATH + "/cancel";
 
 export type StatementDiagnosticsReport = {
   id: string;
@@ -32,70 +33,26 @@ export type StatementDiagnosticsReport = {
 export type StatementDiagnosticsResponse = StatementDiagnosticsReport[];
 
 export async function getStatementDiagnosticsReports(): Promise<StatementDiagnosticsResponse> {
-  let result: StatementDiagnosticsResponse = [];
-
-  const createReq = () => {
-    let offset = "";
-    const args = [];
-    if (result.length > 0) {
-      // Using the id is more performant and reliable than offset.
-      // Schema is PRIMARY KEY (id) with INT8 DEFAULT unique_rowid() NOT NULL.
-      offset = " AND (id::STRING < $1) ";
-      const last = result[result.length - 1];
-      args.push(last.id);
-    }
-    const query = `SELECT
-      id::STRING,
-      statement_fingerprint,
-      completed,
-      statement_diagnostics_id::STRING,
-      requested_at,
-      min_execution_latency,
-      expires_at
-    FROM
-      system.statement_diagnostics_requests 
-    WHERE
-     (expires_at > now() OR expires_at IS NULL OR completed = true) ${offset}
-     order by id desc`;
-
-    return createSqlExecutionRequest(undefined, [
-      {
-        sql: query,
-        arguments: args,
-      },
-    ]);
-  };
-
-  const err = await executeInternalSqlHelper<StatementDiagnosticsReport>(
-    createReq,
-    (response: SqlTxnResult<StatementDiagnosticsReport>[]) => {
-      if (!response) {
-        return;
-      }
-
-      if (txnResultSetIsEmpty(response)) {
-        return;
-      }
-
-      response.forEach(x => {
-        if (x.rows && x.rows.length > 0) {
-          result = result.concat(x.rows);
-        }
-      });
-    },
+  const response = await fetchData(
+    cockroach.server.serverpb.StatementDiagnosticsReportsResponse,
+    STATEMENT_DIAGNOSTICS_PATH,
   );
-
-  if (err) {
-    throw err;
-  }
-
-  return result;
+  return response.reports.map(report => {
+    return {
+      id: report.id.toString(),
+      statement_fingerprint: report.statement_fingerprint,
+      completed: report.completed,
+      statement_diagnostics_id: report.statement_diagnostics_id.toString(),
+      requested_at: moment.unix(report.requested_at.seconds.toNumber()),
+      min_execution_latency: moment.duration(
+        report.min_execution_latency.seconds.toNumber(),
+        "seconds",
+      ),
+      expires_at: moment.unix(report.expires_at.seconds.toNumber()),
+    };
+  });
 }
 
-type CheckPendingStmtDiagnosticRow = {
-  count: number;
-};
-
 export type InsertStmtDiagnosticRequest = {
   stmtFingerprint: string;
   samplingProbability?: number;
@@ -108,97 +65,24 @@ export type InsertStmtDiagnosticResponse = {
   req_resp: boolean;
 };
 
-export function createStatementDiagnosticsReport({
-  stmtFingerprint,
-  samplingProbability,
-  minExecutionLatencySeconds,
-  expiresAfterSeconds,
-  planGist,
-}: InsertStmtDiagnosticRequest): Promise<InsertStmtDiagnosticResponse> {
-  const args: any = [stmtFingerprint];
-  let query = `SELECT crdb_internal.request_statement_bundle($1, $2, $3::INTERVAL, $4::INTERVAL) as req_resp`;
-
-  if (planGist) {
-    args.push(planGist);
-    query = `SELECT crdb_internal.request_statement_bundle($1, $2, $3, $4::INTERVAL, $5::INTERVAL) as req_resp`;
-  }
-
-  if (samplingProbability) {
-    args.push(samplingProbability);
-  } else {
-    args.push(0);
-  }
-  if (minExecutionLatencySeconds) {
-    args.push(Duration(minExecutionLatencySeconds * 1e9));
-  } else {
-    args.push("0");
-  }
-  if (expiresAfterSeconds && expiresAfterSeconds !== 0) {
-    args.push(Duration(expiresAfterSeconds * 1e9));
-  } else {
-    args.push("0");
-  }
-
-  const createStmtDiag = {
-    sql: query,
-    arguments: args,
-  };
-
-  const req: SqlExecutionRequest = {
-    execute: true,
-    statements: [createStmtDiag],
-  };
-
-  return checkExistingDiagRequest(stmtFingerprint).then(_ => {
-    return executeInternalSql<InsertStmtDiagnosticResponse>(req).then(res => {
-      // If request succeeded but query failed, throw error (caught by saga/cacheDataReducer).
-      if (res.error) {
-        throw res.error;
-      }
-
-      if (
-        res.execution?.txn_results[0]?.rows?.length === 0 ||
-        res.execution?.txn_results[0]?.rows[0]["req_resp"] === false
-      ) {
-        throw new Error("Failed to insert statement diagnostics request");
-      }
-
-      return res.execution.txn_results[0].rows[0];
-    });
-  });
-}
-
-function checkExistingDiagRequest(stmtFingerprint: string): Promise<void> {
-  // Query to check if there's already a pending request for this fingerprint.
-  const checkPendingStmtDiag = {
-    sql: `SELECT count(1) FROM system.statement_diagnostics_requests
-        WHERE
-        completed = false AND
-        statement_fingerprint = $1 AND
-      (expires_at IS NULL OR expires_at > now())`,
-    arguments: [stmtFingerprint],
-  };
-
-  const req: SqlExecutionRequest = {
-    execute: true,
-    statements: [checkPendingStmtDiag],
-  };
-
-  return executeInternalSql<CheckPendingStmtDiagnosticRow>(req).then(res => {
-    // If request succeeded but query failed, throw error (caught by saga/cacheDataReducer).
-    if (res.error) {
-      throw res.error;
-    }
-
-    if (res.execution?.txn_results[0]?.rows?.length === 0) {
-      throw new Error("Failed to check pending statement diagnostics");
-    }
-
-    if (res.execution.txn_results[0].rows[0].count > 0) {
-      throw new Error(
-        "A pending request for the requested fingerprint already exists. Cancel the existing request first and try again.",
-      );
-    }
+export async function createStatementDiagnosticsReport(
+  req: InsertStmtDiagnosticRequest,
+): Promise<InsertStmtDiagnosticResponse> {
+  return fetchData(
+    cockroach.server.serverpb.CreateStatementDiagnosticsReportResponse,
+    STATEMENT_DIAGNOSTICS_PATH,
+    cockroach.server.serverpb.CreateStatementDiagnosticsReportRequest,
+    new cockroach.server.serverpb.CreateStatementDiagnosticsReportRequest({
+      statement_fingerprint: req.stmtFingerprint,
+      sampling_probability: req.samplingProbability,
+      min_execution_latency: NumberToDuration(req.minExecutionLatencySeconds),
+      expires_after: NumberToDuration(req.expiresAfterSeconds),
+      plan_gist: req.planGist,
+    }),
+  ).then(response => {
+    return {
+      req_resp: response.report !== null,
+    };
   });
 }
 
@@ -210,42 +94,22 @@ export type CancelStmtDiagnosticResponse = {
   stmt_diag_req_id: string;
 };
 
-export function cancelStatementDiagnosticsReport({
-  requestId,
-}: CancelStmtDiagnosticRequest): Promise<CancelStmtDiagnosticResponse> {
-  const query = `UPDATE system.statement_diagnostics_requests 
-SET expires_at = '1970-01-01' 
-WHERE completed = false 
-AND id = $1::INT8 
-AND (expires_at IS NULL OR expires_at > now()) RETURNING id as stmt_diag_req_id`;
-  const req: SqlExecutionRequest = {
-    execute: true,
-    statements: [
-      {
-        sql: query,
-        arguments: [requestId],
-      },
-    ],
-  };
-
-  return executeInternalSql<CancelStmtDiagnosticResponse>(req).then(res => {
-    // If request succeeded but query failed, throw error (caught by saga/cacheDataReducer).
-    if (res.error) {
-      throw res.error;
+export async function cancelStatementDiagnosticsReport(
+  req: CancelStmtDiagnosticRequest,
+): Promise<CancelStmtDiagnosticResponse> {
+  return fetchData(
+    cockroach.server.serverpb.CancelStatementDiagnosticsReportResponse,
+    CANCEL_STATEMENT_DIAGNOSTICS_PATH,
+    cockroach.server.serverpb.CancelStatementDiagnosticsReportRequest,
+    new cockroach.server.serverpb.CancelStatementDiagnosticsReportRequest({
+      request_id: Long.fromString(req.requestId),
+    }),
+  ).then(response => {
+    if (response.error) {
+      throw new Error(response.error);
     }
-
-    if (!res.execution?.txn_results?.length) {
-      throw new Error(
-        "cancelStatementDiagnosticsReport - unexpected zero txn_results",
-      );
-    }
-
-    if (res.execution.txn_results[0].rows?.length === 0) {
-      throw new Error(
-        `No pending request found for the request id: ${requestId}`,
-      );
-    }
-
-    return res.execution.txn_results[0].rows[0];
+    return {
+      stmt_diag_req_id: req.requestId,
+    };
   });
 }

pkg/ui/workspaces/e2e-tests/cypress/e2e/health-check/authenticated.cy.ts

@@ -8,9 +8,14 @@
 // by the Apache License, Version 2.0, included in the file
 // licenses/APL.txt.
 
+import { SQLPrivilege } from "../../support/types";
+
 describe("health check: authenticated user", () => {
   it("serves a DB Console overview page", () => {
-    cy.login();
+    cy.getUserWithExactPrivileges([SQLPrivilege.ADMIN]);
+    cy.fixture("users").then((users) => {
+      cy.login(users[0].username, users[0].password);
+    });
 
     // Ensure that something reasonable renders at / when authenticated, making
     // just enough assertions to ensure the right page loaded. If this test

pkg/ui/workspaces/e2e-tests/cypress/e2e/statementBundles/statementBundles.cy.ts

@@ -0,0 +1,91 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+import { SQLPrivilege } from "../../support/types";
+
+// TODO (xinhaoz): This test currently only works when running pnpm run cy:run
+// directly against a local cluster set up with sql activity in the last hour
+// and the expected cypress users in fixtures. We need to automate this server
+// setup for e2e testing.
+describe("statement bundles", () => {
+  const runTestsForPrivilegedUser = (privilege: SQLPrivilege) => {
+    describe(`${privilege} user`, () => {
+      beforeEach(() => {
+        cy.getUserWithExactPrivileges([privilege]).then((user) => {
+          cy.login(user.username, user.password);
+        });
+      });
+
+      it("can request statement bundles", () => {
+        cy.visit("#/sql-activity");
+        cy.contains("button", "Apply").click();
+        // Open modal.
+        cy.contains("button", "Activate").click();
+        // Wait for modal.
+        cy.findByText(/activate statement diagnostics/i).should("be.visible");
+        // Click the Activate button within the modal
+        cy.get(`[role="dialog"]`) // Adjust this selector to match your modal's structure
+          .contains("button", "Activate")
+          .click();
+        cy.findByText(/statement diagnostics were successfully activated/i);
+      });
+
+      it("can view statement bundles", () => {
+        cy.visit("#/reports/statements/diagnosticshistory");
+        cy.get("table tbody tr").should("have.length.at.least", 1);
+      });
+
+      it("can cancel statement bundles", () => {
+        cy.visit("#/reports/statements/diagnosticshistory");
+        cy.get("table tbody tr").should("have.length.at.least", 1);
+        cy.contains("button", "Cancel").click();
+        cy.findByText(/statement diagnostics were successfully cancelled/i);
+      });
+    });
+  };
+
+  const runTestsForNonPrivilegedUser = (privilege?: SQLPrivilege) => {
+    beforeEach(() => {
+      cy.getUserWithExactPrivileges(privilege ? [privilege] : []).then(
+        (user) => {
+          cy.login(user.username, user.password);
+        },
+      );
+    });
+
+    it("cannot request statement bundles", () => {
+      cy.visit("#/sql-activity");
+      cy.contains("button", "Apply").click();
+      // Should not see an Activate button.
+      cy.contains("button", "Activate").should("not.exist");
+    });
+
+    it("cannot view statement bundles", () => {
+      cy.visit("#/reports/statements/diagnosticshistory");
+      cy.findByText(/no statement diagnostics to show/i);
+    });
+  };
+
+  describe("view activity user", () => {
+    runTestsForPrivilegedUser(SQLPrivilege.VIEWACTIVITY);
+  });
+
+  describe("admin user", () => {
+    runTestsForPrivilegedUser(SQLPrivilege.ADMIN);
+  });
+
+  describe("non-privileged VIEWACTIVITYREDACTED user", () => {
+    runTestsForNonPrivilegedUser(SQLPrivilege.VIEWACTIVITYREDACTED);
+  });
+
+  describe("non-privileged user", () => {
+    runTestsForNonPrivilegedUser();
+  });
+});