From de5b3cd39cf27f1949cecb5afe94e69588ff354e Mon Sep 17 00:00:00 2001 From: rongzhang Date: Wed, 6 Nov 2024 21:26:48 +0000 Subject: [PATCH] feat: datadoc restore feature --- .../server/datasources_socketio/datadoc.py | 11 ++++ querybook/server/lib/github/serializers.py | 2 +- querybook/server/logic/datadoc.py | 53 +++++++++++++++++++ querybook/server/logic/datadoc_collab.py | 39 ++++++++++++++ .../components/DataDocGitHub/CommitCard.tsx | 8 ++- .../DataDocGitHub/GitHubVersions.tsx | 17 +++--- querybook/webapp/const/analytics.ts | 1 + .../webapp/hooks/dataDoc/useRestoreDataDoc.ts | 52 ++++++++++++++++++ .../webapp/lib/data-doc/datadoc-socketio.ts | 33 ++++++++++++ querybook/webapp/redux/dataDoc/action.ts | 11 ++++ .../dataDocWebsocket/dataDocWebsocket.ts | 21 ++++++++ 11 files changed, 236 insertions(+), 12 deletions(-) create mode 100644 querybook/webapp/hooks/dataDoc/useRestoreDataDoc.ts diff --git a/querybook/server/datasources_socketio/datadoc.py b/querybook/server/datasources_socketio/datadoc.py index 9397162bd..d7455cd9a 100644 --- a/querybook/server/datasources_socketio/datadoc.py +++ b/querybook/server/datasources_socketio/datadoc.py @@ -216,6 +216,17 @@ def update_data_doc(id, fields): datadoc_collab.update_datadoc(id, fields, sid=request.sid) +@register_socket("restore_data_doc", namespace=DATA_DOC_NAMESPACE) +@data_doc_socket +def restore_data_doc(datadoc_id: int, commit_sha: str, commit_message: str): + datadoc_collab.restore_data_doc( + datadoc_id=datadoc_id, + commit_sha=commit_sha, + commit_message=commit_message, + sid=request.sid, + ) + + @register_socket("update_data_cell", namespace=DATA_DOC_NAMESPACE) @data_doc_socket def update_data_cell(doc_id, cell_id, fields): diff --git a/querybook/server/lib/github/serializers.py b/querybook/server/lib/github/serializers.py index 73b146264..1341350aa 100644 --- a/querybook/server/lib/github/serializers.py +++ b/querybook/server/lib/github/serializers.py @@ -223,7 +223,7 @@ def deserialize_datadoc_content(content_str: str) -> List[DataCell]: cell = DataCell( id=metadata.get("id"), cell_type=cell_type_enum, - context=context if cell_type_enum != DataCellType.chart else None, + context=context if cell_type_enum != DataCellType.chart else "", created_at=parse_datetime_as_utc(metadata.get("created_at")), updated_at=parse_datetime_as_utc(metadata.get("updated_at")), meta=metadata.get("meta", {}), diff --git a/querybook/server/logic/datadoc.py b/querybook/server/logic/datadoc.py index 37ff2ddbf..ec7c15c03 100644 --- a/querybook/server/logic/datadoc.py +++ b/querybook/server/logic/datadoc.py @@ -248,6 +248,59 @@ def clone_data_doc(id, owner_uid, commit=True, session=None): return new_data_doc +@with_session +def restore_data_doc_from_commit( + datadoc_id: int, commit_datadoc: DataDoc, commit=True, session=None +) -> DataDoc: + data_doc = get_data_doc_by_id(datadoc_id, session=session) + assert data_doc is not None, "DataDoc not found" + + # Update the DataDoc's title and meta + data_doc = update_data_doc( + datadoc_id, + title=commit_datadoc.title, + meta=commit_datadoc.meta, + commit=False, + session=session, + ) + + # Delete existing DataDocCells and DataCells + for existing_cell in data_doc.cells: + delete_data_doc_cell( + data_doc_id=data_doc.id, + data_cell_id=existing_cell.id, + commit=False, + session=session, + ) + + # Create new DataCells from commit and add them to the DataDoc + for index, cell in enumerate(commit_datadoc.cells): + data_cell = create_data_cell( + cell_type=cell.cell_type.name, + context=cell.context, + meta=cell.meta, + commit=False, + session=session, + ) + insert_data_doc_cell( + data_doc_id=data_doc.id, + cell_id=data_cell.id, + index=index, + commit=False, + session=session, + ) + + if commit: + session.commit() + update_es_data_doc_by_id(data_doc.id) + update_es_queries_by_datadoc_id(data_doc.id) + else: + session.flush() + + session.refresh(data_doc) + return data_doc + + """ ---------------------------------------------------------------------------------------------------------- DATA CELL diff --git a/querybook/server/logic/datadoc_collab.py b/querybook/server/logic/datadoc_collab.py index 8e1154548..d384a40ed 100644 --- a/querybook/server/logic/datadoc_collab.py +++ b/querybook/server/logic/datadoc_collab.py @@ -4,9 +4,13 @@ ) from app.flask_app import socketio from app.db import with_session +from clients.github_client import GitHubClient from const.data_doc import DATA_DOC_NAMESPACE +from datasources.github import with_github_client from logic import datadoc as logic +from logic import user as user_logic from logic.datadoc_permission import assert_can_read, assert_can_write +from flask_login import current_user @with_session @@ -43,6 +47,41 @@ def update_datadoc(doc_id, fields, sid="", session=None): return doc_dict +@with_session +@with_github_client +def restore_data_doc( + github_client: GitHubClient, + datadoc_id: int, + commit_sha: str, + commit_message: str, + sid="", + session=None, +): + assert_can_write(datadoc_id, session=session) + verify_data_doc_permission(datadoc_id, session=session) + + commit_datadoc = github_client.get_datadoc_at_commit(commit_sha) + restored_datadoc = logic.restore_data_doc_from_commit( + datadoc_id, commit_datadoc, commit=True, session=session + ) + + user = user_logic.get_user_by_id(current_user.id, session=session) + assert user is not None, "User does not exist" + + # Emit the restored DataDoc to clients + socketio.emit( + "data_doc_restored", + ( + sid, + restored_datadoc.to_dict(with_cells=True), + commit_message, + user.get_name(), + ), + namespace=DATA_DOC_NAMESPACE, + room=datadoc_id, + ) + + @with_session def insert_data_cell( doc_id, index, cell_type, context=None, meta=None, sid="", session=None diff --git a/querybook/webapp/components/DataDocGitHub/CommitCard.tsx b/querybook/webapp/components/DataDocGitHub/CommitCard.tsx index 097853e54..fe0555d97 100644 --- a/querybook/webapp/components/DataDocGitHub/CommitCard.tsx +++ b/querybook/webapp/components/DataDocGitHub/CommitCard.tsx @@ -11,8 +11,8 @@ import './GitHub.scss'; interface IProps { version: ICommit; - onRestore: (sha: string, message: string) => void; - onCompare: (version: ICommit) => void; + onRestore: (version: ICommit) => Promise; + onCompare: (version?: ICommit) => void; } export const CommitCard: React.FC = ({ @@ -41,9 +41,7 @@ export const CommitCard: React.FC = ({ - onRestore(version.sha, version.commit.message) - } + onClick={() => onRestore(version)} className="ml8" title="Restore Version" hoverColor="var(--color-accent-dark)" diff --git a/querybook/webapp/components/DataDocGitHub/GitHubVersions.tsx b/querybook/webapp/components/DataDocGitHub/GitHubVersions.tsx index dec48794f..48f0d1262 100644 --- a/querybook/webapp/components/DataDocGitHub/GitHubVersions.tsx +++ b/querybook/webapp/components/DataDocGitHub/GitHubVersions.tsx @@ -1,6 +1,7 @@ import React, { useCallback, useState } from 'react'; import { QueryComparison } from 'components/TranspileQueryModal/QueryComparison'; +import { useRestoreDataDoc } from 'hooks/dataDoc/useRestoreDataDoc'; import { usePaginatedResource } from 'hooks/usePaginatedResource'; import { useResource } from 'hooks/useResource'; import { GitHubResource, ICommit } from 'resource/github'; @@ -62,11 +63,15 @@ export const GitHubVersions: React.FunctionComponent = ({ } ); + const restoreDataDoc = useRestoreDataDoc(); + const handleRestore = useCallback( - async (commitSha: string, commitMessage: string) => { - alert('Restore feature not implemented yet'); + async (commit: ICommit) => { + const commitId = commit.sha; + const commitMessage = commit.commit.message; + await restoreDataDoc(docId, commitId, commitMessage); }, - [] + [docId, restoreDataDoc] ); const toggleCompare = useCallback( @@ -135,7 +140,7 @@ export const GitHubVersions: React.FunctionComponent = ({ key={version.sha} version={version} onRestore={handleRestore} - onCompare={() => toggleCompare(version)} + onCompare={toggleCompare} /> ))} @@ -197,8 +202,8 @@ export const GitHubVersions: React.FunctionComponent = ({ /> ) : ( diff --git a/querybook/webapp/const/analytics.ts b/querybook/webapp/const/analytics.ts index 509679ba1..294dc3ada 100644 --- a/querybook/webapp/const/analytics.ts +++ b/querybook/webapp/const/analytics.ts @@ -120,6 +120,7 @@ export enum ElementType { // Github Integration GITHUB_CONNECT_BUTTON = 'GITHUB_CONNECT_BUTTON', GITHUB_LINK_BUTTON = 'GITHUB_LINK_BUTTON', + GITHUB_RESTORE_DATADOC_BUTTON = 'GITHUB_RESTORE_DATADOC_BUTTON', } export interface EventData { diff --git a/querybook/webapp/hooks/dataDoc/useRestoreDataDoc.ts b/querybook/webapp/hooks/dataDoc/useRestoreDataDoc.ts new file mode 100644 index 000000000..a8136e787 --- /dev/null +++ b/querybook/webapp/hooks/dataDoc/useRestoreDataDoc.ts @@ -0,0 +1,52 @@ +import { useCallback } from 'react'; +import toast from 'react-hot-toast'; +import { useDispatch } from 'react-redux'; + +import { ComponentType, ElementType } from 'const/analytics'; +import { trackClick } from 'lib/analytics'; +import { sendConfirm } from 'lib/querybookUI'; +import { restoreDataDoc } from 'redux/dataDoc/action'; +import { Dispatch } from 'redux/store/types'; + +export function useRestoreDataDoc() { + const dispatch: Dispatch = useDispatch(); + + const handleConfirm = useCallback( + (docId: number, commitId: string, commitMessage: string) => () => { + trackClick({ + component: ComponentType.GITHUB, + element: ElementType.GITHUB_RESTORE_DATADOC_BUTTON, + }); + + toast.promise( + dispatch(restoreDataDoc(docId, commitId, commitMessage)), + { + loading: 'Restoring DataDoc...', + success: 'DataDoc has been successfully restored!', + error: 'Failed to restore DataDoc. Please try again.', + } + ); + }, + [dispatch] + ); + + return useCallback( + async ( + docId: number, + commitId: string, + commitMessage: string + ): Promise => { + sendConfirm({ + header: 'Restore DataDoc?', + message: + 'You are about to restore this DataDoc to the selected commit. Restoring will overwrite your current work. Please ensure you have committed any ongoing changes before proceeding.', + onConfirm: handleConfirm(docId, commitId, commitMessage), + confirmColor: 'cancel', + cancelColor: 'default', + confirmText: 'Confirm Restore', + confirmIcon: 'AlertOctagon', + }); + }, + [handleConfirm] + ); +} diff --git a/querybook/webapp/lib/data-doc/datadoc-socketio.ts b/querybook/webapp/lib/data-doc/datadoc-socketio.ts index 517f8688c..b85b0befc 100644 --- a/querybook/webapp/lib/data-doc/datadoc-socketio.ts +++ b/querybook/webapp/lib/data-doc/datadoc-socketio.ts @@ -26,6 +26,15 @@ export interface IDataDocSocketEvent { (rawDataDoc, isSameOrigin: boolean) => any >; + dataDocRestored?: IDataDocSocketEventPromise< + ( + rawDataDoc: any, + commitMessage: string, + username: string, + isSameOrigin: boolean + ) => any + >; + updateDataCell?: IDataDocSocketEventPromise< (rawDataCell, isSameOrigin: boolean) => any >; @@ -161,6 +170,17 @@ export class DataDocSocket { ); }; + public restoreDataDoc = ( + docId: number, + commitId: string, + commitMessage: string + ) => { + this.socket.emit('restore_data_doc', docId, commitId, commitMessage); + return this.makePromise>( + 'dataDocRestored' + ); + }; + public updateDataCell = ( cellId: number, fields: { meta?: IDataCellMeta; context?: string } @@ -320,6 +340,19 @@ export class DataDocSocket { ); }); + this.socket.on( + 'data_doc_restored', + (originator, rawDataDoc, commitMessage, username) => { + this.resolvePromiseAndEvent( + 'dataDocRestored', + originator, + rawDataDoc, + commitMessage, + username + ); + } + ); + this.socket.on('data_cell_updated', (originator, rawDataCell) => { this.resolvePromiseAndEvent( 'updateDataCell', diff --git a/querybook/webapp/redux/dataDoc/action.ts b/querybook/webapp/redux/dataDoc/action.ts index 7886bb1c9..a3d8c48d9 100644 --- a/querybook/webapp/redux/dataDoc/action.ts +++ b/querybook/webapp/redux/dataDoc/action.ts @@ -273,6 +273,17 @@ export function deleteDataDoc(docId: number): ThunkResult> { }; } +export function restoreDataDoc( + docId: number, + commitId: string, + commitMessage: string +): ThunkResult> { + return async (dispatch) => { + await dataDocSocket.restoreDataDoc(docId, commitId, commitMessage); + await dispatch(fetchDataDoc(docId)); + }; +} + export function insertDataDocCell( docId: number, index: number, diff --git a/querybook/webapp/redux/dataDocWebsocket/dataDocWebsocket.ts b/querybook/webapp/redux/dataDocWebsocket/dataDocWebsocket.ts index e22c73831..407fccc7b 100644 --- a/querybook/webapp/redux/dataDocWebsocket/dataDocWebsocket.ts +++ b/querybook/webapp/redux/dataDocWebsocket/dataDocWebsocket.ts @@ -1,3 +1,5 @@ +import toast from 'react-hot-toast'; + import { IAccessRequest } from 'const/accessRequest'; import { IDataDocEditor } from 'const/datadoc'; import dataDocSocket, { @@ -61,6 +63,25 @@ export function openDataDoc(docId: number): ThunkResult> { } }, }, + + dataDocRestored: { + resolve: ( + rawDataDoc, + commitMessage, + username, + isSameOrigin + ) => { + dispatch(fetchDataDoc(docId)); + + // Show a notification to other users + if (!isSameOrigin) { + toast.success( + `DataDoc restored by ${username}: "${commitMessage}"` + ); + } + }, + }, + updateDataCell: { resolve: (rawDataCell, isSameOrigin) => { if (!isSameOrigin) {