From 8d550a7ea5708ccae1136f4c0445fd9e4573341c Mon Sep 17 00:00:00 2001 From: HugoCasa Date: Fri, 4 Aug 2023 16:48:09 +0200 Subject: [PATCH] feat: add toggle for postgres between public and all schemas (#1991) --- .../lib/components/DBSchemaExplorer.svelte | 38 ++++++++++++++++--- .../lib/components/codeGen/ScriptFix.svelte | 5 ++- .../lib/components/codeGen/ScriptGen.svelte | 32 +++++++++++----- frontend/src/lib/components/codeGen/lib.ts | 34 ++++++++++++----- frontend/src/lib/stores.ts | 16 +++++++- 5 files changed, 97 insertions(+), 28 deletions(-) diff --git a/frontend/src/lib/components/DBSchemaExplorer.svelte b/frontend/src/lib/components/DBSchemaExplorer.svelte index cbbebae501..ce85b3230a 100644 --- a/frontend/src/lib/components/DBSchemaExplorer.svelte +++ b/frontend/src/lib/components/DBSchemaExplorer.svelte @@ -1,11 +1,13 @@ {#if $dbSchema && resourcePath} @@ -176,7 +196,13 @@ export async function main(args: any) { - + {#if resourceType === 'postgresql'} + + + + + {/if} + {/if} diff --git a/frontend/src/lib/components/codeGen/ScriptFix.svelte b/frontend/src/lib/components/codeGen/ScriptFix.svelte index 411d111706..31947b7ced 100644 --- a/frontend/src/lib/components/codeGen/ScriptFix.svelte +++ b/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 diff --git a/frontend/src/lib/components/codeGen/ScriptGen.svelte b/frontend/src/lib/components/codeGen/ScriptGen.svelte index 0847523143..16e0d42cb7 100644 --- a/frontend/src/lib/components/codeGen/ScriptGen.svelte +++ b/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 @@ - {#if lang === 'postgresql' && $dbSchema} -

- Will take into account the DB schema - - In order to better generate the script, we pass the selected DB schema to GPT-4. - -

+ {#if ['postgresql', 'mysql'].includes(lang) && $dbSchema} +
+

+ Will take into account the DB schema + + In order to better generate the script, we pass the selected DB schema to GPT-4. + +

+ {#if lang === 'postgresql'} + + + + + {/if} +
{/if} {:else}

{ 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 } diff --git a/frontend/src/lib/stores.ts b/frontend/src/lib/stores.ts index 56f62019ae..e0de369418 100644 --- a/frontend/src/lib/stores.ts +++ b/frontend/src/lib/stores.ts @@ -66,7 +66,21 @@ export const hubScripts = writable< | undefined >(undefined) export const existsOpenaiResourcePath = writable(false) -export const dbSchema = writable(undefined) +export type DBSchema = { + [schemaKey: string]: { + [tableKey: string]: { + [columnKey: string]: { + type: string + default: string + required: boolean + } + } + } +} + +export const dbSchema = writable(undefined) + +export const dbSchemaPublicOnly = writable(true) export function switchWorkspace(workspace: string | undefined) { localStorage.removeItem('flow')