feat: datadoc restore feature

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):

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", {}),

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

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

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<void>;
+    onCompare: (version?: ICommit) => void;
 }
 
 export const CommitCard: React.FC<IProps> = ({
@@ -41,9 +41,7 @@ export const CommitCard: React.FC<IProps> = ({
                 </Link>
             </Button>
             <AsyncButton
-                onClick={async () =>
-                    onRestore(version.sha, version.commit.message)
-                }
+                onClick={() => onRestore(version)}
                 className="ml8"
                 title="Restore Version"
                 hoverColor="var(--color-accent-dark)"

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<IProps> = ({
         }
     );
 
+    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<IProps> = ({
                     key={version.sha}
                     version={version}
                     onRestore={handleRestore}
-                    onCompare={() => toggleCompare(version)}
+                    onCompare={toggleCompare}
                 />
             ))}
         </div>
@@ -197,8 +202,8 @@ export const GitHubVersions: React.FunctionComponent<IProps> = ({
                         />
                     ) : (
                         <QueryComparison
-                            fromQuery={comparisonData?.current_content}
-                            toQuery={comparisonData?.commit_content}
+                            fromQuery={comparisonData?.commit_content}
+                            toQuery={comparisonData?.current_content}
                             fromQueryTitle={`Commit: ${selectedCommit.commit.message}`}
                             toQueryTitle="Current DataDoc"
                         />

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 {

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<void> => {
+            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]
+    );
+}

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<IDataDocSocketPromise<[rawDataDoc: any]>>(
+            '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',

querybook/webapp/redux/dataDoc/action.ts

@@ -273,6 +273,17 @@ export function deleteDataDoc(docId: number): ThunkResult<Promise<void>> {
     };
 }
 
+export function restoreDataDoc(
+    docId: number,
+    commitId: string,
+    commitMessage: string
+): ThunkResult<Promise<void>> {
+    return async (dispatch) => {
+        await dataDocSocket.restoreDataDoc(docId, commitId, commitMessage);
+        await dispatch(fetchDataDoc(docId));
+    };
+}
+
 export function insertDataDocCell(
     docId: number,
     index: number,

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<Promise<any>> {
                     }
                 },
             },
+
+            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) {