Merge pull request #277 from dolthub/taylor/upload-restore

web,graphql: File upload fixes, add restore all tables button

graphql-server/schema.gql

@@ -346,6 +346,7 @@ type Mutation {
   resetDatabase(newDatabase: String): Boolean!
   loadDataFile(schemaName: String, tableName: String!, refName: String!, databaseName: String!, importOp: ImportOperation!, fileType: FileType!, file: Upload!, modifier: LoadDataModifier): Boolean!
   mergePull(fromBranchName: String!, toBranchName: String!, databaseName: String!, author: AuthorInfo): Boolean!
+  restoreAllTables(databaseName: String!, refName: String!): Boolean!
   createTag(tagName: String!, databaseName: String!, message: String, fromRefName: String!, author: AuthorInfo): String!
   deleteTag(databaseName: String!, tagName: String!): Boolean!
 }

graphql-server/src/queryFactory/dolt/index.ts

@@ -389,6 +389,51 @@ export class DoltQueryFactory
       args.refName,
     );
   }
+
+  async restoreAllTables(args: t.RefArgs): t.PR {
+    return this.queryQR(
+      async qr => {
+        console.log("[restore_all]: starting transaction");
+        await qr.query("BEGIN");
+
+        console.log("[restore_all]: calling");
+        const res = await qr.query(qh.callResetHard);
+
+        if (res.length && res[0].status !== "0") {
+          console.log("[restore_all]: reset not successful, rolling back");
+          await qr.query("ROLLBACK");
+          throw new Error("Reset --hard not successful");
+        }
+
+        // Handles any new tables that weren't restored by dolt_reset(--hard)
+        const status = await dem.getDoltStatus(qr.manager);
+
+        if (status.length) {
+          status.forEach(async r => {
+            console.log("[restore_all]: checking out new table", r.table_name);
+            const checkRes = await qr.query(qh.callCheckoutTable, [
+              r.table_name,
+            ]);
+            if (checkRes.length && checkRes[0].status !== "0") {
+              console.log(
+                "[restore_all]: checkout not successful, rolling back",
+              );
+              await qr.query("ROLLBACK");
+              throw new Error(
+                `Checking out table not successful: ${checkRes[0].message}`,
+              );
+            }
+          });
+        }
+
+        console.log("[restore_all]: committing");
+        await qr.query("COMMIT");
+        return res;
+      },
+      args.databaseName,
+      args.refName,
+    );
+  }
 }
 
 async function getTableInfoWithQR(

graphql-server/src/queryFactory/dolt/queries.ts

@@ -139,3 +139,7 @@ export function getOrderByFromDiffCols(cols: RawRows): string {
   const orderBy = diffCols.map(c => `\`${c}\` ASC`).join(", ");
   return orderBy === "" ? "" : `ORDER BY ${orderBy} `;
 }
+
+export const callResetHard = `CALL DOLT_RESET("--hard")`;
+
+export const callCheckoutTable = `CALL DOLT_CHECKOUT(?)`;

graphql-server/src/queryFactory/doltgres/index.ts

@@ -11,7 +11,7 @@ import { handleTableNotFound } from "../../utils";
 import * as dem from "../dolt/doltEntityManager";
 import { getAuthorString, handleRefNotFound, unionCols } from "../dolt/utils";
 import { PostgresQueryFactory } from "../postgres";
-import { getSchema } from "../postgres/utils";
+import { getSchema, tableWithoutSchema } from "../postgres/utils";
 import * as t from "../types";
 import * as qh from "./queries";
 
@@ -382,6 +382,49 @@ export class DoltgresQueryFactory
       args.refName,
     );
   }
