From 13d2b866f952e112a2da20c22c09460d9081a042 Mon Sep 17 00:00:00 2001
From: Joshen Lim <joshenlimek@gmail.com>
Date: Thu, 9 Feb 2023 15:04:30 +0800
Subject: [PATCH] Fix table editor foreign key selector cannot select tables
 outside of public schema

---
 .../ColumnEditor/ColumnEditor.tsx             |  1 -
 .../ForeignKeySelector/ForeignKeySelector.tsx | 50 +++++++++++++++----
 .../TableEditor/ColumnManagement.tsx          |  1 -
 .../TableEditorLayout/TableEditorMenu.tsx     |  2 +-
 4 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/studio/components/interfaces/TableGridEditor/SidePanelEditor/ColumnEditor/ColumnEditor.tsx b/studio/components/interfaces/TableGridEditor/SidePanelEditor/ColumnEditor/ColumnEditor.tsx
index 61a169e9555..251104fc29e 100644
--- a/studio/components/interfaces/TableGridEditor/SidePanelEditor/ColumnEditor/ColumnEditor.tsx
+++ b/studio/components/interfaces/TableGridEditor/SidePanelEditor/ColumnEditor/ColumnEditor.tsx
@@ -375,7 +375,6 @@ const ColumnEditor: FC<Props> = ({
       )}
 
       <ForeignKeySelector
-        tables={tables}
         column={columnFields}
         visible={isEditingRelation}
         closePanel={() => setIsEditingRelation(false)}
diff --git a/studio/components/interfaces/TableGridEditor/SidePanelEditor/ForeignKeySelector/ForeignKeySelector.tsx b/studio/components/interfaces/TableGridEditor/SidePanelEditor/ForeignKeySelector/ForeignKeySelector.tsx
index e8146e33c43..f43723572de 100644
--- a/studio/components/interfaces/TableGridEditor/SidePanelEditor/ForeignKeySelector/ForeignKeySelector.tsx
+++ b/studio/components/interfaces/TableGridEditor/SidePanelEditor/ForeignKeySelector/ForeignKeySelector.tsx
@@ -1,16 +1,16 @@
 import React, { FC, useEffect, useState } from 'react'
 import { get, find, isEmpty, sortBy } from 'lodash'
 import { Dictionary } from 'components/grid'
-import { SidePanel, Input, Listbox, IconHelpCircle } from 'ui'
-import type { PostgresTable, PostgresColumn } from '@supabase/postgres-meta'
+import { SidePanel, Input, Listbox, IconHelpCircle, IconDatabase } from 'ui'
+import type { PostgresTable, PostgresColumn, PostgresSchema } from '@supabase/postgres-meta'
 
+import { useStore } from 'hooks'
 import ActionBar from '../ActionBar'
 import { ForeignKey } from './ForeignKeySelector.types'
 import { ColumnField } from '../SidePanelEditor.types'
 import InformationBox from 'components/ui/InformationBox'
 
 interface Props {
-  tables: PostgresTable[]
   column: ColumnField
   metadata?: any
   visible: boolean
@@ -18,16 +18,15 @@ interface Props {
   saveChanges: (value: { table: PostgresTable; column: PostgresColumn } | undefined) => void
 }
 
-const ForeignKeySelector: FC<Props> = ({
-  tables = [] as PostgresTable[],
-  column,
-  visible = false,
-  closePanel,
-  saveChanges,
-}) => {
+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 schemas = meta.schemas.list()
+  const tables = meta.tables.list((table: PostgresTable) => table.schema === selectedSchema)
+
   const foreignKey = column?.foreignKey
   const selectedTable: PostgresTable | undefined = find(tables, {
     name: selectedForeignKey?.table,
@@ -37,6 +36,10 @@ const ForeignKeySelector: FC<Props> = ({
     name: selectedForeignKey?.column,
   })
 
+  useEffect(() => {
+    meta.tables.loadBySchema(selectedSchema)
+  }, [selectedSchema])
+
   useEffect(() => {
     // Reset the state of the side panel
     if (visible) {
@@ -87,6 +90,9 @@ const ForeignKeySelector: FC<Props> = ({
 
   const onSaveChanges = (resolve: () => void) => {
     const errors = {} as Dictionary<any>
+    if (!selectedForeignKey?.table) {
+      errors['table'] = 'Please select a table'
+    }
     if (selectedForeignKey?.table && !selectedForeignKey.column) {
       errors['column'] = `The table ${selectedForeignKey.table} has no columns`
     }
@@ -140,6 +146,30 @@ const ForeignKeySelector: FC<Props> = ({
             urlLabel="Postgres Foreign Key Documentation"
           />
 
+          <Listbox
+            label="Select a schema"
+            value={selectedSchema}
+            error={errors.schema}
+            onChange={(value: string) => setSelectedSchema(value)}
+          >
+            {schemas.map((schema: PostgresSchema) => {
+              return (
+                <Listbox.Option
+                  key={schema.id}
+                  value={schema.name}
+                  label={schema.name}
+                  addOnBefore={() => <IconDatabase size={16} strokeWidth={1.5} />}
+                >
+                  <div className="flex items-center gap-2">
+                    {/* For aria searching to target the schema name instead of schema */}
+                    <span className="hidden">{schema.name}</span>
+                    <span className="text-scale-1200">{schema.name}</span>
+                  </div>
+                </Listbox.Option>
+              )
+            })}
+          </Listbox>
+
           <Listbox
             label="Select a table to reference to"
             value={selectedTable?.id}
diff --git a/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/ColumnManagement.tsx b/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/ColumnManagement.tsx
index f601b5962bc..a834c258255 100644
--- a/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/ColumnManagement.tsx
+++ b/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/ColumnManagement.tsx
@@ -353,7 +353,6 @@ const ColumnManagement: FC<Props> = ({
         </div>
       </div>
       <ForeignKeySelector
-        tables={tables}
         column={selectedColumnToEditRelation as ColumnField}
         visible={!isUndefined(selectedColumnToEditRelation)}
         closePanel={() => setSelectedColumnToEditRelation(undefined)}
diff --git a/studio/components/layouts/TableEditorLayout/TableEditorMenu.tsx b/studio/components/layouts/TableEditorLayout/TableEditorMenu.tsx
index 72d18037fa5..dcdfc12bd0d 100644
--- a/studio/components/layouts/TableEditorLayout/TableEditorMenu.tsx
+++ b/studio/components/layouts/TableEditorLayout/TableEditorMenu.tsx
@@ -43,7 +43,7 @@ const TableEditorMenu: FC<Props> = ({
   onDeleteTable = () => {},
   onDuplicateTable = () => {},
 }) => {
-  const { meta, ui } = useStore()
+  const { meta } = useStore()
   const { id, ref } = useParams()
 
   const schemas: PostgresSchema[] = meta.schemas.list()