feat: add test connection for bigquery (#1988)

Co-authored-by: Ruben Fiszel <ruben@windmill.dev>

frontend/src/lib/components/DBSchemaExplorer.svelte

@@ -6,6 +6,7 @@
 	import Drawer from './common/drawer/Drawer.svelte'
 	import DrawerContent from './common/drawer/DrawerContent.svelte'
 	import ObjectViewer from './propertyPicker/ObjectViewer.svelte'
+	import { tryEvery } from '$lib/utils'
 	import ToggleButton from './common/toggleButton-v2/ToggleButton.svelte'
 	import ToggleButtonGroup from './common/toggleButton-v2/ToggleButtonGroup.svelte'
 
@@ -120,41 +121,36 @@ export async function main(args: any) {
 				}
 			}
 		})
-		let i = 1
-		const inter = setInterval(async () => {
-			try {
+
+		tryEvery({
+			tryCode: async () => {
 				const testResult = await JobService.getCompletedJob({
 					workspace: $workspaceStore!,
 					id: job
 				})
-				if (testResult) {
-					if (!testResult.success) {
-						console.error(testResult.result?.['error']?.['message'])
-					} else {
-						dbSchema.set(testResult.result)
-					}
-					clearInterval(inter)
+				if (!testResult.success) {
+					console.error(testResult.result?.['error']?.['message'])
+				} else {
+					dbSchema.set(testResult.result)
 				}
-			} catch (err) {
-				if (i >= 5) {
-					console.error('Could not query DB schema within 5s')
-					clearInterval(inter)
-					try {
-						await JobService.cancelQueuedJob({
-							workspace: $workspaceStore!,
-							id: job,
-							requestBody: {
-								reason: 'Could not query DB schema within 5s'
-							}
-						})
-					} catch (err) {
-						console.error(err)
-					}
+			},
+			timeoutCode: async () => {
+				console.error('Could not query DB schema within 5s')
+				try {
+					await JobService.cancelQueuedJob({
+						workspace: $workspaceStore!,
+						id: job,
+						requestBody: {
+							reason: 'Could not query DB schema within 5s'
+						}
+					})
+				} catch (err) {
+					console.error(err)
 				}
-			} finally {
-				i += 1
-			}
-		}, 1000)
+			},
+			interval: 500,
+			timeout: 5000
+		})
 	}
 
 	function formatSchema(

frontend/src/lib/components/TestConnection.svelte

@@ -5,6 +5,7 @@
 	import Button from './common/button/Button.svelte'
 	import { sendUserToast } from '$lib/toast'
 	import { workspaceStore } from '$lib/stores'
+	import { tryEvery } from '$lib/utils'
 
 	export let resource_type: string | undefined
 	export let args: Record<string, any> | any = {}
@@ -45,6 +46,10 @@ export async function main(database: any) {
 }`,
 			lang: 'deno'
 		},
+		bigquery: {
+			code: `select 1`,
+			lang: 'bigquery'
+		},
 		snowflake: {
 			code: `select 1`,
 			lang: 'snowflake'
@@ -68,33 +73,37 @@ export async function main(database: any) {
 				}
 			}
 		})
-		await new Promise((r) => setTimeout(r, 5000))
-		loading = false
-		try {
-			const testResult = await JobService.getCompletedJob({
-				workspace: $workspaceStore!,
-				id: job
-			})
-			if (testResult) {
+
+		tryEvery({
+			tryCode: async () => {
+				const testResult = await JobService.getCompletedJob({
+					workspace: $workspaceStore!,
+					id: job
+				})
+				loading = false
 				sendUserToast(
 					testResult.success ? 'Connection successful' : testResult.result?.['error']?.['message'],
 					!testResult.success
 				)
-			}
-		} catch (e) {
-			sendUserToast('Connection did not resolve after 5s', true)
-			try {
-				await JobService.cancelQueuedJob({
-					workspace: $workspaceStore!,
-					id: job,
-					requestBody: {
-						reason: 'Connection did not resolve after 5s'
-					}
-				})
-			} catch (err) {
-				console.error(err)
-			}
-		}
+			},
+			timeoutCode: async () => {
+				loading = false
+				sendUserToast('Connection did not resolve after 5s', true)
+				try {
+					await JobService.cancelQueuedJob({
+						workspace: $workspaceStore!,
+						id: job,
+						requestBody: {
+							reason: 'Connection did not resolve after 5s'
+						}
+					})
+				} catch (err) {
+					console.error(err)
+				}
+			},
+			interval: 500,
+			timeout: 5000
+		})
 	}
 </script>
 

frontend/src/lib/infer.ts

@@ -54,7 +54,7 @@ export async function inferArgs(
 		} else if (language == 'bigquery') {
 			inferedSchema = JSON.parse(parse_bigquery(code))
 			inferedSchema.args = [
-				{ name: 'database', typ: { resource: 'gcp_service_account' } },
+				{ name: 'database', typ: { resource: 'bigquery' } },
 				...inferedSchema.args
 			]
 		} else if (language == 'snowflake') {

frontend/src/lib/utils.ts

@@ -604,3 +604,30 @@ export function isCodeInjection(expr: string | undefined): boolean {
 
 	return dynamicTemplateRegex.test(expr)
 }
+
+export async function tryEvery({
+	tryCode,
+	timeoutCode,
+	interval,
+	timeout
+}: {
+	tryCode: () => Promise<any>
+	timeoutCode: () => void
+	interval: number
+	timeout: number
+}) {
+	const times = Math.floor(timeout / interval)
+
+	let i = 0
+	while (i < times) {
+		await sleep(interval)
+		try {
+			await tryCode()
+			break
+		} catch (err) {}
+		i++
+	}
+	if (i >= times) {
+		timeoutCode()
+	}
+}