+
+  async restoreAllTables(args: t.RefArgs): t.PR {
+    return this.queryQR(
+      async qr => {
+        console.log("[restore_all]: starting transaction");
+        await qr.query("BEGIN");
+
+        console.log("[restore_all]: calling");
+        const res = await qr.query(qh.callResetHard);
+        if (res.length && res[0].dolt_reset[0] !== "0") {
+          console.log("[restore_all]: reset not successful, rolling back");
+          await qr.query("ROLLBACK");
+          throw new Error("Reset --hard not successful");
+        }
+
+        // Handles any new tables that weren't restored by dolt_reset(--hard)
+        const status = await dem.getDoltStatus(qr.manager);
+        if (status.length) {
+          status.forEach(async r => {
+            console.log("[restore_all]: checking out new table", r.table_name);
+            const checkRes = await qr.query(qh.callCheckoutTable, [
+              tableWithoutSchema(r.table_name),
+            ]);
+            if (checkRes.length && checkRes[0].dolt_checkout[0] !== "0") {
+              console.log(
+                "[restore_all]: checkout not successful, rolling back",
+              );
+              await qr.query("ROLLBACK");
+              throw new Error(
+                `Checking out table not successful: ${checkRes[0].message}`,
+              );
+            }
+          });
+        }
+
+        console.log("[restore_all]: committing");
+        await qr.query("COMMIT");
+        return res;
+      },
+      args.databaseName,
+      args.refName,
+    );
+  }
 }
 
 async function getTableInfoWithQR(

graphql-server/src/queryFactory/doltgres/queries.ts

@@ -156,3 +156,7 @@ export function getOrderByFromDiffCols(cols: t.RawRows): string {
   const orderBy = diffCols.map(c => `"${c}" ASC`).join(", ");
   return orderBy === "" ? "" : `ORDER BY ${orderBy} `;
 }
+
+export const callResetHard = `SELECT DOLT_RESET('--hard')`;
+
+export const callCheckoutTable = `SELECT DOLT_CHECKOUT($1::text)`;

graphql-server/src/queryFactory/index.ts

@@ -142,4 +142,6 @@ export declare class QueryFactory {
   ): Promise<{ rows: t.RawRows; columns: t.RawRows }>;
 
   getRowDiffs(args: t.RowDiffArgs): t.DiffRes;
+
+  restoreAllTables(args: t.RefArgs): t.PR;
 }

graphql-server/src/queryFactory/mysql/index.ts

@@ -297,4 +297,8 @@ export class MySQLQueryFactory
   async getRowDiffs(_args: t.RowDiffArgs): t.DiffRes {
     throw notDoltError("get row sided diffs");
   }
+
+  async restoreAllTables(_args: t.RefArgs): t.PR {
+    throw notDoltError("restore all tables");
+  }
 }

graphql-server/src/status/status.resolver.ts

@@ -1,4 +1,4 @@
-import { Args, Query, Resolver } from "@nestjs/graphql";
+import { Args, Mutation, Query, Resolver } from "@nestjs/graphql";
 import { ConnectionProvider } from "../connections/connection.provider";
 import { RefArgs } from "../utils/commonTypes";
 import { Status, fromStatusRows } from "./status.model";
@@ -13,4 +13,11 @@ export class StatusResolver {
     const res = await conn.getStatus(args);
     return fromStatusRows(res, args.databaseName, args.refName);
   }
+
+  @Mutation(_returns => Boolean)
+  async restoreAllTables(@Args() args: RefArgs): Promise<boolean> {
+    const conn = this.conn.connection();
+    await conn.restoreAllTables(args);
+    return true;
+  }
 }

web/.eslintrc.js

@@ -13,6 +13,9 @@ module.exports = {
         extensions: [".js", ".jsx", ".ts", ".tsx"],
       },
     },
