From 21e0465dba299e548f94e17ae5456933aa1b2398 Mon Sep 17 00:00:00 2001
From: Joshen Lim <joshenlimek@gmail.com>
Date: Fri, 10 Feb 2023 13:36:32 +0800
Subject: [PATCH] Fix:

---
 .../ui/src/components/Listbox/Listbox2.tsx    | 26 ++++----------
 .../ForeignKeySelector/ForeignKeySelector.tsx | 35 ++++++++++++-------
 2 files changed, 29 insertions(+), 32 deletions(-)

diff --git a/packages/ui/src/components/Listbox/Listbox2.tsx b/packages/ui/src/components/Listbox/Listbox2.tsx
index 5366c6062b5..78f9a4bbd06 100644
--- a/packages/ui/src/components/Listbox/Listbox2.tsx
+++ b/packages/ui/src/components/Listbox/Listbox2.tsx
@@ -18,8 +18,7 @@ function classNames(...classes: any) {
   return classes.filter(Boolean).join(' ')
 }
 
-export interface Props
-  extends Omit<React.InputHTMLAttributes<HTMLButtonElement>, 'size'> {
+export interface Props extends Omit<React.InputHTMLAttributes<HTMLButtonElement>, 'size'> {
   className?: string
   children: React.ReactNode
   descriptionText?: string
@@ -73,14 +72,8 @@ function Listbox({
 
   const triggerRef = useRef<HTMLButtonElement>(null)
 
-  const {
-    formContextOnChange,
-    values,
-    errors,
-    handleBlur,
-    touched,
-    fieldLevelValidation,
-  } = useFormContext()
+  const { formContextOnChange, values, errors, handleBlur, touched, fieldLevelValidation } =
+    useFormContext()
 
   if (values && !value) {
     defaultValue = values[id || name]
@@ -257,9 +250,7 @@ function Listbox({
           className={__styles.options_container}
         >
           <div>
-            <SelectContext.Provider
-              value={{ onChange: handleOnChange, selected }}
-            >
+            <SelectContext.Provider value={{ onChange: handleOnChange, selected }}>
               {children}
             </SelectContext.Provider>
           </div>
@@ -313,9 +304,7 @@ function SelectOption({
             <div className={__styles.option_inner}>
               {addOnBefore && addOnBefore({ active, selected })}
               <span>
-                {typeof children === 'function'
-                  ? children({ active, selected })
-                  : children}
+                {typeof children === 'function' ? children({ active, selected }) : children}
               </span>
             </div>
 
@@ -326,10 +315,7 @@ function SelectOption({
                   active ? __styles.option_check_active : ''
                 )}
               >
-                <IconCheck
-                  className={__styles.option_check_icon}
-                  aria-hidden="true"
-                />
+                <IconCheck className={__styles.option_check_icon} aria-hidden="true" />
               </span>
             ) : null}
           </DropdownMenuPrimitive.Item>
diff --git a/studio/components/interfaces/TableGridEditor/SidePanelEditor/ForeignKeySelector/ForeignKeySelector.tsx b/studio/components/interfaces/TableGridEditor/SidePanelEditor/ForeignKeySelector/ForeignKeySelector.tsx
index f43723572de..823ad2382d4 100644
--- a/studio/components/interfaces/TableGridEditor/SidePanelEditor/ForeignKeySelector/ForeignKeySelector.tsx
+++ b/studio/components/interfaces/TableGridEditor/SidePanelEditor/ForeignKeySelector/ForeignKeySelector.tsx
@@ -21,11 +21,16 @@ interface Props {
 const ForeignKeySelector: FC<Props> = ({ column, visible = false, closePanel, saveChanges }) => {
   const { meta } = useStore()
   const [errors, setErrors] = useState<any>({})
-  const [selectedSchema, setSelectedSchema] = useState('public')
-  const [selectedForeignKey, setSelectedForeignKey] = useState<ForeignKey>()
+  const [selectedForeignKey, setSelectedForeignKey] = useState<ForeignKey>({
+    schema: 'public',
+    table: '',
+    column: '',
+  })
 
   const schemas = meta.schemas.list()
-  const tables = meta.tables.list((table: PostgresTable) => table.schema === selectedSchema)
+  const tables = meta.tables.list(
+    (table: PostgresTable) => table.schema === selectedForeignKey.schema
+  )
 
   const foreignKey = column?.foreignKey
   const selectedTable: PostgresTable | undefined = find(tables, {
@@ -36,14 +41,11 @@ const ForeignKeySelector: FC<Props> = ({ column, visible = false, closePanel, sa
     name: selectedForeignKey?.column,
   })
 
-  useEffect(() => {
-    meta.tables.loadBySchema(selectedSchema)
-  }, [selectedSchema])
-
   useEffect(() => {
     // Reset the state of the side panel
     if (visible) {
       setErrors({})
+
       if (foreignKey) {
         setSelectedForeignKey({
           schema: foreignKey.target_table_schema,
@@ -52,7 +54,7 @@ const ForeignKeySelector: FC<Props> = ({ column, visible = false, closePanel, sa
         })
       } else {
         setSelectedForeignKey({
-          schema: '',
+          schema: 'public',
           table: '',
           column: '',
         })
@@ -60,6 +62,12 @@ const ForeignKeySelector: FC<Props> = ({ column, visible = false, closePanel, sa
     }
   }, [visible])
 
+  const updateSelectedSchema = (schema: string) => {
+    meta.tables.loadBySchema(schema)
+    const updatedForeignKey = { ...selectedForeignKey, schema, table: '', column: '' } as ForeignKey
+    setSelectedForeignKey(updatedForeignKey)
+  }
+
   const updateSelectedTable = (tableId: number) => {
     setErrors({})
     if (!tableId) {
@@ -147,10 +155,11 @@ const ForeignKeySelector: FC<Props> = ({ column, visible = false, closePanel, sa
           />
 
           <Listbox
+            id="schema"
             label="Select a schema"
-            value={selectedSchema}
+            value={selectedForeignKey.schema}
             error={errors.schema}
-            onChange={(value: string) => setSelectedSchema(value)}
+            onChange={(value: string) => updateSelectedSchema(value)}
           >
             {schemas.map((schema: PostgresSchema) => {
               return (
@@ -171,12 +180,13 @@ const ForeignKeySelector: FC<Props> = ({ column, visible = false, closePanel, sa
           </Listbox>
 
           <Listbox
+            id="table"
             label="Select a table to reference to"
-            value={selectedTable?.id}
+            value={selectedTable?.id ?? 1}
             error={errors.table}
             onChange={(value: string) => updateSelectedTable(Number(value))}
           >
-            <Listbox.Option key="empty" value="" label="---">
+            <Listbox.Option key="empty" value={1} label="---">
               ---
             </Listbox.Option>
             {/* @ts-ignore */}
@@ -209,6 +219,7 @@ const ForeignKeySelector: FC<Props> = ({ column, visible = false, closePanel, sa
                 />
               ) : (
                 <Listbox
+                  id="column"
                   value={selectedColumn?.id}
                   // @ts-ignore
                   label={