feat: add toggle for postgres between public and all schemas (#1991)

frontend/src/lib/components/DBSchemaExplorer.svelte

@@ -1,11 +1,13 @@
 <script lang="ts">
 	import { JobService, Preview } from '$lib/gen'
-	import { dbSchema, workspaceStore } from '$lib/stores'
+	import { dbSchema, dbSchemaPublicOnly, workspaceStore, type DBSchema } from '$lib/stores'
 	import { onDestroy } from 'svelte'
 	import Button from './common/button/Button.svelte'
 	import Drawer from './common/drawer/Drawer.svelte'
 	import DrawerContent from './common/drawer/DrawerContent.svelte'
 	import ObjectViewer from './propertyPicker/ObjectViewer.svelte'
+	import ToggleButton from './common/toggleButton-v2/ToggleButton.svelte'
+	import ToggleButtonGroup from './common/toggleButton-v2/ToggleButtonGroup.svelte'
 
 	export let resourceType: string | undefined
 	export let resourcePath: String | undefined = undefined
@@ -155,12 +157,30 @@ export async function main(args: any) {
 		}, 1000)
 	}
 
-	$: resourcePath && resourceType && ['postgresql', 'mysql'].includes(resourceType) && getSchema()
-	$: !resourcePath && $dbSchema && dbSchema.set(undefined)
+	function formatSchema(
+		schema: DBSchema,
+		resourceType: string | undefined,
+		dbSchemaPublicOnly: boolean
+	) {
+		if (resourceType === 'postgresql' && dbSchemaPublicOnly) {
+			return schema.public || schema
+		} else if (resourceType === 'mysql' && Object.keys(schema).length === 1) {
+			return schema[Object.keys(schema)[0]]
+		} else {
+			return schema
+		}
+	}
 
-	onDestroy(() => {
+	$: resourcePath && ['postgresql', 'mysql'].includes(resourceType || '') && getSchema()
+
+	function clearSchema() {
 		dbSchema.set(undefined)
-	})
+		dbSchemaPublicOnly.set(true)
+	}
+
+	$: !resourcePath && $dbSchema && clearSchema()
+
+	onDestroy(clearSchema)
 </script>
 
 {#if $dbSchema && resourcePath}
@@ -176,7 +196,13 @@ export async function main(args: any) {
 	</Button>
 	<Drawer bind:this={drawer} size="800px">
 		<DrawerContent title="DB Schema Explorer" on:close={drawer.closeDrawer}>
-			<ObjectViewer json={$dbSchema} pureViewer />
+			{#if resourceType === 'postgresql'}
+				<ToggleButtonGroup class="mb-4" bind:selected={$dbSchemaPublicOnly}>
+					<ToggleButton value={true} label="Public" />
+					<ToggleButton value={false} label="All" />
+				</ToggleButtonGroup>
+			{/if}
+			<ObjectViewer json={formatSchema($dbSchema, resourceType, $dbSchemaPublicOnly)} pureViewer />
 		</DrawerContent>
 	</Drawer>
 {/if}

frontend/src/lib/components/codeGen/ScriptFix.svelte

@@ -6,7 +6,7 @@
 	import { sendUserToast } from '$lib/toast'
 	import type Editor from '../Editor.svelte'
 	import { faCheck, faClose, faMagicWandSparkles } from '@fortawesome/free-solid-svg-icons'
-	import { dbSchema, existsOpenaiResourcePath } from '$lib/stores'
+	import { dbSchema, dbSchemaPublicOnly, existsOpenaiResourcePath } from '$lib/stores'
 	import type DiffEditor from '../DiffEditor.svelte'
 	import { scriptLangToEditorLang } from '$lib/scripts'
 	import Popover from '../Popover.svelte'
@@ -39,7 +39,8 @@
 				language: lang,
 				code: editor?.getCode() || '',
 				error,
-				dbSchema: lang === 'postgresql' ? $dbSchema : undefined
+				dbSchema: $dbSchema,
+				dbSchemaPublicOnly: $dbSchemaPublicOnly
 			})
 			generatedCode = result.code
 			explanation = result.explanation

frontend/src/lib/components/codeGen/ScriptGen.svelte

@@ -8,12 +8,14 @@
 	import { faCheck, faClose, faMagicWandSparkles } from '@fortawesome/free-solid-svg-icons'
 	import Popup from '../common/popup/Popup.svelte'
 	import { Icon } from 'svelte-awesome'
-	import { dbSchema, existsOpenaiResourcePath } from '$lib/stores'
+	import { dbSchema, dbSchemaPublicOnly, existsOpenaiResourcePath } from '$lib/stores'
 	import type DiffEditor from '../DiffEditor.svelte'
 	import { scriptLangToEditorLang } from '$lib/scripts'
 	import type { Selection } from 'monaco-editor/esm/vs/editor/editor.api'
 	import type SimpleEditor from '../SimpleEditor.svelte'
 	import Tooltip from '../Tooltip.svelte'
+	import ToggleButtonGroup from '../common/toggleButton-v2/ToggleButtonGroup.svelte'
+	import ToggleButton from '../common/toggleButton-v2/ToggleButton.svelte'
 
 	// props
 	export let iconOnly: boolean = false
@@ -45,14 +47,16 @@
 					language: lang,
 					description: funcDesc,
 					selectedCode,
-					dbSchema: lang === 'postgresql' ? $dbSchema : undefined
+					dbSchema: $dbSchema,
+					dbSchemaPublicOnly: $dbSchemaPublicOnly
 				})
 				generatedCode = originalCode.replace(selectedCode, result.code + '\n')
 			} else {
 				const result = await generateScript({
 					language: lang,
 					description: funcDesc,
-					dbSchema: lang === 'postgresql' ? $dbSchema : undefined
+					dbSchema: $dbSchema,
+					dbSchemaPublicOnly: $dbSchemaPublicOnly
 				})
 				generatedCode = result.code
 			}
