feat: table & column completions in sql editor

frontend/src2/components/Code.vue

@@ -20,9 +20,8 @@ import { autocompletion, closeBrackets } from '@codemirror/autocomplete'
 import { javascript } from '@codemirror/lang-javascript'
 import { python } from '@codemirror/lang-python'
 import { MySQL, sql } from '@codemirror/lang-sql'
-import { HighlightStyle, syntaxHighlighting, syntaxTree } from '@codemirror/language'
+import { syntaxTree } from '@codemirror/language'
 import { EditorView } from '@codemirror/view'
-import { tags } from '@lezer/highlight'
 import { onMounted, ref, watch } from 'vue'
 import { Codemirror } from 'vue-codemirror'
 import { tomorrow } from 'thememirror'

frontend/src2/data_source/data_source.ts

@@ -63,6 +63,10 @@ export function getDataSourceOptions() {
 	})
 }
 
+function getSchema(data_source: string) {
+	return call('insights.api.data_sources.get_schema', { data_source })
+}
+
 export default function useDataSourceStore() {
 	if (!sources.value.length) {
 		fetchSources()
@@ -75,6 +79,8 @@ export default function useDataSourceStore() {
 		getSource,
 		getOptions: getDataSourceOptions,
 
+		getSchema,
+
 		testing,
 		testConnection,
 

frontend/src2/query/Query.vue

@@ -1,27 +1,15 @@
 <script setup lang="ts">
-import { useMagicKeys, whenever } from '@vueuse/core'
-import { onBeforeUnmount, provide } from 'vue'
+import { provide } from 'vue'
 import { WorkbookQuery } from '../types/workbook.types'
 import NativeQueryEditor from './components/NativeQueryEditor.vue'
-import useQuery from './query'
 import QueryBuilder from './components/QueryBuilder.vue'
+import useQuery from './query'
 
 const props = defineProps<{ query: WorkbookQuery }>()
 const query = useQuery(props.query)
 provide('query', query)
 window.query = query
 query.execute()
-
-const keys = useMagicKeys()
-const cmdZ = keys['Meta+Z']
-const cmdShiftZ = keys['Meta+Shift+Z']
-const stopUndoWatcher = whenever(cmdZ, () => query.canUndo() && query.history.undo())
-const stopRedoWatcher = whenever(cmdShiftZ, () => query.canRedo() && query.history.redo())
-
-onBeforeUnmount(() => {
-	stopUndoWatcher()
-	stopRedoWatcher()
-})
 </script>
 
 <template>

frontend/src2/query/components/NativeQueryEditor.vue

@@ -2,12 +2,14 @@
 import { useTimeAgo } from '@vueuse/core'
 import { LoadingIndicator } from 'frappe-ui'
 import { Play, RefreshCw, Wand2 } from 'lucide-vue-next'
-import { computed, inject, ref } from 'vue'
+import { computed, inject, ref, watchEffect } from 'vue'
 import Code from '../../components/Code.vue'
 import DataTable from '../../components/DataTable.vue'
 import { Query } from '../query'
 import ContentEditable from '../../components/ContentEditable.vue'
 import DataSourceSelector from './source_selector/DataSourceSelector.vue'
+import { wheneverChanges } from '../../helpers'
+import useDataSourceStore from '../../data_source/data_source'
 
 const query = inject<Query>('query')!
 
@@ -24,6 +26,47 @@ const previewRowCount = computed(() => query.result.rows.length.toLocaleString()
 const totalRowCount = computed(() =>
 	query.result.totalRowCount ? query.result.totalRowCount.toLocaleString() : ''
 )
+
+const dataSourceSchema = ref<Record<string, any>>({})
+const dataSourceStore = useDataSourceStore()
+wheneverChanges(
+	data_source,
+	() => {
+		if (!data_source.value) {
+			dataSourceSchema.value = {}
+			return
+		}
+		dataSourceStore.getSchema(data_source.value).then((schema: any) => {
+			dataSourceSchema.value = schema
+		})
+	},
+	{ immediate: true }
+)
+const completions = computed(() => {
+	if (!Object.keys(dataSourceSchema.value).length)
+		return {
+			schema: {},
+			tables: [],
+		}
+
+	const schema: Record<string, any> = {}
+	Object.entries(dataSourceSchema.value).forEach(([table, tableData]) => {
+		schema[table] = tableData.columns.map((column: any) => ({
+			label: column.label,
+			detail: column.label,
+		}))
+	})
+
+	const tables = Object.entries(dataSourceSchema.value).map(([table, tableData]) => ({
+		label: table,
+		detail: tableData.label,
+	}))
+
+	return {
+		schema,
+		tables,
+	}
+})
 </script>
 
 <template>
@@ -38,7 +81,13 @@ const totalRowCount = computed(() =>
 				></ContentEditable>
 			</div>
 			<div class="flex-1 overflow-hidden">
-				<Code v-model="sql" language="sql" />
+				<Code
+					:key="completions.tables.length"
+					v-model="sql"
+					language="sql"
+					:schema="completions.schema"
+					:tables="completions.tables"
+				/>
 			</div>
 			<div class="flex flex-shrink-0 gap-1 border-t p-1">
 				<Button @click="execute" label="Execute">

frontend/src2/query/components/QueryBuilder.vue

@@ -1,13 +1,26 @@
 <script setup lang="ts">
-import { inject } from 'vue'
+import { inject, onBeforeUnmount } from 'vue'
 import { Query } from '../query'
 import QueryBuilderSourceSelector from './QueryBuilderSourceSelector.vue'
 import QueryBuilderTable from './QueryBuilderTable.vue'
 import QueryBuilderToolbar from './QueryBuilderToolbar.vue'
 import QueryInfo from './QueryInfo.vue'
 import QueryOperations from './QueryOperations.vue'
+import { useMagicKeys } from '@vueuse/core'
+import { whenever } from '@vueuse/core'
 
 const query = inject<Query>('query')!
+
+const keys = useMagicKeys()
+const cmdZ = keys['Meta+Z']
+const cmdShiftZ = keys['Meta+Shift+Z']
+const stopUndoWatcher = whenever(cmdZ, () => query.canUndo() && query.history.undo())
+const stopRedoWatcher = whenever(cmdShiftZ, () => query.canRedo() && query.history.redo())
+
+onBeforeUnmount(() => {
+	stopUndoWatcher()
+	stopRedoWatcher()
+})
 </script>
 
 <template>

insights/api/data_sources.py

@@ -448,3 +448,35 @@ def get_data_sources_of_tables(table_names: list[str]):
         data_sources[table.data_source].append(table.name)
 
     return data_sources
+
+
+@insights_whitelist()
+@site_cache(ttl=24 * 60 * 60)
+@validate_type
+def get_schema(data_source: str):
+    check_data_source_permission(data_source)
+    ds = frappe.get_doc("Insights Data Source v3", data_source)
+    db = ds._get_ibis_backend()
+
+    tables = get_data_source_tables(data_source)
+    schema = {}
+
+    for table in tables:
+        table_name = table.table_name
+        schema[table_name] = {
+            "table": table_name,
+            "label": table.label,
+            "data_source": data_source,
+            "columns": [],
+        }
+        _table = db.table(table_name)
+        for column, datatype in _table.schema().items():
+            schema[table_name]["columns"].append(
+                frappe._dict(
+                    column=column,
+                    label=column,
+                    type=to_insights_type(datatype),
+                )
+            )
+
+    return schema