From 6c1b6114dc6f309bbf6083c1240cfb0c9c8ea4ee Mon Sep 17 00:00:00 2001 From: Lindsey Jin Date: Tue, 21 Dec 2021 22:09:03 -0500 Subject: [PATCH 1/2] ui/server: add ability to create conditional statement diagnostics Resolves #57634 Previously, we did not have the ability to create conditional statement diagnostics in the frontend. This commit adds in the ability to specify a minimum execution latency and an expiry time when creating a statement diagnostics request. These changes apply to both DB and CC console. Since expired requests are not surfaced at all in the frontend, we have also modified the statement diagnostics API response to not return already expired and incomplete requests. Lastly, this commit also deletes some unused code related to statement diagnostics modals. Release note (ui change): Add ability to create conditional statement diagnostics by adding two new fields 1) minimum execution latency, which specifies the limit for when a statement should be tracked and 2) expiry time, which specifies when a diagnostics request should expire. --- pkg/server/statement_diagnostics_requests.go | 12 +- pkg/server/status_test.go | 44 +++++ .../src/api/statementDiagnosticsApi.ts | 7 +- .../workspaces/cluster-ui/src/modal/modal.tsx | 4 +- .../diagnostics/diagnosticsView.spec.tsx | 24 +-- .../diagnostics/diagnosticsView.tsx | 54 +++--- .../src/statementDetails/statementDetails.tsx | 30 +++- .../statementDetailsConnected.ts | 21 ++- .../activateStatementDiagnosticsModal.scss | 108 ++++++++++++ .../activateStatementDiagnosticsModal.tsx | 161 ++++++++++++++++-- .../src/statementsPage/statementsPage.tsx | 9 +- .../statementsPageConnected.tsx | 20 ++- .../statementDiagnostics.reducer.ts | 6 +- .../statementDiagnostics.sagas.spec.ts | 36 ++-- .../cluster-ui/src/util/convert.spec.ts | 18 ++ .../workspaces/cluster-ui/src/util/convert.ts | 14 ++ .../src/redux/statements/statementsActions.ts | 14 +- .../redux/statements/statementsSagas.spec.ts | 30 +++- .../src/redux/statements/statementsSagas.ts | 21 ++- .../statementDiagnosticsHistory/index.tsx | 2 +- .../diagnostics/activateDiagnosticsModal.tsx | 105 ------------ .../diagnostics/diagnosticStatusBadge.styl | 13 -- .../diagnostics/diagnosticStatusBadge.tsx | 103 ----------- .../diagnostics/diagnosticStatuses.ts | 11 -- .../src/views/statements/diagnostics/index.ts | 12 -- 25 files changed, 534 insertions(+), 345 deletions(-) create mode 100644 pkg/ui/workspaces/cluster-ui/src/statementsDiagnostics/activateStatementDiagnosticsModal.scss delete mode 100644 pkg/ui/workspaces/db-console/src/views/statements/diagnostics/activateDiagnosticsModal.tsx delete mode 100644 pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatusBadge.styl delete mode 100644 pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatusBadge.tsx delete mode 100644 pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatuses.ts delete mode 100644 pkg/ui/workspaces/db-console/src/views/statements/diagnostics/index.ts diff --git a/pkg/server/statement_diagnostics_requests.go b/pkg/server/statement_diagnostics_requests.go index 3984e0b45bc2..9d82ef69ccc8 100644 --- a/pkg/server/statement_diagnostics_requests.go +++ b/pkg/server/statement_diagnostics_requests.go @@ -20,6 +20,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/server/serverpb" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/cockroach/pkg/sql/sessiondata" + "github.com/cockroachdb/cockroach/pkg/util/timeutil" "github.com/cockroachdb/errors" ) @@ -115,10 +116,11 @@ func (s *statusServer) CancelStatementDiagnosticsReport( return &response, nil } -// StatementDiagnosticsRequests retrieves all of the statement -// diagnostics requests in the `system.statement_diagnostics_requests` table. +// StatementDiagnosticsRequests retrieves all statement diagnostics +// requests in the `system.statement_diagnostics_requests` table that +// have not yet expired. func (s *statusServer) StatementDiagnosticsRequests( - ctx context.Context, req *serverpb.StatementDiagnosticsReportsRequest, + ctx context.Context, _ *serverpb.StatementDiagnosticsReportsRequest, ) (*serverpb.StatementDiagnosticsReportsResponse, error) { ctx = propagateGatewayMetadata(ctx) ctx = s.AnnotateCtx(ctx) @@ -179,6 +181,10 @@ func (s *statusServer) StatementDiagnosticsRequests( } if expiresAt, ok := row[6].(*tree.DTimestampTZ); ok { req.ExpiresAt = expiresAt.Time + // Don't return already expired requests. + if req.ExpiresAt.Before(timeutil.Now()) { + continue + } } } diff --git a/pkg/server/status_test.go b/pkg/server/status_test.go index fd84684bc9ef..3befff45479a 100644 --- a/pkg/server/status_test.go +++ b/pkg/server/status_test.go @@ -2446,6 +2446,50 @@ func TestStatementDiagnosticsCompleted(t *testing.T) { } } +func TestStatementDiagnosticsDoesNotReturnExpiredRequests(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{}) + defer s.Stopper().Stop(context.Background()) + db := sqlutils.MakeSQLRunner(sqlDB) + + statementFingerprint := "INSERT INTO test VALUES (_)" + expiresAfter := 5 * time.Millisecond + + // Create statement diagnostics request with defined expiry time. + req := &serverpb.CreateStatementDiagnosticsReportRequest{ + StatementFingerprint: statementFingerprint, + MinExecutionLatency: 500 * time.Millisecond, + ExpiresAfter: expiresAfter, + } + var resp serverpb.CreateStatementDiagnosticsReportResponse + if err := postStatusJSONProto(s, "stmtdiagreports", req, &resp); err != nil { + t.Fatal(err) + } + + // Wait for request to expire. + time.Sleep(expiresAfter) + + // Check that created statement diagnostics report is incomplete. + report := db.QueryStr(t, ` +SELECT completed +FROM system.statement_diagnostics_requests +WHERE statement_fingerprint = $1`, statementFingerprint) + + require.Equal(t, report[0][0], "false") + + // Check that expired report is not returned in API response. + var respGet serverpb.StatementDiagnosticsReportsResponse + if err := getStatusJSONProto(s, "stmtdiagreports", &respGet); err != nil { + t.Fatal(err) + } + + for _, report := range respGet.Reports { + require.NotEqual(t, report.StatementFingerprint, statementFingerprint) + } +} + func TestJobStatusResponse(t *testing.T) { defer leaktest.AfterTest(t)() defer log.Scope(t).Close(t) diff --git a/pkg/ui/workspaces/cluster-ui/src/api/statementDiagnosticsApi.ts b/pkg/ui/workspaces/cluster-ui/src/api/statementDiagnosticsApi.ts index 0a27c0c6c76c..51d5fd0c26ca 100644 --- a/pkg/ui/workspaces/cluster-ui/src/api/statementDiagnosticsApi.ts +++ b/pkg/ui/workspaces/cluster-ui/src/api/statementDiagnosticsApi.ts @@ -14,6 +14,7 @@ import { fetchData } from "src/api"; const STATEMENT_DIAGNOSTICS_PATH = "/_status/stmtdiagreports"; const CREATE_STATEMENT_DIAGNOSTICS_REPORT_PATH = "/_status/stmtdiagreports"; +type CreateStatementDiagnosticsReportRequestMessage = cockroach.server.serverpb.CreateStatementDiagnosticsReportRequest; type CreateStatementDiagnosticsReportResponseMessage = cockroach.server.serverpb.CreateStatementDiagnosticsReportResponse; export function getStatementDiagnosticsReports(): Promise< @@ -26,14 +27,12 @@ export function getStatementDiagnosticsReports(): Promise< } export function createStatementDiagnosticsReport( - statementsFingerprint: string, + req: CreateStatementDiagnosticsReportRequestMessage, ): Promise { return fetchData( cockroach.server.serverpb.CreateStatementDiagnosticsReportResponse, CREATE_STATEMENT_DIAGNOSTICS_REPORT_PATH, cockroach.server.serverpb.CreateStatementDiagnosticsReportRequest, - { - statement_fingerprint: statementsFingerprint, - }, + req, ); } diff --git a/pkg/ui/workspaces/cluster-ui/src/modal/modal.tsx b/pkg/ui/workspaces/cluster-ui/src/modal/modal.tsx index 71eae0e2bb05..11a3d2f6175f 100644 --- a/pkg/ui/workspaces/cluster-ui/src/modal/modal.tsx +++ b/pkg/ui/workspaces/cluster-ui/src/modal/modal.tsx @@ -22,6 +22,7 @@ export interface ModalProps { okText?: string; cancelText?: string; visible: boolean; + className?: string; } const cx = classNames.bind(styles); @@ -34,11 +35,12 @@ export const Modal: React.FC = ({ cancelText, visible, title, + className, }) => { return ( {title}} - className={cx("crl-modal")} + className={cx("crl-modal", className)} visible={visible} closeIcon={
diff --git a/pkg/ui/workspaces/cluster-ui/src/statementDetails/diagnostics/diagnosticsView.spec.tsx b/pkg/ui/workspaces/cluster-ui/src/statementDetails/diagnostics/diagnosticsView.spec.tsx index e133da5c7706..27bd9813cfad 100644 --- a/pkg/ui/workspaces/cluster-ui/src/statementDetails/diagnostics/diagnosticsView.spec.tsx +++ b/pkg/ui/workspaces/cluster-ui/src/statementDetails/diagnostics/diagnosticsView.spec.tsx @@ -11,7 +11,7 @@ import React from "react"; import { assert } from "chai"; import { mount, ReactWrapper } from "enzyme"; -import sinon, { SinonSpy } from "sinon"; +import sinon from "sinon"; import Long from "long"; import { MemoryRouter } from "react-router-dom"; import { cockroach } from "@cockroachlabs/crdb-protobuf-client"; @@ -25,6 +25,8 @@ type IStatementDiagnosticsReport = cockroach.server.serverpb.IStatementDiagnosti const sandbox = sinon.createSandbox(); +const activateDiagnosticsRef = { current: { showModalFor: jest.fn() } }; + function generateDiagnosticsRequest( extendObject: Partial = {}, ): IStatementDiagnosticsReport { @@ -42,12 +44,10 @@ function generateDiagnosticsRequest( describe("DiagnosticsView", () => { let wrapper: ReactWrapper; - let activateFn: SinonSpy; const statementFingerprint = "some-id"; beforeEach(() => { sandbox.reset(); - activateFn = sandbox.spy(); }); describe("With Empty state", () => { @@ -55,8 +55,8 @@ describe("DiagnosticsView", () => { wrapper = mount( {}} @@ -65,10 +65,12 @@ describe("DiagnosticsView", () => { ); }); - it("calls activate callback with statementFingerprintId when click on Activate button", () => { + it("opens the statement diagnostics modal when Activate button is clicked", () => { const activateButtonComponent = wrapper.find(Button).first(); activateButtonComponent.simulate("click"); - activateFn.calledOnceWith(statementFingerprint); + expect(activateDiagnosticsRef.current.showModalFor).toBeCalledWith( + statementFingerprint, + ); }); }); @@ -82,8 +84,8 @@ describe("DiagnosticsView", () => { wrapper = mount( {}} @@ -96,12 +98,14 @@ describe("DiagnosticsView", () => { assert.isTrue(wrapper.find(Table).exists()); }); - it("calls activate callback with statementFingerprintId when click on Activate button", () => { + it("opens the statement diagnostics modal when Activate button is clicked", () => { const activateButtonComponent = wrapper .findWhere(n => n.prop("children") === "Activate diagnostics") .first(); activateButtonComponent.simulate("click"); - activateFn.calledOnceWith(statementFingerprint); + expect(activateDiagnosticsRef.current.showModalFor).toBeCalledWith( + statementFingerprint, + ); }); it("Activate button is hidden if diagnostics is requested and waiting query", () => { @@ -112,8 +116,8 @@ describe("DiagnosticsView", () => { wrapper = mount( {}} diff --git a/pkg/ui/workspaces/cluster-ui/src/statementDetails/diagnostics/diagnosticsView.tsx b/pkg/ui/workspaces/cluster-ui/src/statementDetails/diagnostics/diagnosticsView.tsx index f86dbda0ec61..3d2cb93750ca 100644 --- a/pkg/ui/workspaces/cluster-ui/src/statementDetails/diagnostics/diagnosticsView.tsx +++ b/pkg/ui/workspaces/cluster-ui/src/statementDetails/diagnostics/diagnosticsView.tsx @@ -8,7 +8,7 @@ // by the Apache License, Version 2.0, included in the file // licenses/APL.txt. -import React, { useCallback } from "react"; +import React from "react"; import { Link } from "react-router-dom"; import moment from "moment"; import classnames from "classnames/bind"; @@ -17,18 +17,19 @@ import { Button, Icon } from "@cockroachlabs/ui-components"; import { Text, TextTypes } from "src/text"; import { Table, ColumnsConfig } from "src/table"; import { SummaryCard } from "src/summaryCard"; -import { DiagnosticStatusBadge } from "src/statementsDiagnostics"; +import { + ActivateDiagnosticsModalRef, + DiagnosticStatusBadge, +} from "src/statementsDiagnostics"; import emptyListResultsImg from "src/assets/emptyState/empty-list-results.svg"; import { getDiagnosticsStatus, sortByCompletedField, sortByRequestedAtField, } from "./diagnosticsUtils"; -import { statementDiagnostics } from "src/util/docs"; import { EmptyTable } from "src/empty"; import styles from "./diagnosticsView.module.scss"; import { getBasePath } from "../../api"; -import { Anchor } from "../../anchor"; type IStatementDiagnosticsReport = cockroach.server.serverpb.IStatementDiagnosticsReport; @@ -36,10 +37,10 @@ export interface DiagnosticsViewStateProps { hasData: boolean; diagnosticsReports: cockroach.server.serverpb.IStatementDiagnosticsReport[]; showDiagnosticsViewLink?: boolean; + activateDiagnosticsRef: React.RefObject; } export interface DiagnosticsViewDispatchProps { - activate: (statementFingerprint: string) => void; dismissAlertMessage: () => void; onDownloadDiagnosticBundleClick?: (statementFingerprint: string) => void; onSortingChange?: ( @@ -72,30 +73,24 @@ const NavButton: React.FC = props => ( ); export const EmptyDiagnosticsView = ({ - activate, statementFingerprint, showDiagnosticsViewLink, + activateDiagnosticsRef, }: DiagnosticsViewProps) => { - const onActivateButtonClick = useCallback(() => { - activate(statementFingerprint); - }, [activate, statementFingerprint]); return ( - When you activate statement diagnostics, CockroachDB will wait for the - next query that matches this statement fingerprint. A download button - will appear on the statement list and detail pages when the query is - ready. The statement diagnostic will include EXPLAIN plans, table - statistics, and traces.{" "} - Learn More - - } footer={
- {showDiagnosticsViewLink && ( @@ -186,11 +181,6 @@ export class DiagnosticsView extends React.Component< }, ]; - onActivateButtonClick = () => { - const { activate, statementFingerprint } = this.props; - activate(statementFingerprint); - }; - componentWillUnmount() { this.props.dismissAlertMessage(); } @@ -202,7 +192,13 @@ export class DiagnosticsView extends React.Component< }; render() { - const { hasData, diagnosticsReports, showDiagnosticsViewLink } = this.props; + const { + hasData, + diagnosticsReports, + showDiagnosticsViewLink, + statementFingerprint, + activateDiagnosticsRef, + } = this.props; const canRequestDiagnostics = diagnosticsReports.every( diagnostic => diagnostic.completed, @@ -227,7 +223,11 @@ export class DiagnosticsView extends React.Component< Statement diagnostics {canRequestDiagnostics && (
); }, diff --git a/pkg/ui/workspaces/cluster-ui/src/statementsPage/statementsPage.tsx b/pkg/ui/workspaces/cluster-ui/src/statementsPage/statementsPage.tsx index 6a30b77b2b71..67ee37c7058b 100644 --- a/pkg/ui/workspaces/cluster-ui/src/statementsPage/statementsPage.tsx +++ b/pkg/ui/workspaces/cluster-ui/src/statementsPage/statementsPage.tsx @@ -56,9 +56,10 @@ import { import { ISortedTablePagination } from "../sortedtable"; import styles from "./statementsPage.module.scss"; import { EmptyStatementsPlaceholder } from "./emptyStatementsPlaceholder"; -import { cockroach } from "@cockroachlabs/crdb-protobuf-client"; +import { cockroach, google } from "@cockroachlabs/crdb-protobuf-client"; type IStatementDiagnosticsReport = cockroach.server.serverpb.IStatementDiagnosticsReport; +type IDuration = google.protobuf.IDuration; import sortableTableStyles from "src/sortedtable/sortedtable.module.scss"; import ColumnsSelector from "../columnsSelector/columnsSelector"; import { SelectOption } from "../multiSelectCheckbox/multiSelectCheckbox"; @@ -88,7 +89,11 @@ export interface StatementsPageDispatchProps { refreshStatementDiagnosticsRequests: () => void; resetSQLStats: () => void; dismissAlertMessage: () => void; - onActivateStatementDiagnostics: (statement: string) => void; + onActivateStatementDiagnostics: ( + statement: string, + minExecLatency: IDuration, + expiresAfter: IDuration, + ) => void; onDiagnosticsModalOpen?: (statement: string) => void; onSearchComplete?: (query: string) => void; onPageChanged?: (newPage: number) => void; diff --git a/pkg/ui/workspaces/cluster-ui/src/statementsPage/statementsPageConnected.tsx b/pkg/ui/workspaces/cluster-ui/src/statementsPage/statementsPageConnected.tsx index 00bfabc7de46..90fc8dd71af9 100644 --- a/pkg/ui/workspaces/cluster-ui/src/statementsPage/statementsPageConnected.tsx +++ b/pkg/ui/workspaces/cluster-ui/src/statementsPage/statementsPageConnected.tsx @@ -40,6 +40,12 @@ import { selectIsTenant } from "../store/uiConfig"; import { nodeRegionsByIDSelector } from "../store/nodes"; import { StatementsRequest } from "src/api/statementsApi"; import { TimeScale } from "../timeScaleDropdown"; +import { cockroach, google } from "@cockroachlabs/crdb-protobuf-client"; + +type IDuration = google.protobuf.IDuration; + +const CreateStatementDiagnosticsReportRequest = + cockroach.server.serverpb.CreateStatementDiagnosticsReportRequest; export const ConnectedStatementsPage = withRouter( connect< @@ -82,9 +88,19 @@ export const ConnectedStatementsPage = withRouter( value: false, }), ), - onActivateStatementDiagnostics: (statementFingerprint: string) => { + onActivateStatementDiagnostics: ( + statementFingerprint: string, + minExecLatency: IDuration, + expiresAfter: IDuration, + ) => { dispatch( - statementDiagnosticsActions.createReport(statementFingerprint), + statementDiagnosticsActions.createReport( + new CreateStatementDiagnosticsReportRequest({ + statement_fingerprint: statementFingerprint, + min_execution_latency: minExecLatency, + expires_after: expiresAfter, + }), + ), ); dispatch( analyticsActions.track({ diff --git a/pkg/ui/workspaces/cluster-ui/src/store/statementDiagnostics/statementDiagnostics.reducer.ts b/pkg/ui/workspaces/cluster-ui/src/store/statementDiagnostics/statementDiagnostics.reducer.ts index fc169776da7b..a2998302ba7d 100644 --- a/pkg/ui/workspaces/cluster-ui/src/store/statementDiagnostics/statementDiagnostics.reducer.ts +++ b/pkg/ui/workspaces/cluster-ui/src/store/statementDiagnostics/statementDiagnostics.reducer.ts @@ -12,6 +12,7 @@ import { createSlice, PayloadAction } from "@reduxjs/toolkit"; import { cockroach } from "@cockroachlabs/crdb-protobuf-client"; import { DOMAIN_NAME, noopReducer } from "../utils"; +type CreateStatementDiagnosticsReportRequest = cockroach.server.serverpb.CreateStatementDiagnosticsReportRequest; type StatementDiagnosticsReportsResponse = cockroach.server.serverpb.StatementDiagnosticsReportsResponse; export type StatementDiagnosticsState = { @@ -48,7 +49,10 @@ const statementDiagnosticsSlice = createSlice({ refresh: noopReducer, request: noopReducer, invalidated: noopReducer, - createReport: (_state, _action: PayloadAction) => {}, + createReport: ( + _state, + _action: PayloadAction, + ) => {}, createReportCompleted: noopReducer, createReportFailed: (_state, _action: PayloadAction) => {}, openNewDiagnosticsModal: (_state, _action: PayloadAction) => {}, diff --git a/pkg/ui/workspaces/cluster-ui/src/store/statementDiagnostics/statementDiagnostics.sagas.spec.ts b/pkg/ui/workspaces/cluster-ui/src/store/statementDiagnostics/statementDiagnostics.sagas.spec.ts index 0f6a8ece1619..ad0b3865c397 100644 --- a/pkg/ui/workspaces/cluster-ui/src/store/statementDiagnostics/statementDiagnostics.sagas.spec.ts +++ b/pkg/ui/workspaces/cluster-ui/src/store/statementDiagnostics/statementDiagnostics.sagas.spec.ts @@ -11,7 +11,7 @@ import { expectSaga } from "redux-saga-test-plan"; import { throwError } from "redux-saga-test-plan/providers"; import { call } from "redux-saga/effects"; -import { cockroach } from "@cockroachlabs/crdb-protobuf-client"; +import { cockroach, google } from "@cockroachlabs/crdb-protobuf-client"; import Long from "long"; import { createDiagnosticsReportSaga, @@ -24,12 +24,26 @@ import { getStatementDiagnosticsReports, } from "src/api/statementDiagnosticsApi"; +const CreateStatementDiagnosticsReportRequest = + cockroach.server.serverpb.CreateStatementDiagnosticsReportRequest; const CreateStatementDiagnosticsReportResponse = cockroach.server.serverpb.CreateStatementDiagnosticsReportResponse; describe("statementsDiagnostics sagas", () => { describe("createDiagnosticsReportSaga", () => { const statementFingerprint = "SELECT * FROM table"; + const minExecLatency = new google.protobuf.Duration({ + seconds: new Long(100), + }); + const expiresAfter = new google.protobuf.Duration({ + seconds: new Long(0), // Setting expiresAfter to 0 means the request won't expire. + }); + + const request = new CreateStatementDiagnosticsReportRequest({ + statement_fingerprint: statementFingerprint, + min_execution_latency: minExecLatency, + expires_after: expiresAfter, + }); const report = new CreateStatementDiagnosticsReportResponse({ report: { @@ -44,15 +58,9 @@ describe("statementsDiagnostics sagas", () => { ); it("successful request", () => { - expectSaga( - createDiagnosticsReportSaga, - actions.createReport(statementFingerprint), - ) + expectSaga(createDiagnosticsReportSaga, actions.createReport(request)) .provide([ - [ - call(createStatementDiagnosticsReport, statementFingerprint), - report, - ], + [call(createStatementDiagnosticsReport, request), report], [call(getStatementDiagnosticsReports), reportsResponse], ]) .put(actions.createReportCompleted()) @@ -68,15 +76,9 @@ describe("statementsDiagnostics sagas", () => { it("failed request", () => { const error = new Error("Failed request"); - expectSaga( - createDiagnosticsReportSaga, - actions.createReport(statementFingerprint), - ) + expectSaga(createDiagnosticsReportSaga, actions.createReport(request)) .provide([ - [ - call(createStatementDiagnosticsReport, statementFingerprint), - throwError(error), - ], + [call(createStatementDiagnosticsReport, request), throwError(error)], [call(getStatementDiagnosticsReports), reportsResponse], ]) .put(actions.createReportFailed(error)) diff --git a/pkg/ui/workspaces/cluster-ui/src/util/convert.spec.ts b/pkg/ui/workspaces/cluster-ui/src/util/convert.spec.ts index 66b5628a5f80..b3334767de29 100644 --- a/pkg/ui/workspaces/cluster-ui/src/util/convert.spec.ts +++ b/pkg/ui/workspaces/cluster-ui/src/util/convert.spec.ts @@ -18,6 +18,7 @@ import { TimestampToNumber, LongToMoment, DurationToNumber, + NumberToDuration, } from "./convert"; import { fromNumber } from "long"; @@ -113,4 +114,21 @@ describe("Test convert functions", (): void => { expect(DurationToNumber(null, 5)).toEqual(5); }); }); + + describe("NumberToDuration", (): void => { + it("should convert a number representing seconds to a Duration", (): void => { + SECONDS.forEach((seconds: number) => { + const duration = NumberToDuration(seconds); + expect(DurationToNumber(duration)).toEqual(seconds); + }); + }); + + it("should convert a number representing seconds and milliseconds to a Duration", (): void => { + const MixedTimes = [0.005, 11.035, 1.003, 0, 0.025]; + MixedTimes.forEach((time: number) => { + const duration = NumberToDuration(time); + expect(DurationToNumber(duration)).toEqual(time); + }); + }); + }); }); diff --git a/pkg/ui/workspaces/cluster-ui/src/util/convert.ts b/pkg/ui/workspaces/cluster-ui/src/util/convert.ts index 7431f551d9c3..8021a656bf22 100644 --- a/pkg/ui/workspaces/cluster-ui/src/util/convert.ts +++ b/pkg/ui/workspaces/cluster-ui/src/util/convert.ts @@ -11,6 +11,7 @@ import moment from "moment"; import * as protos from "@cockroachlabs/crdb-protobuf-client"; +import { fromNumber } from "long"; /** * NanoToMilli converts a nanoseconds value into milliseconds. @@ -113,6 +114,19 @@ export function DurationToNumber( return duration.seconds.toNumber() + NanoToMilli(duration.nanos) * 1e-3; } +/** + * NumberToDuration converts a number representing a duration in seconds + * to a Duration object. + */ +export function NumberToDuration( + seconds?: number, +): protos.google.protobuf.IDuration { + return new protos.google.protobuf.Duration({ + seconds: fromNumber(seconds), + nanos: SecondsToNano(seconds - Math.floor(seconds)), + }); +} + // durationFromISO8601String function converts a string date in ISO8601 format to moment.Duration export const durationFromISO8601String = (value: string): moment.Duration => { if (!value) { diff --git a/pkg/ui/workspaces/db-console/src/redux/statements/statementsActions.ts b/pkg/ui/workspaces/db-console/src/redux/statements/statementsActions.ts index a3b332e57dbc..7ba41c63873a 100644 --- a/pkg/ui/workspaces/db-console/src/redux/statements/statementsActions.ts +++ b/pkg/ui/workspaces/db-console/src/redux/statements/statementsActions.ts @@ -10,6 +10,8 @@ import { Action } from "redux"; import { PayloadAction } from "src/interfaces/action"; +import { google } from "@cockroachlabs/crdb-protobuf-client"; +import IDuration = google.protobuf.IDuration; import { TimeScale } from "src/redux/timewindow"; export const CREATE_STATEMENT_DIAGNOSTICS_REPORT = @@ -25,13 +27,23 @@ export type DiagnosticsReportPayload = { statementFingerprint: string; }; +export type CreateStatementDiagnosticsReportPayload = { + statementFingerprint: string; + minExecLatency: IDuration; + expiresAfter: IDuration; +}; + export function createStatementDiagnosticsReportAction( statementFingerprint: string, -): PayloadAction { + minExecLatency: IDuration, + expiresAfter: IDuration, +): PayloadAction { return { type: CREATE_STATEMENT_DIAGNOSTICS_REPORT, payload: { statementFingerprint, + minExecLatency, + expiresAfter, }, }; } diff --git a/pkg/ui/workspaces/db-console/src/redux/statements/statementsSagas.spec.ts b/pkg/ui/workspaces/db-console/src/redux/statements/statementsSagas.spec.ts index 4d51d9654832..e9cc8b57f698 100644 --- a/pkg/ui/workspaces/db-console/src/redux/statements/statementsSagas.spec.ts +++ b/pkg/ui/workspaces/db-console/src/redux/statements/statementsSagas.spec.ts @@ -18,20 +18,28 @@ import { createStatementDiagnosticsReportAction, } from "./statementsActions"; import { createStatementDiagnosticsReport } from "src/util/api"; -import { cockroach } from "src/js/protos"; +import { cockroach, google } from "src/js/protos"; import CreateStatementDiagnosticsReportRequest = cockroach.server.serverpb.CreateStatementDiagnosticsReportRequest; import { throwError } from "redux-saga-test-plan/providers"; +import Duration = google.protobuf.Duration; +import Long from "long"; describe("statementsSagas", () => { describe("requestDiagnostics generator", () => { - it("calls api#createStatementDiagnosticsReport with statement fingerprint as payload", () => { + it("calls api#createStatementDiagnosticsReport with statement fingerprint, min exec latency, and expires after fields as payload", () => { const statementFingerprint = "some-id"; + const minExecLatency = new Duration({ seconds: new Long(10) }); + const expiresAfter = new Duration({ seconds: new Long(15 * 60) }); const action = createStatementDiagnosticsReportAction( statementFingerprint, + minExecLatency, + expiresAfter, ); - const diagnosticsReportRequest = new CreateStatementDiagnosticsReportRequest( + const createDiagnosticsReportRequest = new CreateStatementDiagnosticsReportRequest( { statement_fingerprint: statementFingerprint, + min_execution_latency: minExecLatency, + expires_after: expiresAfter, }, ); @@ -39,7 +47,7 @@ describe("statementsSagas", () => { .provide([ [call.fn(createStatementDiagnosticsReport), Promise.resolve()], ]) - .call(createStatementDiagnosticsReport, diagnosticsReportRequest) + .call(createStatementDiagnosticsReport, createDiagnosticsReportRequest) .put(createStatementDiagnosticsReportCompleteAction()) .dispatch(action) .run(); @@ -48,10 +56,18 @@ describe("statementsSagas", () => { it("calls dispatched failed action if api#createStatementDiagnosticsReport request failed ", () => { const statementFingerprint = "some-id"; - const action = createStatementDiagnosticsReportAction(statementFingerprint); - const diagnosticsReportRequest = new CreateStatementDiagnosticsReportRequest( + const minExecLatency = new Duration({ seconds: new Long(10) }); + const expiresAfter = new Duration({ seconds: new Long(15 * 60) }); + const action = createStatementDiagnosticsReportAction( + statementFingerprint, + minExecLatency, + expiresAfter, + ); + const createDiagnosticsReportRequest = new CreateStatementDiagnosticsReportRequest( { statement_fingerprint: statementFingerprint, + min_execution_latency: minExecLatency, + expires_after: expiresAfter, }, ); @@ -59,7 +75,7 @@ describe("statementsSagas", () => { .provide([ [call.fn(createStatementDiagnosticsReport), throwError(new Error())], ]) - .call(createStatementDiagnosticsReport, diagnosticsReportRequest) + .call(createStatementDiagnosticsReport, createDiagnosticsReportRequest) .put(createStatementDiagnosticsReportFailedAction()) .dispatch(action) .run(); diff --git a/pkg/ui/workspaces/db-console/src/redux/statements/statementsSagas.ts b/pkg/ui/workspaces/db-console/src/redux/statements/statementsSagas.ts index e6a76ec4688c..b00cfc0a81ae 100644 --- a/pkg/ui/workspaces/db-console/src/redux/statements/statementsSagas.ts +++ b/pkg/ui/workspaces/db-console/src/redux/statements/statementsSagas.ts @@ -14,7 +14,7 @@ import { PayloadAction } from "src/interfaces/action"; import { createStatementDiagnosticsReport } from "src/util/api"; import { CREATE_STATEMENT_DIAGNOSTICS_REPORT, - DiagnosticsReportPayload, + CreateStatementDiagnosticsReportPayload, createStatementDiagnosticsReportCompleteAction, createStatementDiagnosticsReportFailedAction, SET_COMBINED_STATEMENTS_TIME_SCALE, @@ -34,14 +34,21 @@ import { TimeScale, toDateRange } from "@cockroachlabs/cluster-ui"; import Long from "long"; export function* createDiagnosticsReportSaga( - action: PayloadAction, + action: PayloadAction, ) { - const { statementFingerprint } = action.payload; - const diagnosticsReportRequest = new CreateStatementDiagnosticsReportRequest({ - statement_fingerprint: statementFingerprint, - }); + const { statementFingerprint, minExecLatency, expiresAfter } = action.payload; + const createDiagnosticsReportRequest = new CreateStatementDiagnosticsReportRequest( + { + statement_fingerprint: statementFingerprint, + min_execution_latency: minExecLatency, + expires_after: expiresAfter, + }, + ); try { - yield call(createStatementDiagnosticsReport, diagnosticsReportRequest); + yield call( + createStatementDiagnosticsReport, + createDiagnosticsReportRequest, + ); yield put(createStatementDiagnosticsReportCompleteAction()); yield put(invalidateStatementDiagnosticsRequests()); // PUT expects action with `type` field which isn't defined in `refresh` ThunkAction interface diff --git a/pkg/ui/workspaces/db-console/src/views/reports/containers/statementDiagnosticsHistory/index.tsx b/pkg/ui/workspaces/db-console/src/views/reports/containers/statementDiagnosticsHistory/index.tsx index 137a083ffbfa..db162fedce55 100644 --- a/pkg/ui/workspaces/db-console/src/views/reports/containers/statementDiagnosticsHistory/index.tsx +++ b/pkg/ui/workspaces/db-console/src/views/reports/containers/statementDiagnosticsHistory/index.tsx @@ -29,7 +29,6 @@ import { invalidateStatementDiagnosticsRequests, refreshStatementDiagnosticsRequests, } from "src/redux/apiReducers"; -import { DiagnosticStatusBadge } from "src/views/statements/diagnostics/diagnosticStatusBadge"; import "./statementDiagnosticsHistoryView.styl"; import { cockroach } from "src/js/protos"; import IStatementDiagnosticsReport = cockroach.server.serverpb.IStatementDiagnosticsReport; @@ -43,6 +42,7 @@ import { EmptyTable, shortStatement, getDiagnosticsStatus, + DiagnosticStatusBadge, SortedTable, SortSetting, ColumnDescriptor, diff --git a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/activateDiagnosticsModal.tsx b/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/activateDiagnosticsModal.tsx deleted file mode 100644 index a4b44c5023ba..000000000000 --- a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/activateDiagnosticsModal.tsx +++ /dev/null @@ -1,105 +0,0 @@ -// Copyright 2020 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 React, { - forwardRef, - useState, - useCallback, - useImperativeHandle, -} from "react"; -import { connect } from "react-redux"; -import { Anchor, Modal, Text } from "src/components"; -import { createStatementDiagnosticsReportAction } from "src/redux/statements"; -import { AppDispatch } from "src/redux/state"; -import { - invalidateStatementDiagnosticsRequests, - refreshStatementDiagnosticsRequests, -} from "src/redux/apiReducers"; -import { statementDiagnostics } from "src/util/docs"; -import { - trackActivateDiagnostics, - trackDiagnosticsModalOpen, -} from "src/util/analytics"; -export type ActivateDiagnosticsModalProps = MapDispatchToProps; - -export interface ActivateDiagnosticsModalRef { - showModalFor: (statement: string) => void; -} - -const ActivateDiagnosticsModal = ( - props: ActivateDiagnosticsModalProps, - ref: React.RefObject, -) => { - const { activate } = props; - const [visible, setVisible] = useState(false); - const [statement, setStatement] = useState(); - - const onOkHandler = useCallback(() => { - activate(statement); - trackActivateDiagnostics(statement); - setVisible(false); - }, [activate, statement]); - - const onCancelHandler = useCallback(() => setVisible(false), []); - - useImperativeHandle(ref, () => { - return { - showModalFor: (forwardStatement: string) => { - setStatement(forwardStatement); - trackDiagnosticsModalOpen(forwardStatement); - setVisible(true); - }, - }; - }); - - return ( - - - When you activate statement diagnostics, CockroachDB will wait for the - next query that matches this statement fingerprint. - -

- - A download button will appear on the statement list and detail pages - when the query is ready. The download will include EXPLAIN plans, table - statistics, and traces.{" "} - Learn more - - - ); -}; - -interface MapDispatchToProps { - activate: (statement: string) => void; - refreshDiagnosticsReports: () => void; -} - -const mapDispatchToProps = (dispatch: AppDispatch): MapDispatchToProps => ({ - activate: (statement: string) => - dispatch(createStatementDiagnosticsReportAction(statement)), - refreshDiagnosticsReports: () => { - dispatch(invalidateStatementDiagnosticsRequests()); - dispatch(refreshStatementDiagnosticsRequests()); - }, -}); - -export default connect( - null, - mapDispatchToProps, - null, - { forwardRef: true }, -)(forwardRef(ActivateDiagnosticsModal)); diff --git a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatusBadge.styl b/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatusBadge.styl deleted file mode 100644 index 775aa7ad27d4..000000000000 --- a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatusBadge.styl +++ /dev/null @@ -1,13 +0,0 @@ -// Copyright 2020 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. - -.diagnostic-status-badge - &__content - width fit-content diff --git a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatusBadge.tsx b/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatusBadge.tsx deleted file mode 100644 index ee7585410a22..000000000000 --- a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatusBadge.tsx +++ /dev/null @@ -1,103 +0,0 @@ -// Copyright 2020 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 React from "react"; - -import { Badge, Anchor, Tooltip } from "src/components"; -import { statementDiagnostics } from "src/util/docs"; -import { DiagnosticStatuses } from "./diagnosticStatuses"; -import "./diagnosticStatusBadge.styl"; - -interface OwnProps { - status: DiagnosticStatuses; - enableTooltip?: boolean; -} - -function mapDiagnosticsStatusToBadge(diagnosticsStatus: DiagnosticStatuses) { - switch (diagnosticsStatus) { - case "READY": - return "success"; - case "WAITING": - return "info"; - case "ERROR": - return "danger"; - default: - return "info"; - } -} - -function mapStatusToDescription(diagnosticsStatus: DiagnosticStatuses) { - switch (diagnosticsStatus) { - case "READY": - return ( -

-

- {"The most recent "} - - diagnostics - - { - " for this SQL statement fingerprint are ready to download. Access the full history of diagnostics for the fingerprint in the Statement Details page." - } -

-
- ); - case "WAITING": - return ( -
-

- CockroachDB is waiting for the next SQL statement that matches this - fingerprint. -

-

- {"When the most recent "} - - diagnostics - - {" are ready to download, a link will appear in this row."} -

-
- ); - case "ERROR": - return ( -
- { - "There was an error when attempting to collect diagnostics for this statement. Please try activating again. " - } - - Learn more - -
- ); - default: - return ""; - } -} - -export function DiagnosticStatusBadge(props: OwnProps) { - const { status, enableTooltip } = props; - const tooltipContent = mapStatusToDescription(status); - - if (!enableTooltip) { - return ; - } - - return ( - -
- -
- - ); -} - -DiagnosticStatusBadge.defaultProps = { - enableTooltip: true, -}; diff --git a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatuses.ts b/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatuses.ts deleted file mode 100644 index 96ae8058a55a..000000000000 --- a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/diagnosticStatuses.ts +++ /dev/null @@ -1,11 +0,0 @@ -// Copyright 2020 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. - -export type DiagnosticStatuses = "READY" | "WAITING" | "ERROR"; diff --git a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/index.ts b/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/index.ts deleted file mode 100644 index 7b0ee512512e..000000000000 --- a/pkg/ui/workspaces/db-console/src/views/statements/diagnostics/index.ts +++ /dev/null @@ -1,12 +0,0 @@ -// Copyright 2020 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. - -export * from "./activateDiagnosticsModal"; -export * from "./diagnosticStatusBadge"; From 9cb4298896f221255d2f00c3341e21362c885cb4 Mon Sep 17 00:00:00 2001 From: Yevgeniy Miretskiy Date: Tue, 21 Dec 2021 10:07:52 -0500 Subject: [PATCH 2/2] streamingccl: Ensure appropriate privileges when replicating. Add checks to replication stream manager to ensure appropriate privileges and require enterprise license when executing streaming replication. Release Notes: None --- .../streamingccl/streamproducer/BUILD.bazel | 4 + .../streamproducer/replication_manager.go | 40 +++++++-- .../replication_manager_test.go | 88 +++++++++++++++++++ pkg/sql/sem/builtins/replication_builtins.go | 12 +-- pkg/streaming/api.go | 9 +- 5 files changed, 136 insertions(+), 17 deletions(-) create mode 100644 pkg/ccl/streamingccl/streamproducer/replication_manager_test.go diff --git a/pkg/ccl/streamingccl/streamproducer/BUILD.bazel b/pkg/ccl/streamingccl/streamproducer/BUILD.bazel index 8c07b30b9157..7b6ffe19290d 100644 --- a/pkg/ccl/streamingccl/streamproducer/BUILD.bazel +++ b/pkg/ccl/streamingccl/streamproducer/BUILD.bazel @@ -54,6 +54,7 @@ go_test( srcs = [ "main_test.go", "producer_job_test.go", + "replication_manager_test.go", "replication_stream_test.go", ], embed = [":streamproducer"], @@ -76,8 +77,11 @@ go_test( "//pkg/security/securitytest", "//pkg/server", "//pkg/settings/cluster", + "//pkg/sql", "//pkg/sql/catalog/catalogkv", "//pkg/sql/distsql", + "//pkg/sql/sem/tree", + "//pkg/sql/sessiondatapb", "//pkg/streaming", "//pkg/testutils", "//pkg/testutils/serverutils", diff --git a/pkg/ccl/streamingccl/streamproducer/replication_manager.go b/pkg/ccl/streamingccl/streamproducer/replication_manager.go index 1ce6a3e8e2e3..c792778f0cfe 100644 --- a/pkg/ccl/streamingccl/streamproducer/replication_manager.go +++ b/pkg/ccl/streamingccl/streamproducer/replication_manager.go @@ -9,8 +9,12 @@ package streamproducer import ( + "github.com/cockroachdb/cockroach/pkg/ccl/utilccl" "github.com/cockroachdb/cockroach/pkg/jobs/jobspb" "github.com/cockroachdb/cockroach/pkg/kv" + "github.com/cockroachdb/cockroach/pkg/sql" + "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode" + "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/cockroach/pkg/streaming" "github.com/cockroachdb/cockroach/pkg/util/hlc" @@ -19,7 +23,7 @@ import ( type replicationStreamManagerImpl struct{} // CompleteStreamIngestion implements ReplicationStreamManager interface. -func (r replicationStreamManagerImpl) CompleteStreamIngestion( +func (r *replicationStreamManagerImpl) CompleteStreamIngestion( evalCtx *tree.EvalContext, txn *kv.Txn, streamID streaming.StreamID, @@ -29,14 +33,14 @@ func (r replicationStreamManagerImpl) CompleteStreamIngestion( } // StartReplicationStream implements ReplicationStreamManager interface. -func (r replicationStreamManagerImpl) StartReplicationStream( +func (r *replicationStreamManagerImpl) StartReplicationStream( evalCtx *tree.EvalContext, txn *kv.Txn, tenantID uint64, ) (streaming.StreamID, error) { return startReplicationStreamJob(evalCtx, txn, tenantID) } // UpdateReplicationStreamProgress implements ReplicationStreamManager interface. -func (r replicationStreamManagerImpl) UpdateReplicationStreamProgress( +func (r *replicationStreamManagerImpl) UpdateReplicationStreamProgress( evalCtx *tree.EvalContext, streamID streaming.StreamID, frontier hlc.Timestamp, txn *kv.Txn, ) (jobspb.StreamReplicationStatus, error) { return heartbeatReplicationStream(evalCtx, streamID, frontier, txn) @@ -45,14 +49,36 @@ func (r replicationStreamManagerImpl) UpdateReplicationStreamProgress( // StreamPartition returns a value generator which yields events for the specified partition. // opaqueSpec contains streampb.PartitionSpec protocol message. // streamID specifies the streaming job this partition belongs too. -func (r replicationStreamManagerImpl) StreamPartition( +func (r *replicationStreamManagerImpl) StreamPartition( evalCtx *tree.EvalContext, streamID streaming.StreamID, opaqueSpec []byte, ) (tree.ValueGenerator, error) { return streamPartition(evalCtx, streamID, opaqueSpec) } -func init() { - streaming.GetReplicationStreamManagerHook = func() (streaming.ReplicationStreamManager, error) { - return &replicationStreamManagerImpl{}, nil +func newReplicationStreamManagerWithPrivilegesCheck( + evalCtx *tree.EvalContext, +) (streaming.ReplicationStreamManager, error) { + isAdmin, err := evalCtx.SessionAccessor.HasAdminRole(evalCtx.Context) + if err != nil { + return nil, err + } + + if !isAdmin { + return nil, + pgerror.New(pgcode.InsufficientPrivilege, "replication restricted to ADMIN role") + } + + execCfg := evalCtx.Planner.ExecutorConfig().(*sql.ExecutorConfig) + enterpriseCheckErr := utilccl.CheckEnterpriseEnabled( + execCfg.Settings, execCfg.ClusterID(), execCfg.Organization(), "REPLICATION") + if enterpriseCheckErr != nil { + return nil, pgerror.Wrap(enterpriseCheckErr, + pgcode.InsufficientPrivilege, "replication requires enterprise license") } + + return &replicationStreamManagerImpl{}, nil +} + +func init() { + streaming.GetReplicationStreamManagerHook = newReplicationStreamManagerWithPrivilegesCheck } diff --git a/pkg/ccl/streamingccl/streamproducer/replication_manager_test.go b/pkg/ccl/streamingccl/streamproducer/replication_manager_test.go new file mode 100644 index 000000000000..2fe61b6c5a5f --- /dev/null +++ b/pkg/ccl/streamingccl/streamproducer/replication_manager_test.go @@ -0,0 +1,88 @@ +// Copyright 2021 The Cockroach Authors. +// +// Licensed as a CockroachDB Enterprise file under the Cockroach Community +// License (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt + +package streamproducer + +import ( + "context" + "fmt" + "testing" + + "github.com/cockroachdb/cockroach/pkg/base" + "github.com/cockroachdb/cockroach/pkg/ccl/utilccl" + "github.com/cockroachdb/cockroach/pkg/security" + "github.com/cockroachdb/cockroach/pkg/sql" + "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" + "github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb" + "github.com/cockroachdb/cockroach/pkg/streaming" + "github.com/cockroachdb/cockroach/pkg/testutils/serverutils" + "github.com/cockroachdb/cockroach/pkg/testutils/sqlutils" + "github.com/cockroachdb/cockroach/pkg/util/leaktest" + "github.com/cockroachdb/cockroach/pkg/util/log" + "github.com/cockroachdb/cockroach/pkg/util/protoutil" + "github.com/stretchr/testify/require" +) + +func TestReplicationManagerRequiresAdminRole(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + ctx := context.Background() + s, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{}) + defer s.Stopper().Stop(ctx) + tDB := sqlutils.MakeSQLRunner(sqlDB) + + var sessionData sessiondatapb.SessionData + { + var sessionSerialized []byte + tDB.QueryRow(t, "SELECT crdb_internal.serialize_session()").Scan(&sessionSerialized) + require.NoError(t, protoutil.Unmarshal(sessionSerialized, &sessionData)) + } + + getManagerForUser := func(u string) (streaming.ReplicationStreamManager, error) { + sqlUser, err := security.MakeSQLUsernameFromUserInput(u, security.UsernameValidation) + require.NoError(t, err) + execCfg := s.ExecutorConfig().(sql.ExecutorConfig) + txn := kvDB.NewTxn(ctx, "test") + p, cleanup := sql.NewInternalPlanner("test", txn, sqlUser, &sql.MemoryMetrics{}, &execCfg, sessionData) + defer cleanup() + ec := p.(interface{ EvalContext() *tree.EvalContext }).EvalContext() + return newReplicationStreamManagerWithPrivilegesCheck(ec) + } + + for _, tc := range []struct { + user string + expErr string + isEnterprise bool + }{ + {user: "admin", expErr: "", isEnterprise: true}, + {user: "root", expErr: "", isEnterprise: true}, + {user: "nobody", expErr: "replication restricted to ADMIN role", isEnterprise: true}, + {user: "admin", expErr: "use of REPLICATION requires an enterprise license", isEnterprise: false}, + {user: "root", expErr: "use of REPLICATION requires an enterprise license", isEnterprise: false}, + {user: "nobody", expErr: "replication restricted to ADMIN role", isEnterprise: false}, + } { + t.Run(fmt.Sprintf("%s/ent=%t", tc.user, tc.isEnterprise), func(t *testing.T) { + if tc.isEnterprise { + defer utilccl.TestingEnableEnterprise()() + } else { + defer utilccl.TestingDisableEnterprise()() + } + + m, err := getManagerForUser(tc.user) + if tc.expErr == "" { + require.NoError(t, err) + require.NotNil(t, m) + } else { + require.Regexp(t, tc.expErr, err) + require.Nil(t, m) + } + }) + } + +} diff --git a/pkg/sql/sem/builtins/replication_builtins.go b/pkg/sql/sem/builtins/replication_builtins.go index 176ed70f041e..1e4dd805dd16 100644 --- a/pkg/sql/sem/builtins/replication_builtins.go +++ b/pkg/sql/sem/builtins/replication_builtins.go @@ -43,7 +43,7 @@ var replicationBuiltins = map[string]builtinDefinition{ }, ReturnType: tree.FixedReturnType(types.Int), Fn: func(evalCtx *tree.EvalContext, args tree.Datums) (tree.Datum, error) { - mgr, err := streaming.GetReplicationStreamManager() + mgr, err := streaming.GetReplicationStreamManager(evalCtx) if err != nil { return nil, err } @@ -80,7 +80,7 @@ var replicationBuiltins = map[string]builtinDefinition{ }, ReturnType: tree.FixedReturnType(types.Int), Fn: func(evalCtx *tree.EvalContext, args tree.Datums) (tree.Datum, error) { - mgr, err := streaming.GetReplicationStreamManager() + mgr, err := streaming.GetReplicationStreamManager(evalCtx) if err != nil { return nil, err } @@ -114,7 +114,7 @@ var replicationBuiltins = map[string]builtinDefinition{ }, ReturnType: tree.FixedReturnType(types.String), Fn: func(evalCtx *tree.EvalContext, args tree.Datums) (tree.Datum, error) { - mgr, err := streaming.GetReplicationStreamManager() + mgr, err := streaming.GetReplicationStreamManager(evalCtx) if err != nil { return nil, err } @@ -150,13 +150,13 @@ var replicationBuiltins = map[string]builtinDefinition{ []*types.T{types.Bytes}, []string{"stream_event"}, ), - func(ctx *tree.EvalContext, args tree.Datums) (tree.ValueGenerator, error) { - mgr, err := streaming.GetReplicationStreamManager() + func(evalCtx *tree.EvalContext, args tree.Datums) (tree.ValueGenerator, error) { + mgr, err := streaming.GetReplicationStreamManager(evalCtx) if err != nil { return nil, err } return mgr.StreamPartition( - ctx, + evalCtx, streaming.StreamID(tree.MustBeDInt(args[0])), []byte(tree.MustBeDBytes(args[1])), ) diff --git a/pkg/streaming/api.go b/pkg/streaming/api.go index f193d32ad04c..19579d92a12a 100644 --- a/pkg/streaming/api.go +++ b/pkg/streaming/api.go @@ -27,8 +27,9 @@ func (j StreamID) SafeValue() {} // InvalidStreamID is the zero value for StreamID corresponding to no stream. const InvalidStreamID StreamID = 0 -// GetReplicationStreamManagerHook is the hook to get a collection of APIs that streaming replication supports. -var GetReplicationStreamManagerHook func() (ReplicationStreamManager, error) +// GetReplicationStreamManagerHook is the hook to get access to the replication API. +// Used by builtin functions to trigger streaming replication. +var GetReplicationStreamManagerHook func(evalCtx *tree.EvalContext) (ReplicationStreamManager, error) // ReplicationStreamManager represents a collection of APIs that streaming replication supports. type ReplicationStreamManager interface { @@ -65,9 +66,9 @@ type ReplicationStreamManager interface { } // GetReplicationStreamManager returns a ReplicationStreamManager if a CCL binary is loaded. -func GetReplicationStreamManager() (ReplicationStreamManager, error) { +func GetReplicationStreamManager(evalCtx *tree.EvalContext) (ReplicationStreamManager, error) { if GetReplicationStreamManagerHook == nil { return nil, errors.New("replication streaming requires a CCL binary") } - return GetReplicationStreamManagerHook() + return GetReplicationStreamManagerHook(evalCtx) }