move ColumnMenu state to internal, provide ref

www/src/components/Table/ColumnEditor/index.tsx

@@ -65,13 +65,14 @@ const useStyles = makeStyles(theme =>
   })
 );
 
+// TODO: REMOVE THIS OLD COMPONENT
 export default function ColumnEditor() {
   const {
     tableState,
     tableActions,
     selectedColumnHeader,
     setSelectedColumnHeader,
-  } = useFiretableContext();
+  } = useFiretableContext() as any;
   const actions = tableActions!.column;
   const { column, anchorEl } = selectedColumnHeader ?? {};
 

www/src/components/Table/ColumnHeader.tsx

@@ -90,17 +90,21 @@ const ColumnHeader: Column<any>["headerRenderer"] = ({ column }) => {
   const classes = useStyles();
 
   const {
-    setSelectedColumnHeader,
     tableState,
     tableActions,
     userClaims,
+    columnMenuRef,
   } = useFiretableContext();
-  if (!setSelectedColumnHeader || !tableState || !tableActions) return null;
+  if (!columnMenuRef || !tableState || !tableActions) return null;
   const { orderBy } = tableState;
 
   const handleClick = (
     event: React.MouseEvent<HTMLButtonElement, MouseEvent>
-  ) => setSelectedColumnHeader({ column, anchorEl: event.currentTarget });
+  ) =>
+    columnMenuRef?.current?.setSelectedColumnHeader({
+      column,
+      anchorEl: event.currentTarget,
+    });
 
   const isSorted = orderBy?.[0]?.key === (column.key as string);
   const isAsc = isSorted && orderBy?.[0]?.direction === "asc";

www/src/components/Table/ColumnMenu/index.tsx

@@ -1,9 +1,6 @@
 import React, { useState } from "react";
 
 import { createStyles, makeStyles, Menu } from "@material-ui/core";
-
-import MenuContents from "./MenuContents";
-
 import LockOpenIcon from "@material-ui/icons/LockOpen";
 import LockIcon from "@material-ui/icons/Lock";
 import VisibilityOffIcon from "@material-ui/icons/VisibilityOff";
@@ -18,12 +15,17 @@ import ReorderIcon from "@material-ui/icons/Reorder";
 import ColumnPlusBeforeIcon from "assets/icons/ColumnPlusBefore";
 import ColumnPlusAfterIcon from "assets/icons/ColumnPlusAfter";
 import ColumnRemoveIcon from "assets/icons/ColumnRemove";
+
+import MenuContents from "./MenuContents";
 import NameChange from "./NameChange";
 import NewColumn from "./NewColumn";
 import TypeChange from "./TypeChange";
+
 import { useFiretableContext } from "contexts/firetableContext";
 import { FIELDS } from "constants/fields";
 import _find from "lodash/find";
+import { Column } from "react-data-grid";
+import { PopoverProps } from "@material-ui/core";
 
 const INITIAL_MODAL = { type: "", data: {} };
 
@@ -32,6 +34,18 @@ enum ModalStates {
   typeChange = "TYPE_CHANGE",
   new = "NEW_COLUMN",
 }
