fix: optimistic update (#1052)

* fix: optimistic update

* added guard and changed useUpdateAtom to useSetAtom

* removed useUpdateAtom

* extract updateTableSchema init

* remove unnecessary deps

Co-authored-by: Han Tuerker <burhan.tuerker@gmail.com>

src/hooks/useFirestoreDocWithAtom.ts

@@ -170,7 +170,7 @@ export default useFirestoreDocWithAtom;
  * Create the Firestore document reference.
  * Put code in a function so the results can be compared by useMemoValue.
  */
-const getDocRef = <T>(
+export const getDocRef = <T>(
   firebaseDb: Firestore,
   path: string | undefined,
   pathSegments?: Array<string | undefined>

src/sources/TableSourceFirestore/TableSourceFirestore.tsx

@@ -1,6 +1,15 @@
-import { memo, useCallback } from "react";
-import { useAtom } from "jotai";
-import { FirestoreError } from "firebase/firestore";
+import { memo, useCallback, useEffect } from "react";
+import { useAtom, useSetAtom } from "jotai";
+import useMemoValue from "use-memo-value";
+import { cloneDeep, set } from "lodash-es";
+import {
+  FirestoreError,
+  deleteField,
+  refEqual,
+  setDoc,
+} from "firebase/firestore";
+import { useSnackbar } from "notistack";
+import { useErrorHandler } from "react-error-boundary";
 
 import {
   tableScope,
@@ -16,29 +25,81 @@ import {
   tableNextPageAtom,
   serverDocCountAtom,
 } from "@src/atoms/tableScope";
-import useFirestoreDocWithAtom from "@src/hooks/useFirestoreDocWithAtom";
+
+import useFirestoreDocWithAtom, {
+  getDocRef,
+} from "@src/hooks/useFirestoreDocWithAtom";
 import useFirestoreCollectionWithAtom from "@src/hooks/useFirestoreCollectionWithAtom";
 
 import useAuditChange from "./useAuditChange";
 import useBulkWriteDb from "./useBulkWriteDb";
 import { handleFirestoreError } from "./handleFirestoreError";
 
-import { useSnackbar } from "notistack";
-import { useErrorHandler } from "react-error-boundary";
 import { getTableSchemaPath } from "@src/utils/table";
+import { TableSchema } from "@src/types/table";
+import { firebaseDbAtom } from "@src/sources/ProjectSourceFirebase";
+import { projectScope } from "@src/atoms/projectScope";
 
 /**
  * When rendered, provides atom values for top-level tables and sub-tables
  */
 export const TableSourceFirestore = memo(function TableSourceFirestore() {
-  // Get tableSettings from tableId and tables in projectScope
+  const [firebaseDb] = useAtom(firebaseDbAtom, projectScope);
   const [tableSettings] = useAtom(tableSettingsAtom, tableScope);
+  const setTableSchema = useSetAtom(tableSchemaAtom, tableScope);
+  const setUpdateTableSchema = useSetAtom(updateTableSchemaAtom, tableScope);
+
+  const { enqueueSnackbar } = useSnackbar();
+
   if (!tableSettings) throw new Error("No table config");
   if (!tableSettings.collection)
     throw new Error("Invalid table config: no collection");
 
+  const tableSchemaDocRef = useMemoValue(
+    getDocRef<TableSchema>(firebaseDb, getTableSchemaPath(tableSettings)),
+    (next, prev) => refEqual(next as any, prev as any)
+  );
   const isCollectionGroup = tableSettings.tableType === "collectionGroup";
 
+  useEffect(() => {
+    if (!tableSchemaDocRef) return;
+
+    setUpdateTableSchema(
+      () => (update: TableSchema, deleteFields?: string[]) => {
+        const updateToDb = cloneDeep(update);
+
+        if (Array.isArray(deleteFields)) {
+          for (const field of deleteFields) {
+            // Use deterministic set firestore sentinel's on schema columns config
+            // Required for nested columns
+            // i.e field = "columns.base.nested.nested"
+            // key: columns, rest: base.nested.nested
+            // set columns["base.nested.nested"] instead columns.base.nested.nested
+            const [key, ...rest] = field.split(".");
+            if (key === "columns") {
+              (updateToDb as any).columns[rest.join(".")] = deleteField();
+            } else {
+              set(updateToDb, field, deleteField());
+            }
+          }
+        }
+
+        // Update UI state to reflect changes immediately to prevent flickering effects
+        setTableSchema((tableSchema) => ({ ...tableSchema, ...update }));
+
+        return setDoc(tableSchemaDocRef, updateToDb, { merge: true }).catch(
+          (e) => {
+            enqueueSnackbar((e as Error).message, { variant: "error" });
+          }
+        );
+      }
+    );
+
+    return () => {
+      setUpdateTableSchema(undefined);
+    };
+  }, [tableSchemaDocRef, setTableSchema, setUpdateTableSchema, enqueueSnackbar]);
+
   // Get tableSchema and store in tableSchemaAtom.
   // If it doesn’t exist, initialize columns
   useFirestoreDocWithAtom(
@@ -47,7 +108,6 @@ export const TableSourceFirestore = memo(function TableSourceFirestore() {
     getTableSchemaPath(tableSettings),
     {
       createIfNonExistent: { columns: {} },
-      updateDataAtom: updateTableSchemaAtom,
       disableSuspense: true,
     }
   );
@@ -58,7 +118,6 @@ export const TableSourceFirestore = memo(function TableSourceFirestore() {
   const [page] = useAtom(tablePageAtom, tableScope);
   // Get documents from collection and store in tableRowsDbAtom
   // and handle some errors with snackbars
-  const { enqueueSnackbar } = useSnackbar();
   const elevateError = useErrorHandler();
   const handleErrorCallback = useCallback(
     (error: FirestoreError) =>