+    next: {
+      rootDir: "renderer",
+    },
   },
   env: {
     browser: true,

web/renderer/components/DatabaseTableHeader/CreateViewButton/CreateViewModal.tsx

@@ -3,7 +3,7 @@ import { useSqlEditorContext } from "@contexts/sqleditor";
 import { FormInput, FormModal, Loader } from "@dolthub/react-components";
 import useSqlBuilder from "@hooks/useSqlBuilder";
 import { ModalProps } from "@lib/modalProps";
-import { DatabaseParams } from "@lib/params";
+import { OptionalRefParams } from "@lib/params";
 import dynamic from "next/dynamic";
 import { SyntheticEvent, useState } from "react";
 import css from "./index.module.css";
@@ -13,7 +13,7 @@ const AceEditor = dynamic(async () => import("@components/AceEditor"), {
 });
 
 type Props = {
-  params: DatabaseParams & { refName?: string };
+  params: OptionalRefParams;
   query: string;
 } & ModalProps;
 

web/renderer/components/DatabaseTableHeader/CreateViewButton/index.tsx

@@ -1,12 +1,12 @@
 import HideForNoWritesWrapper from "@components/util/HideForNoWritesWrapper";
 import { Button } from "@dolthub/react-components";
-import { DatabaseParams } from "@lib/params";
+import { OptionalRefParams } from "@lib/params";
 import { useState } from "react";
 import CreateViewModal from "./CreateViewModal";
 import css from "./index.module.css";
 
 type Props = {
-  params: DatabaseParams & { refName?: string };
+  params: OptionalRefParams;
   query: string;
 };
 

web/renderer/components/FormSelectForRefs/useGetBranchOptionsForSelect.ts

@@ -1,13 +1,17 @@
 import { FormSelectTypes } from "@dolthub/react-components";
 import { excerpt } from "@dolthub/web-utils";
-import { useBranchesForSelectorQuery } from "@gen/graphql-types";
+import {
+  BranchForBranchSelectorFragment,
+  useBranchesForSelectorQuery,
+} from "@gen/graphql-types";
 import { ApolloErrorType } from "@lib/errors/types";
 import { DatabaseParams } from "@lib/params";
 
 type ReturnType = {
   branchOptions: Array<FormSelectTypes.Option<string>>;
   error?: ApolloErrorType;
   loading: boolean;
+  defaultBranch?: string;
 };
 
 export default function useGetBranchOptionsForSelect(
@@ -26,5 +30,19 @@ export default function useGetBranchOptionsForSelect(
       };
     }) ?? [];
 
-  return { branchOptions, error: branchRes.error, loading: branchRes.loading };
+  return {
+    branchOptions,
+    error: branchRes.error,
+    loading: branchRes.loading,
+    defaultBranch: getDefaultBranch(branchRes.data?.allBranches),
+  };
+}
+
+function getDefaultBranch(
+  branches?: BranchForBranchSelectorFragment[],
+): string | undefined {
+  if (!branches?.length) return undefined;
+  const mainBranch = branches.find(b => b.branchName === "main");
+  if (mainBranch) return mainBranch.branchName;
+  return branches[0].branchName;
 }

web/renderer/components/StatusWithOptions/ResetModal.tsx

@@ -1,10 +1,12 @@
 import Link from "@components/links/Link";
-import { Button, Modal } from "@dolthub/react-components";
-import { StatusFragment } from "@gen/graphql-types";
+import { Button, Loader, Modal } from "@dolthub/react-components";
+import { StatusFragment, useRestoreAllMutation } from "@gen/graphql-types";
+import useMutation from "@hooks/useMutation";
 import useSqlBuilder from "@hooks/useSqlBuilder";
 import { ModalProps } from "@lib/modalProps";
 import { RefParams } from "@lib/params";
 import { getPostgresTableName } from "@lib/postgres";
+import { refetchSqlUpdateQueriesCacheEvict } from "@lib/refetchQueries";
 import { sqlQuery } from "@lib/urls";
 import css from "./index.module.css";
 