+
+type SelectedColumnHeader = {
+  column: Column<any> & { [key: string]: any };
+  anchorEl: PopoverProps["anchorEl"];
+};
+export type ColumnMenuRef = {
+  selectedColumnHeader: SelectedColumnHeader | null;
+  setSelectedColumnHeader: React.Dispatch<
+    React.SetStateAction<SelectedColumnHeader | null>
+  >;
+};
+
 const useStyles = makeStyles(theme =>
   createStyles({
     paper: {
@@ -44,12 +58,12 @@ const useStyles = makeStyles(theme =>
 export default function ColumnMenu() {
   const classes = useStyles();
   const [modal, setModal] = useState(INITIAL_MODAL);
-  const {
-    tableState,
-    tableActions,
-    selectedColumnHeader,
-    setSelectedColumnHeader,
-  } = useFiretableContext();
+  const { tableState, tableActions, columnMenuRef } = useFiretableContext();
+
+  const [selectedColumnHeader, setSelectedColumnHeader] = useState();
+  if (columnMenuRef)
+    columnMenuRef.current = { selectedColumnHeader, setSelectedColumnHeader };
+
   if (!tableState || !tableActions) return null;
   const { orderBy } = tableState;
 

www/src/components/Table/FinalColumnHeader.tsx

@@ -39,12 +39,16 @@ const useStyles = makeStyles(theme =>
 const FinalColumnHeader: Column<any>["headerRenderer"] = ({ column }) => {
   const classes = useStyles();
 
-  const { setSelectedColumnHeader } = useFiretableContext();
-  if (!setSelectedColumnHeader) return null;
+  const { columnMenuRef } = useFiretableContext();
+  if (!columnMenuRef) return null;
 
   const handleClick = (
     event: React.MouseEvent<HTMLButtonElement, MouseEvent>
-  ) => setSelectedColumnHeader({ column, anchorEl: event.currentTarget });
+  ) =>
+    columnMenuRef?.current?.setSelectedColumnHeader({
+      column,
+      anchorEl: event.currentTarget,
+    });
 
   return (
     <Grid

www/src/contexts/firetableContext.tsx

@@ -12,11 +12,7 @@ import useSettings from "hooks/useSettings";
 import { useAppContext } from "./appContext";
 import { useSnackContext } from "./snackContext";
 import { SideDrawerRef } from "components/SideDrawer";
-
-type SelectedColumnHeader = {
-  column: Column<any> & { [key: string]: any };
-  anchorEl: PopoverProps["anchorEl"];
-};
+import { ColumnMenuRef } from "components/Table/ColumnMenu";
 
 export type Table = {
   collection: string;
@@ -40,16 +36,12 @@ interface FiretableContextProps {
   createTable: Function;
   userClaims: any;
 
-  // TODO: Investigate if this can be moved out of this context
-  selectedColumnHeader: SelectedColumnHeader | null;
-  setSelectedColumnHeader: React.Dispatch<
-    React.SetStateAction<SelectedColumnHeader | null>
-  >;
-
   // A ref to the data grid. Contains data grid functions
   dataGridRef: React.RefObject<DataGridHandle>;
   // A ref to the side drawer state. Prevents unnecessary re-renders
   sideDrawerRef: React.MutableRefObject<SideDrawerRef | undefined>;
+  // A ref to the column menu. Prevents unnecessary re-renders
+  columnMenuRef: React.MutableRefObject<ColumnMenuRef | undefined>;
 }
 
 const firetableContext = React.createContext<Partial<FiretableContextProps>>(
@@ -67,10 +59,6 @@ export const FiretableContextProvider: React.FC = ({ children }) => {
   const [settings, createTable] = useSettings();
   const [userRoles, setUserRoles] = useState<null | string[]>();
   const [userClaims, setUserClaims] = useState<any>();
-  const [
-    selectedColumnHeader,
-    setSelectedColumnHeader,
-  ] = useState<SelectedColumnHeader | null>(null);
 
   const { currentUser } = useAppContext();
   useEffect(() => {
@@ -135,6 +123,7 @@ export const FiretableContextProvider: React.FC = ({ children }) => {
   // A ref to the data grid. Contains data grid functions
   const dataGridRef = useRef<DataGridHandle>(null);
   const sideDrawerRef = useRef<SideDrawerRef>();
+  const columnMenuRef = useRef<ColumnMenuRef>();
 
   return (
     <firetableContext.Provider
@@ -146,10 +135,9 @@ export const FiretableContextProvider: React.FC = ({ children }) => {
         tables,
         sections,
         userClaims,
-        selectedColumnHeader,
-        setSelectedColumnHeader,
         dataGridRef,
         sideDrawerRef,
+        columnMenuRef,
       }}
     >
       {children}