From 5cca5833e94fc4c8a80e210164da09f2a1ceb677 Mon Sep 17 00:00:00 2001
From: HugoCasa <hugo@casademont.ch>
Date: Fri, 8 Sep 2023 16:49:47 +0200
Subject: [PATCH] feat: snowflake schema explorer + refactoring (#2260)

---
 .../lib/components/DBSchemaExplorer.svelte    | 84 +++++++++++++++----
 frontend/src/lib/components/Editor.svelte     |  6 +-
 frontend/src/lib/components/EditorBar.svelte  |  3 +-
 .../src/lib/components/ModulePreview.svelte   |  3 +-
 .../src/lib/components/ScriptEditor.svelte    |  3 +
 .../InlineScriptEditor.svelte                 |  8 ++
 .../copilot/FlowCopilotStatus.svelte          |  2 +-
 .../lib/components/copilot/ScriptFix.svelte   |  3 +-
 .../lib/components/copilot/ScriptGen.svelte   |  6 +-
 .../flows/content/FlowModuleComponent.svelte  | 11 +++
 .../components/scriptEditor/LogPanel.svelte   |  8 +-
 frontend/src/lib/stores.ts                    | 11 +--
 12 files changed, 113 insertions(+), 35 deletions(-)

diff --git a/frontend/src/lib/components/DBSchemaExplorer.svelte b/frontend/src/lib/components/DBSchemaExplorer.svelte
index 3e747dbcc8..e03218d8e7 100644
--- a/frontend/src/lib/components/DBSchemaExplorer.svelte
+++ b/frontend/src/lib/components/DBSchemaExplorer.svelte
@@ -1,6 +1,12 @@
 <script lang="ts">
 	import { JobService, Preview } from '$lib/gen'
-	import { dbSchemas, workspaceStore, type DBSchema, type GraphqlSchema } from '$lib/stores'
+	import {
+		dbSchemas,
+		workspaceStore,
+		type DBSchema,
+		type GraphqlSchema,
+		type SQLSchema
+	} from '$lib/stores'
 	import Button from './common/button/Button.svelte'
 	import Drawer from './common/drawer/Drawer.svelte'
 	import DrawerContent from './common/drawer/DrawerContent.svelte'
@@ -180,6 +186,51 @@ GROUP BY
                 cols.append(col)
             schema[dataset.dataset_id][row[0]] = cols
     return schema
+`,
+			lang: 'python3'
+		},
+		snowflake: {
+			code: `# requirements:
+# snowflake-connector-python==3.2.0
+from typing import Any
+import snowflake.connector as sf
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+def main(args):
+    if not args["database"]:
+        raise Exception("a selected database is required for the schema explorer")
+    p_key = serialization.load_pem_private_key(
+        args["private_key"].encode(), password=None, backend=default_backend()
+    )
+    pkb = p_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+    ctx = sf.connect(
+        user=args["username"],
+        account=args["account_identifier"],
+        private_key=pkb,
+        warehouse=args["warehouse"],
+        database=args["database"],
+        schema=args["schema"],
+        role=args["role"],
+    )
+    cs = ctx.cursor()
+    rows = cs.execute("select TABLE_SCHEMA, TABLE_NAME, DATA_TYPE, COLUMN_NAME, COLUMN_DEFAULT, IS_NULLABLE from information_schema.columns where table_schema != 'INFORMATION_SCHEMA'")
+    schema = dict()
+    for row in rows:
+        if row[0] not in schema:
+            schema[row[0]] = dict()
+        if row[1] not in schema[row[0]]:
+            schema[row[0]][row[1]] = dict()
+        schema[row[0]][row[1]][row[3]] = {
+            "type": row[2],
+            "required": row[5] == "YES",
+        }
+        if row[4] is not None:
+            schema[row[0]][row[1]][row[3]]["default"] = row[4]
+    return schema
 `,
 			lang: 'python3'
 		}
@@ -210,19 +261,18 @@ GROUP BY
 					if (!testResult.success) {
 						console.error(testResult.result?.['error']?.['message'])
 					} else {
-						if (resourceType === 'postgresql') {
-							$dbSchemas[resourcePath] = {
-								lang: 'postgresql',
-								schema: testResult.result,
-								publicOnly: true
-							}
-						} else if (
-							resourceType !== undefined &&
-							['mysql', 'graphql', 'bigquery'].includes(resourceType)
-						) {
-							$dbSchemas[resourcePath] = {
-								lang: resourceType as 'mysql' | 'graphql' | 'bigquery',
-								schema: testResult.result
+						if (resourceType !== undefined) {
+							if (resourceType !== 'graphql') {
+								$dbSchemas[resourcePath] = {
+									lang: resourceType as SQLSchema['lang'],
+									schema: testResult.result,
+									publicOnly: !!testResult.result.public || !!testResult.result.PUBLIC
+								}
+							} else {
+								$dbSchemas[resourcePath] = {
+									lang: 'graphql',
+									schema: testResult.result
+								}
 							}
 						}
 					}
@@ -253,8 +303,8 @@ GROUP BY
 	}
 
 	function formatSchema(dbSchema: DBSchema) {
-		if (dbSchema.lang === 'postgresql' && dbSchema.publicOnly) {
-			return dbSchema.schema.public || dbSchema
+		if (dbSchema.lang !== 'graphql' && dbSchema.publicOnly) {
+			return dbSchema.schema.public || dbSchema.schema.PUBLIC || dbSchema
 		} else if (dbSchema.lang === 'mysql' && Object.keys(dbSchema.schema).length === 1) {
 			return dbSchema.schema[Object.keys(dbSchema.schema)[0]]
 		} else {
@@ -299,7 +349,7 @@ GROUP BY
 					>Refresh
 				</Button>
 			</svelte:fragment>
-			{#if dbSchema.lang === 'postgresql'}
+			{#if dbSchema.lang !== 'graphql' && (dbSchema.schema?.public || dbSchema.schema?.PUBLIC)}
 				<ToggleButtonGroup class="mb-4" bind:selected={dbSchema.publicOnly}>
 					<ToggleButton value={true} label="Public" />
 					<ToggleButton value={false} label="All" />
diff --git a/frontend/src/lib/components/Editor.svelte b/frontend/src/lib/components/Editor.svelte
index 8dd8800110..06915ee038 100644
--- a/frontend/src/lib/components/Editor.svelte
+++ b/frontend/src/lib/components/Editor.svelte
@@ -70,6 +70,7 @@
 	export let yContent: Text | undefined = undefined
 	export let awareness: any | undefined = undefined
 	export let folding = false
+	export let args: Record<string, any> | undefined = undefined
 
 	languages.typescript.typescriptDefaults.setModeConfiguration({
 		completionItems: false,
@@ -224,7 +225,8 @@
 	let command: Disposable | undefined = undefined
 
 	let sqlSchemaCompletor: Disposable | undefined = undefined
-	$: dbSchema = $dbSchemas[Object.keys($dbSchemas)[0]]
+	$: args &&
+		(dbSchema = $dbSchemas[(lang === 'graphql' ? args.api : args.database)?.replace('$res:', '')])
 	$: dbSchema && ['sql', 'graphql'].includes(lang) && addDBSchemaCompletions()
 	$: (!dbSchema || lang !== 'sql') && sqlSchemaCompletor && sqlSchemaCompletor.dispose()
 	$: (!dbSchema || lang !== 'graphql') && graphqlService && graphqlService.setSchemaConfig([])
@@ -242,7 +244,7 @@
 					introspectionJSON: schema
 				}
 			])
-		} else if (schemaLang === 'mysql' || schemaLang === 'postgresql') {
+		} else {
 			if (sqlSchemaCompletor) {
 				sqlSchemaCompletor.dispose()
 			}
diff --git a/frontend/src/lib/components/EditorBar.svelte b/frontend/src/lib/components/EditorBar.svelte
index 8b6d98d6c8..cccf6d948b 100644
--- a/frontend/src/lib/components/EditorBar.svelte
+++ b/frontend/src/lib/components/EditorBar.svelte
@@ -62,6 +62,7 @@
 	export let collabUsers: { name: string }[] = []
 	export let scriptPath: string | undefined = undefined
 	export let diffEditor: DiffEditor | undefined = undefined
+	export let args: Record<string, any>
 
 	let contextualVariablePicker: ItemPicker
 	let variablePicker: ItemPicker
@@ -530,7 +531,7 @@
 				</div>
 			{/if}
 
-			<ScriptGen {editor} {diffEditor} {lang} {iconOnly} />
+			<ScriptGen {editor} {diffEditor} {lang} {iconOnly} {args} />
 
 			<!-- <Popover
 				notClickable
diff --git a/frontend/src/lib/components/ModulePreview.svelte b/frontend/src/lib/components/ModulePreview.svelte
index fa4d17a37c..c67d6c113e 100644
--- a/frontend/src/lib/components/ModulePreview.svelte
+++ b/frontend/src/lib/components/ModulePreview.svelte
@@ -135,12 +135,13 @@
 							jobId={testJob?.id}
 							result={testJob.result}>
 							<svelte:fragment slot="copilot-fix">
-								{#if lang && editor && diffEditor && testJob?.result?.error}
+								{#if lang && editor && diffEditor && stepArgs && testJob?.result?.error}
 									<ScriptFix
 										error={JSON.stringify(testJob.result.error)}
 										{lang}
 										{editor}
 										{diffEditor}
+										args={stepArgs}
 									/>
 								{/if}
 							</svelte:fragment>
diff --git a/frontend/src/lib/components/ScriptEditor.svelte b/frontend/src/lib/components/ScriptEditor.svelte
index ca9fe4d1ca..fe195b8dea 100644
--- a/frontend/src/lib/components/ScriptEditor.svelte
+++ b/frontend/src/lib/components/ScriptEditor.svelte
@@ -249,6 +249,7 @@
 			kind={asKind(kind)}
 			{template}
 			{diffEditor}
+			{args}
 		/>
 		{#if !noSyncFromGithub}
 			<div class="py-1">
@@ -302,6 +303,7 @@
 						deno={lang == 'deno'}
 						automaticLayout={true}
 						{fixedOverflowWidgets}
+						{args}
 					/>
 					<DiffEditor
 						bind:this={diffEditor}
@@ -364,6 +366,7 @@
 							previewIsLoading={testIsLoading}
 							{editor}
 							{diffEditor}
+							{args}
 						/>
 					</Pane>
 				</Splitpanes>
diff --git a/frontend/src/lib/components/apps/editor/inlineScriptsPanel/InlineScriptEditor.svelte b/frontend/src/lib/components/apps/editor/inlineScriptsPanel/InlineScriptEditor.svelte
index bc9e8207b5..91fa1c8276 100644
--- a/frontend/src/lib/components/apps/editor/inlineScriptsPanel/InlineScriptEditor.svelte
+++ b/frontend/src/lib/components/apps/editor/inlineScriptsPanel/InlineScriptEditor.svelte
@@ -195,6 +195,10 @@
 					editor={inlineScript?.language === 'frontend' ? simpleEditor : editor}
 					{diffEditor}
 					inlineScript
+					args={Object.entries(fields).reduce((acc, [key, obj]) => {
+						acc[key] = obj.type === 'static' ? obj.value : undefined
+						return acc
+					}, {})}
 				/>
 
 				<Button
@@ -281,6 +285,10 @@
 							}
 							$app = $app
 						}}
+						args={Object.entries(fields).reduce((acc, [key, obj]) => {
+							acc[key] = obj.type === 'static' ? obj.value : undefined
+							return acc
+						}, {})}
 					/>
 				{:else}
 					<SimpleEditor
diff --git a/frontend/src/lib/components/copilot/FlowCopilotStatus.svelte b/frontend/src/lib/components/copilot/FlowCopilotStatus.svelte
index fdd70b085c..8e8ee90df6 100644
--- a/frontend/src/lib/components/copilot/FlowCopilotStatus.svelte
+++ b/frontend/src/lib/components/copilot/FlowCopilotStatus.svelte
@@ -27,7 +27,7 @@
 	<ManualPopover bind:this={copilotPopover}>
 		<Button
 			size="xs"
-			btnClasses="mr-2 z-[901]"
+			btnClasses={'mr-2 ' + ($currentStepStore !== undefined ? 'z-[901]' : '')}
 			on:click={() => {
 				if (copilotLoading || ($currentStepStore !== undefined && $currentStepStore !== 'Input')) {
 					abortController?.abort()
diff --git a/frontend/src/lib/components/copilot/ScriptFix.svelte b/frontend/src/lib/components/copilot/ScriptFix.svelte
index 758f69454e..c7fbc84f02 100644
--- a/frontend/src/lib/components/copilot/ScriptFix.svelte
+++ b/frontend/src/lib/components/copilot/ScriptFix.svelte
@@ -23,6 +23,7 @@
 	export let editor: Editor | undefined
 	export let diffEditor: DiffEditor | undefined
 	export let error: string
+	export let args: Record<string, any>
 
 	// state
 	let genLoading: boolean = false
@@ -103,7 +104,7 @@
 
 	$: !$generatedCode && hideDiff()
 
-	$: dbSchema = $dbSchemas[Object.keys($dbSchemas)[0]]
+	$: dbSchema = $dbSchemas[(lang === 'graphql' ? args.api : args.database)?.replace('$res:', '')]
 </script>
 
 {#if SUPPORTED_LANGUAGES.has(lang)}
diff --git a/frontend/src/lib/components/copilot/ScriptGen.svelte b/frontend/src/lib/components/copilot/ScriptGen.svelte
index 8cef021322..3cf19804fd 100644
--- a/frontend/src/lib/components/copilot/ScriptGen.svelte
+++ b/frontend/src/lib/components/copilot/ScriptGen.svelte
@@ -34,6 +34,7 @@
 	export let editor: Editor | SimpleEditor | undefined
 	export let diffEditor: DiffEditor | undefined
 	export let inlineScript = false
+	export let args: Record<string, any>
 
 	// state
 	let funcDesc: string = ''
@@ -143,7 +144,8 @@
 	$: !$generatedCode && hideDiff()
 	$: editor && setSelectionHandler()
 	$: selection && (isEdit = !selection.isEmpty())
-	$: dbSchema = $dbSchemas[Object.keys($dbSchemas)[0]]
+
+	$: dbSchema = $dbSchemas[(lang === 'graphql' ? args.api : args.database)?.replace('$res:', '')]
 </script>
 
 {#if $generatedCode.length > 0 && !genLoading}
@@ -297,7 +299,7 @@
 								In order to better generate the script, we pass the selected DB schema to GPT-4.
 							</Tooltip>
 						</p>
-						{#if dbSchema.lang === 'postgresql'}
+						{#if dbSchema.lang !== 'graphql' && (dbSchema.schema?.public || dbSchema.schema?.PUBLIC)}
 							<ToggleButtonGroup class="w-auto shrink-0" bind:selected={dbSchema.publicOnly}>
 								<ToggleButton value={true} label="Public schema" />
 								<ToggleButton value={false} label="All schemas" />
diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index bcbb478214..dcf65711a3 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -210,6 +210,10 @@
 						iconOnly={width < 850}
 						kind={scriptKind}
 						template={scriptTemplate}
+						args={Object.entries(flowModule.value.input_transforms).reduce((acc, [key, obj]) => {
+							acc[key] = obj.type === 'static' ? obj.value : undefined
+							return acc
+						}, {})}
 					/>
 				</div>
 			{/if}
@@ -254,6 +258,13 @@
 										saveDraft()
 									}}
 									fixedOverflowWidgets={true}
+									args={Object.entries(flowModule.value.input_transforms).reduce(
+										(acc, [key, obj]) => {
+											acc[key] = obj.type === 'static' ? obj.value : undefined
+											return acc
+										},
+										{}
+									)}
 								/>
 								<DiffEditor
 									bind:this={diffEditor}
diff --git a/frontend/src/lib/components/scriptEditor/LogPanel.svelte b/frontend/src/lib/components/scriptEditor/LogPanel.svelte
index 482da39657..2e5d11df4a 100644
--- a/frontend/src/lib/components/scriptEditor/LogPanel.svelte
+++ b/frontend/src/lib/components/scriptEditor/LogPanel.svelte
@@ -29,6 +29,7 @@
 	export let pastPreviews: CompletedJob[] = []
 	export let editor: Editor | undefined = undefined
 	export let diffEditor: DiffEditor | undefined = undefined
+	export let args: Record<string, any> | undefined = undefined
 
 	type DrawerContent = {
 		mode: 'json' | Preview.language | 'plain'
@@ -99,12 +100,13 @@
 									result={previewJob.result}
 								>
 									<svelte:fragment slot="copilot-fix">
-										{#if lang && editor && diffEditor && previewJob?.result?.error}
+										{#if lang && editor && diffEditor && args && previewJob?.result?.error}
 											<ScriptFix
 												error={JSON.stringify(previewJob.result.error)}
 												{lang}
 												{editor}
 												{diffEditor}
+												{args}
 											/>
 										{/if}
 									</svelte:fragment>
@@ -137,7 +139,9 @@
 					{#each pastPreviews as { id, created_at, success, result }}
 						<tr class="">
 							<td class="text-xs">
-								<a class="pr-3" href="/run/{id}?workspace={$workspaceStore}" target="_blank">{id.substring(30)}</a>
+								<a class="pr-3" href="/run/{id}?workspace={$workspaceStore}" target="_blank"
+									>{id.substring(30)}</a
+								>
 							</td>
 							<td class="text-xs">{displayDate(created_at)}</td>
 							<td class="text-xs">
diff --git a/frontend/src/lib/stores.ts b/frontend/src/lib/stores.ts
index a6b9fc2175..800064f21a 100644
--- a/frontend/src/lib/stores.ts
+++ b/frontend/src/lib/stores.ts
@@ -82,14 +82,9 @@ type SQLBaseSchema = {
 }
 
 export interface SQLSchema {
-	lang: 'mysql' | 'bigquery'
+	lang: 'mysql' | 'bigquery' | 'postgresql' | 'snowflake'
 	schema: SQLBaseSchema
-}
-
-export interface PostgresqlSchema {
-	lang: 'postgresql'
-	schema: SQLBaseSchema
-	publicOnly: boolean
+	publicOnly: boolean | undefined
 }
 
 export interface GraphqlSchema {
@@ -97,7 +92,7 @@ export interface GraphqlSchema {
 	schema: IntrospectionQuery
 }
 
-export type DBSchema = SQLSchema | PostgresqlSchema | GraphqlSchema
+export type DBSchema = SQLSchema | GraphqlSchema
 
 interface DBSchemas {
 	[resourcePath: string]: DBSchema