@@ -227,13 +231,21 @@
 						<Icon data={faMagicWandSparkles} />
 					</Button>
 				</div>
-				{#if lang === 'postgresql' && $dbSchema}
-					<p class="text-sm mt-2">
-						Will take into account the DB schema
-						<Tooltip>
-							In order to better generate the script, we pass the selected DB schema to GPT-4.
-						</Tooltip>
-					</p>
+				{#if ['postgresql', 'mysql'].includes(lang) && $dbSchema}
+					<div class="flex flex-row items-center justify-between w-96 mt-2">
+						<p class="text-sm">
+							Will take into account the DB schema
+							<Tooltip>
+								In order to better generate the script, we pass the selected DB schema to GPT-4.
+							</Tooltip>
+						</p>
+						{#if lang === 'postgresql'}
+							<ToggleButtonGroup class="w-auto shrink-0" bind:selected={$dbSchemaPublicOnly}>
+								<ToggleButton value={true} label="Public schema" />
+								<ToggleButton value={false} label="All schemas" />
+							</ToggleButtonGroup>
+						{/if}
+					</div>
 				{/if}
 			{:else}
 				<p class="text-sm"

frontend/src/lib/components/codeGen/lib.ts

@@ -2,7 +2,7 @@ import { OpenAI } from 'openai'
 import { OpenAPI } from '../../gen/core/OpenAPI'
 import { ResourceService, Script, WorkspaceService } from '../../gen'
 
-import { existsOpenaiResourcePath, workspaceStore } from '$lib/stores'
+import { existsOpenaiResourcePath, workspaceStore, type DBSchema } from '$lib/stores'
 import { formatResourceTypes } from './utils'
 
 import { EDIT_CONFIG, FIX_CONFIG, GEN_CONFIG } from './prompts'
@@ -40,12 +40,12 @@ workspaceStore.subscribe(async (value) => {
 
 interface BaseOptions {
 	language: Script.language | 'frontend'
-	dbSchema?: object
+	dbSchema: DBSchema | undefined
+	dbSchemaPublicOnly: boolean
 }
 
 interface ScriptGenerationOptions extends BaseOptions {
 	description: string
-	dbSchema?: object
 }
 
 interface EditScriptOptions extends ScriptGenerationOptions {
@@ -72,25 +72,41 @@ async function addResourceTypes(scriptOptions: BaseOptions, workspace: string, p
 
 function addDBSChema(scriptOptions: BaseOptions, prompt: string) {
 	if (['mysql', 'postgresql'].includes(scriptOptions.language) && scriptOptions.dbSchema) {
-		const { dbSchema } = scriptOptions
-		const smallerSchema = {}
+		const { dbSchema, dbSchemaPublicOnly, language } = scriptOptions
+		let smallerSchema: {
+			[schemaKey: string]: {
+				[tableKey: string]: Array<[string, string, boolean, string?]>
+			}
+		} = {}
 		for (const schemaKey in dbSchema) {
+			smallerSchema[schemaKey] = {}
 			for (const tableKey in dbSchema[schemaKey]) {
-				smallerSchema[tableKey] = []
+				smallerSchema[schemaKey][tableKey] = []
 				for (const colKey in dbSchema[schemaKey][tableKey]) {
 					const col = dbSchema[schemaKey][tableKey][colKey]
-					const p = [colKey, col.type, col.required]
+					const p: [string, string, boolean, string?] = [colKey, col.type, col.required]
 					if (col.default) {
 						p.push(col.default)
 					}
-					smallerSchema[tableKey].push(p)
+					smallerSchema[schemaKey][tableKey].push(p)
 				}
 			}
 		}
+
+		let finalSchema:
+			| typeof smallerSchema
+			| {
+					[tableKey: string]: Array<[string, string, boolean, string?]>
+			  } = smallerSchema
+		if (language === 'postgresql' && dbSchemaPublicOnly) {
+			finalSchema = smallerSchema.public || smallerSchema
+		} else if (language === 'mysql' && Object.keys(smallerSchema).length === 1) {
+			finalSchema = smallerSchema[Object.keys(smallerSchema)[0]]
+		}
 		prompt =
 			prompt +
 			"\nHere's the database schema, each column is in the format [name, type, required, default?]: " +
-			JSON.stringify(smallerSchema)
+			JSON.stringify(finalSchema)
 	}
 	return prompt
 }

frontend/src/lib/stores.ts

@@ -66,7 +66,21 @@ export const hubScripts = writable<
 	| undefined
 >(undefined)
 export const existsOpenaiResourcePath = writable<boolean>(false)
-export const dbSchema = writable<object | undefined>(undefined)
+export type DBSchema = {
+	[schemaKey: string]: {
+		[tableKey: string]: {
+			[columnKey: string]: {
+				type: string
+				default: string
+				required: boolean
+			}
+		}
+	}
+}
+
+export const dbSchema = writable<DBSchema | undefined>(undefined)
+
+export const dbSchemaPublicOnly = writable<boolean>(true)
 
 export function switchWorkspace(workspace: string | undefined) {
 	localStorage.removeItem('flow')