@@ -16,6 +18,21 @@ type Props = ModalProps & {
 
 export default function ResetModal(props: Props) {
   const { getCallProcedure, isPostgres } = useSqlBuilder();
+  const { mutateFn, loading, err, client } = useMutation({
+    hook: useRestoreAllMutation,
+  });
+
+  const onRestoreAll = async () => {
+    try {
+      await mutateFn({ variables: props.params });
+      props.setIsOpen(false);
+      client
+        .refetchQueries(refetchSqlUpdateQueriesCacheEvict)
+        .catch(console.error);
+    } catch (_) {
+      // Handled by useMutation
+    }
+  };
 
   const getTableName = (tn: string): string => {
     if (isPostgres) {
@@ -33,9 +50,15 @@ export default function ResetModal(props: Props) {
       title="Reset uncommitted changes"
       isOpen={props.isOpen}
       onRequestClose={onClose}
-      button={<Button onClick={onClose}>Done</Button>}
+      button={
+        <Button onClick={onRestoreAll} pill>
+          Restore all tables
+        </Button>
+      }
+      err={err}
     >
       <div>
+        <Loader loaded={!loading} />
         <p>
           Choose to unstage staged tables or restore tables to their current
           contents in the current <code>HEAD</code>.

web/renderer/components/StatusWithOptions/queries.ts

@@ -14,3 +14,9 @@ export const GET_STATUS = gql`
     }
   }
 `;
+
+export const RESTORE_ALL = gql`
+  mutation RestoreAll($databaseName: String!, $refName: String!) {
+    restoreAllTables(databaseName: $databaseName, refName: $refName)
+  }
+`;

web/renderer/components/pageComponents/FileUploadPage/FileInfo/index.tsx

@@ -1,4 +1,5 @@
 import { Button } from "@dolthub/react-components";
+import { excerpt } from "@dolthub/web-utils";
 import { FiFile } from "@react-icons/all-files/fi/FiFile";
 import cx from "classnames";
 import { ReactNode } from "react";
@@ -28,7 +29,9 @@ export default function FileInfo(props: Props) {
       >
         <span className={css.fileInfo}>
           <FiFile className={css.fileIcon} />
-          <span data-cy="file-name">{state.selectedFile.name}</span>
+          <span data-cy="file-name">
+            {excerpt(state.selectedFile.name, 30)}
+          </span>
           <span className={css.fileSize}>
             {fileSize(state.selectedFile.size)}
           </span>

web/renderer/components/pageComponents/FileUploadPage/Steps/Branch/index.tsx

@@ -1,6 +1,7 @@
 import useGetBranchOptionsForSelect from "@components/FormSelectForRefs/useGetBranchOptionsForSelect";
 import DatabaseLink from "@components/links/DatabaseLink";
 import { ErrorMsg, FormSelect, Loader } from "@dolthub/react-components";
+import { useEffect } from "react";
 import StepLayout from "../../StepLayout";
 import { useFileUploadContext } from "../../contexts/fileUploadLocalForage";
 import { UploadStage } from "../../enums";
@@ -16,7 +17,14 @@ export default function Branch() {
 function Inner() {
   const { state, error, updateLoad, setItem, dbParams, getUploadUrl } =
     useFileUploadContext();
-  const { branchOptions } = useGetBranchOptionsForSelect(dbParams);
+  const { branchOptions, defaultBranch } =
+    useGetBranchOptionsForSelect(dbParams);
+
+  useEffect(() => {
+    if (defaultBranch && !state.branchName) {
+      setItem("branchName", defaultBranch);
+    }
+  }, [defaultBranch, state.branchName, setItem]);
 
   return (
     <StepLayout

web/renderer/components/pageComponents/FileUploadPage/Steps/Table/index.tsx

@@ -71,7 +71,10 @@ function Inner() {
                 label={isPostgres ? "Table" : undefined}
                 light
               />
-              <UploadQueryInfo tableName={state.existingTable} />
+              <UploadQueryInfo
+                tableName={state.existingTable}
+                hideModifierOptions
+              />
             </div>
           </TableOption>
         </div>

web/renderer/components/pageComponents/FileUploadPage/Steps/Upload/index.tsx

@@ -51,7 +51,10 @@ function Inner() {
           <OpenSpreadsheetZone />
         </div>
         <EditableTable />
-        <UploadQueryInfo tableName={state.tableName} />
+        <UploadQueryInfo
+          tableName={state.tableName}
+          hideModifierOptions={uState.spreadsheetOverlayOpen}
+        />
         <ErrorMsg err={uState.error ?? error} />
       </div>
     </StepLayout>

web/renderer/components/pageComponents/FileUploadPage/UploadQueryInfo.tsx

@@ -6,6 +6,7 @@ import css from "./index.module.css";
 
 type Props = {
   forSpreadsheet?: boolean;
+  hideModifierOptions?: boolean;
   tableName: string;
 };
 
@@ -42,7 +43,7 @@ export default function UploadQueryInfo(props: Props) {
           </div>
         </HelpPopup>
       </div>
-      {!isPostgres && <ModifierOptions />}
+      {!isPostgres && !props.hideModifierOptions && <ModifierOptions />}
     </div>
   );
 }