From 524d8d52e76cfc1dce07eaf34e8d7d0d471775cd Mon Sep 17 00:00:00 2001 From: rongzhang Date: Mon, 9 Dec 2024 18:20:38 +0000 Subject: [PATCH] feat: peer review backend endpoints --- querybook/server/datasources/__init__.py | 3 +- .../server/datasources/query_execution.py | 150 +++++++++++++++--- .../DataDocQueryCell/DataDocQueryCell.tsx | 57 +++++-- .../QueryPeerReviewModal.tsx | 10 +- querybook/webapp/const/analytics.ts | 1 + .../webapp/redux/queryExecutions/action.ts | 8 +- querybook/webapp/resource/queryExecution.ts | 9 +- 7 files changed, 193 insertions(+), 45 deletions(-) diff --git a/querybook/server/datasources/__init__.py b/querybook/server/datasources/__init__.py index bad09b4ef..e41b35317 100644 --- a/querybook/server/datasources/__init__.py +++ b/querybook/server/datasources/__init__.py @@ -19,7 +19,7 @@ from . import survey from . import query_transform from . import github - +from . import query_review # Keep this at the end of imports to make sure the plugin APIs override the default ones try: @@ -50,3 +50,4 @@ query_transform api_plugin github +query_review diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index 97ab2cd8d..b5fd97d32 100644 --- a/querybook/server/datasources/query_execution.py +++ b/querybook/server/datasources/query_execution.py @@ -60,41 +60,141 @@ QUERY_RESULT_LIMIT_CONFIG = get_config_value("query_result_limit") +def process_query_execution(query_execution_id, metadata, data_cell_id, session): + """Process execution metadata and associate query execution with DataDoc.""" + metadata = metadata or {} + + used_api_token = request.headers.get("api-access-token") is not None + if used_api_token: + metadata["used_api_token"] = used_api_token + + if metadata: + logic.create_query_execution_metadata( + query_execution_id, metadata, session=session + ) + + data_doc = None + if data_cell_id: + datadoc_logic.append_query_executions_to_data_cell( + data_cell_id, [query_execution_id], session=session + ) + data_cell = datadoc_logic.get_data_cell_by_id(data_cell_id, session=session) + data_doc = data_cell.doc if data_cell else None + return data_doc + + +def initiate_query_peer_review_workflow( + self, + query_execution_id: int, + uid: int, + peer_review_params: dict, + session=None, +): + """ + Initiates the query peer review workflow by creating a QueryReview, + assigning reviewers, and sending notifications. + """ + reviewer_ids = peer_review_params.get("reviewer_ids", []) + external_recipients = peer_review_params.get("external_recipients", []) + notifier_name = peer_review_params.get("notifier_name", "") + review_request_reason = peer_review_params.get("review_request_reason", "") + + query_review = logic.create_query_review( + query_author_id=uid, + query_execution_id=query_execution_id, + review_request_reason=review_request_reason, + commit=False, + session=session, + ) + + # Add reviewers to the query_review.reviewers relationship + for reviewer_id in reviewer_ids: + reviewer = user_logic.get_user_by_id(reviewer_id, session=session) + if reviewer: + query_review.reviewers.append(reviewer) + + session.commit() + + +def initiate_query_execution( + query_execution, uid, delay_execution, peer_review_params, session +): + """Initiate the query execution based on delay_execution flag.""" + # Initiate peer review workflow + if delay_execution: + initiate_query_peer_review_workflow( + query_execution_id=query_execution.id, + uid=uid, + peer_review_params=peer_review_params, + session=session, + ) + # Start immediate execution + else: + run_query_task.apply_async( + args=[ + query_execution.id, + ] + ) + + @register("/query_execution/", methods=["POST"]) def create_query_execution( - query, engine_id, metadata=None, data_cell_id=None, originator=None + query, + engine_id, + metadata=None, + data_cell_id=None, + originator=None, + peer_review_params=None, + delay_execution=False, ): + """ + Creates a new QueryExecution. + + Args: + query (str): The SQL query to execute. + engine_id (int): The ID of the query engine to use. + metadata (dict, optional): Additional metadata for the query execution. + data_cell_id (int, optional): ID of the DataDoc cell to associate with. + originator (str, optional): Identifier for the originator of the request. + peer_review_params (dict, optional): Parameters for peer review workflow. + delay_execution (bool, optional): Flag to delay query execution. + + Returns: + dict: A dictionary representation of the created QueryExecution. + """ with DBSession() as session: verify_query_engine_permission(engine_id, session=session) uid = current_user.id + status = ( + QueryExecutionStatus.PENDING_REVIEW + if delay_execution + else QueryExecutionStatus.INITIALIZED + ) query_execution = logic.create_query_execution( - query=query, engine_id=engine_id, uid=uid, session=session + query=query, + engine_id=engine_id, + uid=uid, + status=status, + session=session, ) - metadata = metadata or {} - used_api_token = request.headers.get("api-access-token") is not None - if used_api_token: - metadata["used_api_token"] = used_api_token - if metadata: - logic.create_query_execution_metadata( - query_execution.id, metadata, session=session - ) - - data_doc = None - if data_cell_id: - datadoc_logic.append_query_executions_to_data_cell( - data_cell_id, [query_execution.id], session=session - ) - data_cell = datadoc_logic.get_data_cell_by_id(data_cell_id, session=session) - data_doc = data_cell.doc + data_doc = process_query_execution( + query_execution_id=query_execution.id, + metadata=metadata, + data_cell_id=data_cell_id, + session=session, + ) try: - run_query_task.apply_async( - args=[ - query_execution.id, - ] + initiate_query_execution( + query_execution=query_execution, + uid=uid, + delay_execution=delay_execution, + peer_review_params=peer_review_params, + session=session, ) + query_execution_dict = query_execution.to_dict() if data_doc: @@ -311,9 +411,9 @@ def download_statement_execution_result(statement_execution_id): raw = reader.read_raw() response = Response(raw) response.headers["Content-Type"] = "text/csv" - response.headers[ - "Content-Disposition" - ] = f'attachment; filename="{download_file_name}"' + response.headers["Content-Disposition"] = ( + f'attachment; filename="{download_file_name}"' + ) return response diff --git a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx index 260853745..77beef4ab 100644 --- a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx +++ b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx @@ -531,15 +531,42 @@ class DataDocQueryCellComponent extends React.PureComponent { reviewerIds: number[], externalRecipients: string[], notifierName: string, - justification: string, - query: string, - queryEngine: IQueryEngine + justification: string ) { - try { - console.error('TODO: Implement API call here'); - } catch (error) { - console.error('Failed to request peer review:', error); + trackClick({ + component: ComponentType.DATADOC_QUERY_CELL, + element: ElementType.PEER_REVIEW_QUERY_BUTTON, + aux: { + lintError: this.state.hasLintError, + sampleRate: this.sampleRate, + }, + }); + + const transformedQuery = await this.getTransformedQuery(); + if (!transformedQuery) { + return null; } + + const executionMetadata = + this.sampleRate > 0 ? { sample_rate: this.sampleRate } : null; + + const peerReviewParams = { + reviewer_ids: reviewerIds, + external_recipients: externalRecipients, + notifier_name: notifierName, + review_request_reason: justification, + }; + + const delayExecution = true; + + await this.props.createQueryExecution( + transformedQuery, + this.engineId, + this.props.cellId, + executionMetadata, + peerReviewParams, + delayExecution + ); } @bind @@ -1129,8 +1156,20 @@ function mapDispatchToProps(dispatch: Dispatch) { query: string, engineId: number, cellId: number, - metadata: Record - ) => dispatch(createQueryExecution(query, engineId, cellId, metadata)), + metadata: Record, + peerReviewParams?: Record, + delayExecution?: boolean + ) => + dispatch( + createQueryExecution( + query, + engineId, + cellId, + metadata, + peerReviewParams, + delayExecution + ) + ), setTableSidebarId: (id: number) => dispatch(setSidebarTableId(id)), diff --git a/querybook/webapp/components/QueryPeerReviewModal/QueryPeerReviewModal.tsx b/querybook/webapp/components/QueryPeerReviewModal/QueryPeerReviewModal.tsx index edd824b32..0f0a33e62 100644 --- a/querybook/webapp/components/QueryPeerReviewModal/QueryPeerReviewModal.tsx +++ b/querybook/webapp/components/QueryPeerReviewModal/QueryPeerReviewModal.tsx @@ -26,9 +26,7 @@ interface IQueryPeerReviewFormProps { reviewerIds: number[], externalRecipients: string[], notifierName: string, - justification: string, - query: string, - queryEngine: IQueryEngine + justification: string ) => Promise; onHide: () => void; } @@ -87,9 +85,7 @@ const QueryPeerReviewForm: React.FC = ({ reviewerIds, externalRecipients, notifierName, - values.justification, - query, - queryEngine + values.justification ); onHide(); toast.success( @@ -99,7 +95,7 @@ const QueryPeerReviewForm: React.FC = ({ toast.error('Failed to request review.'); } }, - [onConfirm, onHide, query, queryEngine] + [onConfirm, onHide] ); const { description: featureDescription, user_guide_link: helpLink } = diff --git a/querybook/webapp/const/analytics.ts b/querybook/webapp/const/analytics.ts index 3ea663fbf..6d5bc1a95 100644 --- a/querybook/webapp/const/analytics.ts +++ b/querybook/webapp/const/analytics.ts @@ -90,6 +90,7 @@ export enum ElementType { CREATE_DATADOC_BUTTON = 'CREATE_DATADOC_BUTTON', RESULT_EXPORT_BUTTON = 'RESULT_EXPORT_BUTTON', INSERT_SNIPPET_BUTTON = 'INSERT_SNIPPET_BUTTON', + PEER_REVIEW_QUERY_BUTTON = 'PEER_REVIEW_QUERY_BUTTON', // Chart Cell CHART_CONFIG_BUTTON = 'CHART_CONFIG_BUTTON', diff --git a/querybook/webapp/redux/queryExecutions/action.ts b/querybook/webapp/redux/queryExecutions/action.ts index b548f4a41..1ea68e1dd 100644 --- a/querybook/webapp/redux/queryExecutions/action.ts +++ b/querybook/webapp/redux/queryExecutions/action.ts @@ -379,7 +379,9 @@ export function createQueryExecution( query: string, engineId?: number, cellId?: number, - metadata?: Record + metadata?: Record, + peerReviewParams?: Record, + delayExecution?: boolean ): ThunkResult> { return async (dispatch, getState) => { const state = getState(); @@ -389,7 +391,9 @@ export function createQueryExecution( query, selectedEngineId, cellId, - metadata + metadata, + peerReviewParams, + delayExecution ); dispatch(receiveQueryExecution(queryExecution, cellId)); diff --git a/querybook/webapp/resource/queryExecution.ts b/querybook/webapp/resource/queryExecution.ts index 400e123f9..2c4913dc1 100644 --- a/querybook/webapp/resource/queryExecution.ts +++ b/querybook/webapp/resource/queryExecution.ts @@ -71,7 +71,9 @@ export const QueryExecutionResource = { query: string, engineId: number, cellId?: number, - metadata?: Record + metadata?: Record, + peerReviewParams?: Record, + delayExecution?: boolean ) => { const params = { query, @@ -87,6 +89,11 @@ export const QueryExecutionResource = { params['originator'] = dataDocSocket.socketId; } + if (delayExecution != null) { + params['peer_review_params'] = peerReviewParams; + params['delay_execution'] = delayExecution; + } + return ds.save('/query_execution/', params); },