From 7f84b07a0737cecccc89ea1e10fde0fd7b7ccb2b Mon Sep 17 00:00:00 2001
From: Alaister Young <a@alaisteryoung.com>
Date: Mon, 13 Mar 2023 11:52:46 +0800
Subject: [PATCH] adds optimistic updates for foreign row references

---
 .../SidePanelEditor/SidePanelEditor.tsx       | 60 ++++++++++++++++++-
 1 file changed, 58 insertions(+), 2 deletions(-)

diff --git a/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.tsx b/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.tsx
index 64021e6c91c..362d71fdb74 100644
--- a/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.tsx
+++ b/studio/components/interfaces/TableGridEditor/SidePanelEditor/SidePanelEditor.tsx
@@ -1,5 +1,5 @@
 import { useState } from 'react'
-import { useQueryClient } from '@tanstack/react-query'
+import { QueryKey, useQueryClient } from '@tanstack/react-query'
 import { find, isEmpty, isUndefined, noop } from 'lodash'
 import { Dictionary } from 'components/grid'
 import { Modal } from 'ui'
@@ -75,7 +75,63 @@ const SidePanelEditor = ({
 
   const { project } = useProjectContext()
   const { mutateAsync: createTableRow } = useTableRowCreateMutation()
-  const { mutateAsync: updateTableRow } = useTableRowUpdateMutation()
+  const { mutateAsync: updateTableRow } = useTableRowUpdateMutation({
+    async onMutate({ projectRef, table, configuration, payload }) {
+      closePanel()
+
+      const primaryKeyColumns = new Set(Object.keys(configuration.identifiers))
+
+      const queryKey = sqlKeys.query(projectRef, [
+        table.schema,
+        table.name,
+        { table: { name: table.name, schema: table.schema } },
+      ])
+
+      await queryClient.cancelQueries(queryKey)
+
+      const previousRowsQueries = queryClient.getQueriesData<{ result: any[] }>(queryKey)
+
+      queryClient.setQueriesData<{ result: any[] }>(queryKey, (old) => {
+        return {
+          result:
+            old?.result.map((row) => {
+              // match primary keys
+              if (
+                Object.entries(row)
+                  .filter(([key]) => primaryKeyColumns.has(key))
+                  .every(([key, value]) => value === configuration.identifiers[key])
+              ) {
+                return { ...row, ...payload }
+              }
+
+              return row
+            }) ?? [],
+        }
+      })
+
+      return { previousRowsQueries }
+    },
+    onError(error, _variables, context) {
+      const { previousRowsQueries } = context as {
+        previousRowsQueries: [
+          QueryKey,
+          (
+            | {
+                result: any[]
+              }
+            | undefined
+          )
+        ][]
+      }
+
+      previousRowsQueries.forEach(([queryKey, previousRows]) => {
+        if (previousRows) {
+          queryClient.setQueriesData(queryKey, previousRows)
+        }
+        queryClient.invalidateQueries(queryKey)
+      })
+    },
+  })
 
   const saveRow = async (
     payload: any,