diff --git a/package.json b/package.json index 23224b214a..05f0d7ebb1 100644 --- a/package.json +++ b/package.json @@ -10,6 +10,7 @@ "tauri": "tauri" }, "dependencies": { + "@codemirror/autocomplete": "^6.20.1", "@codemirror/lang-sql": "^6.10.0", "@codemirror/state": "^6.6.0", "@codemirror/theme-one-dark": "^6.1.3", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 0a82fe0e4d..b0209c1370 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -8,6 +8,9 @@ importers: .: dependencies: + '@codemirror/autocomplete': + specifier: ^6.20.1 + version: 6.20.1 '@codemirror/lang-sql': specifier: ^6.10.0 version: 6.10.0 diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 220cecf8db..adfbb78d2b 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -1111,7 +1111,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "dbx" -version = "0.2.0" +version = "0.2.1" dependencies = [ "anyhow", "chrono", diff --git a/src/App.vue b/src/App.vue index 26ea24c132..8fb94be70a 100644 --- a/src/App.vue +++ b/src/App.vue @@ -853,6 +853,8 @@ async function setupFileDrop() { import { ref, onMounted, onBeforeUnmount, watch, shallowRef } from "vue"; +import type { CompletionContext } from "@codemirror/autocomplete"; import type { EditorView as EditorViewType } from "@codemirror/view"; import { resolveExecutableSql } from "@/lib/sqlExecutionTarget"; import { formatSqlText, type SqlFormatDialect } from "@/lib/sqlFormatter"; +import { useConnectionStore } from "@/stores/connectionStore"; +import { + buildSqlCompletionItemsFromContext, + getSqlCompletionContext, +} from "@/lib/sqlCompletion"; const props = defineProps<{ modelValue: string; + connectionId?: string; + database?: string; dialect?: "mysql" | "postgres"; formatDialect?: SqlFormatDialect; formatRequestId?: number; @@ -20,6 +28,7 @@ const emit = defineEmits<{ const editorRef = ref(); const view = shallowRef(null); +const connectionStore = useConnectionStore(); const DEFAULT_FONT_SIZE = 13; const MIN_FONT_SIZE = 10; const MAX_FONT_SIZE = 24; @@ -98,6 +107,53 @@ async function formatCurrentSql() { } } +async function provideSqlCompletions( + currentState: import("@codemirror/state").EditorState, + position: number, +) { + if (!props.connectionId || !props.database) return null; + + const completionContext = getSqlCompletionContext(currentState.doc.toString(), position); + const tables = await connectionStore.listCompletionTables(props.connectionId, props.database); + const columnsByTable = new Map>>(); + + if (completionContext.suggestColumns) { + const relatedTables = completionContext.qualifier + ? completionContext.referencedTables.filter((table) => table.alias === completionContext.qualifier || table.name === completionContext.qualifier) + : completionContext.referencedTables; + + await Promise.all(relatedTables.map(async (table) => { + const cacheKey = table.schema ? `${table.schema}.${table.name}` : table.name; + if (columnsByTable.has(cacheKey)) return; + const columns = await connectionStore.listCompletionColumns( + props.connectionId!, + props.database!, + table.name, + table.schema, + ); + columnsByTable.set(cacheKey, columns); + })); + } + + const items = buildSqlCompletionItemsFromContext(completionContext, { + tables, + columnsByTable, + }); + + if (items.length === 0) return null; + + return { + from: position - completionContext.prefix.length, + options: items.map((item) => ({ + label: item.label, + type: item.type === "keyword" ? "keyword" : item.type === "table" ? "class" : "property", + detail: item.detail, + boost: item.boost, + })), + validFor: /^[\w$]*$/, + }; +} + onMounted(async () => { if (!editorRef.value) return; @@ -107,12 +163,14 @@ onMounted(async () => { { sql, MySQL, PostgreSQL }, { basicSetup }, { oneDark }, + { autocompletion, startCompletion }, ] = await Promise.all([ import("@codemirror/view"), import("@codemirror/state"), import("@codemirror/lang-sql"), import("codemirror"), import("@codemirror/theme-one-dark"), + import("@codemirror/autocomplete"), ]); editorViewModule = { EditorView, keymap } as typeof import("@codemirror/view"); fontSizeTheme = new Compartment(); @@ -162,11 +220,24 @@ onMounted(async () => { extensions: [ basicSetup, sql({ dialect }), + autocompletion({ + activateOnTyping: true, + override: [ + async (context: CompletionContext) => provideSqlCompletions(context.state, context.pos), + ], + }), oneDark, runKeymap, EditorView.updateListener.of((update) => { if (update.docChanged) { emit("update:modelValue", update.state.doc.toString()); + let insertedText = ""; + update.changes.iterChanges((_fromA, _toA, _fromB, _toB, inserted) => { + insertedText += inserted.toString(); + }); + if (insertedText.endsWith(".")) { + startCompletion(update.view); + } } if (update.selectionSet || update.docChanged) { emit("selectionChange", selectedSqlFromView(update.view)); diff --git a/src/lib/sqlCompletion.ts b/src/lib/sqlCompletion.ts new file mode 100644 index 0000000000..90c2cd1aea --- /dev/null +++ b/src/lib/sqlCompletion.ts @@ -0,0 +1,194 @@ +const SQL_KEYWORDS = [ + "SELECT", "FROM", "WHERE", "JOIN", "LEFT", "RIGHT", "INNER", "OUTER", + "ON", "GROUP BY", "ORDER BY", "HAVING", "LIMIT", "OFFSET", "INSERT", + "INTO", "VALUES", "UPDATE", "SET", "DELETE", "CREATE", "TABLE", "VIEW", + "AS", "AND", "OR", "NOT", "IN", "IS", "NULL", "LIKE", "DISTINCT", +]; + +const TABLE_TRIGGER_KEYWORDS = new Set(["from", "join", "update", "into"]); + +export interface SqlCompletionTable { + name: string; + schema?: string; + type?: "table" | "view"; +} + +export interface SqlCompletionColumn { + name: string; + table: string; + schema?: string; + dataType?: string; +} + +export interface SqlCompletionItem { + label: string; + type: "keyword" | "table" | "column"; + detail?: string; + boost: number; +} + +export interface SqlCompletionReferencedTable { + name: string; + schema?: string; + alias?: string; +} + +export interface SqlCompletionContext { + prefix: string; + qualifier?: string; + suggestTables: boolean; + suggestColumns: boolean; + referencedTables: SqlCompletionReferencedTable[]; +} + +export function buildSqlCompletionItems( + sql: string, + cursor: number, + input: { + tables: SqlCompletionTable[]; + columnsByTable: Map; + }, +): SqlCompletionItem[] { + const context = getSqlCompletionContext(sql, cursor); + return buildSqlCompletionItemsFromContext(context, input); +} + +export function buildSqlCompletionItemsFromContext( + context: SqlCompletionContext, + input: { + tables: SqlCompletionTable[]; + columnsByTable: Map; + }, +): SqlCompletionItem[] { + const items: SqlCompletionItem[] = []; + + if (context.suggestColumns) { + items.push(...buildColumnItems(context, input.columnsByTable)); + } + + if (context.suggestTables) { + items.push(...buildTableItems(context.prefix, input.tables)); + } + + if (!context.qualifier) { + items.push(...buildKeywordItems(context.prefix)); + } + + return dedupeAndSort(items); +} + +export function getSqlCompletionContext(sql: string, cursor: number): SqlCompletionContext { + const beforeCursor = sql.slice(0, cursor); + const dottedMatch = /([A-Za-z_][\w$]*)\.([A-Za-z_][\w$]*)?$/.exec(beforeCursor); + const plainMatch = /([A-Za-z_][\w$]*)$/.exec(beforeCursor); + const prefix = dottedMatch?.[2] ?? plainMatch?.[1] ?? ""; + const qualifier = dottedMatch?.[1]; + const bareStart = qualifier + ? cursor - prefix.length + : cursor - (plainMatch?.[1]?.length ?? 0); + const beforeToken = beforeCursor.slice(0, Math.max(0, bareStart)).trimEnd(); + const lastWord = /([A-Za-z_][\w$]*)$/.exec(beforeToken)?.[1]?.toLowerCase() ?? ""; + const referencedTables = extractReferencedTables(sql); + + return { + prefix, + qualifier, + suggestTables: TABLE_TRIGGER_KEYWORDS.has(lastWord), + suggestColumns: !!qualifier || referencedTables.length > 0, + referencedTables, + }; +} + +function extractReferencedTables(sql: string): SqlCompletionReferencedTable[] { + const pattern = /\b(?:from|join|update|into)\s+((?:"[^"]+"|`[^`]+`|[A-Za-z_][\w$]*)(?:\.(?:"[^"]+"|`[^`]+`|[A-Za-z_][\w$]*))?)(?:\s+(?:as\s+)?([A-Za-z_][\w$]*))?/gi; + const referenced: SqlCompletionReferencedTable[] = []; + for (const match of sql.matchAll(pattern)) { + const rawName = match[1]; + const alias = match[2]; + const [first, second] = splitQualifiedName(rawName); + if (!first) continue; + const table = second ? { schema: first, name: second, alias } : { name: first, alias }; + referenced.push(table); + } + return referenced; +} + +function splitQualifiedName(input: string): [string | undefined, string | undefined] { + const parts = input.split(".").map((part) => unquoteIdentifier(part.trim())).filter(Boolean); + if (parts.length >= 2) return [parts[0], parts[1]]; + return [parts[0], undefined]; +} + +function unquoteIdentifier(value: string): string { + if ((value.startsWith("\"") && value.endsWith("\"")) || (value.startsWith("`") && value.endsWith("`"))) { + return value.slice(1, -1); + } + return value; +} + +function buildTableItems(prefix: string, tables: SqlCompletionTable[]): SqlCompletionItem[] { + return tables + .filter((table) => matchesPrefix(table.name, prefix)) + .map((table) => ({ + label: table.name, + type: "table" as const, + detail: table.schema ? `${table.schema}.${table.name}` : table.type, + boost: computeBoost(table.name, prefix), + })); +} + +function buildColumnItems( + context: SqlCompletionContext, + columnsByTable: Map, +): SqlCompletionItem[] { + const relatedTables = context.qualifier + ? context.referencedTables.filter((table) => table.alias === context.qualifier || table.name === context.qualifier) + : context.referencedTables; + + const columns = relatedTables.flatMap((table) => { + const key = table.schema ? `${table.schema}.${table.name}` : table.name; + return columnsByTable.get(key) ?? []; + }); + + return columns + .filter((column) => matchesPrefix(column.name, context.prefix)) + .map((column) => ({ + label: column.name, + type: "column" as const, + detail: column.schema ? `${column.schema}.${column.table}` : column.table, + boost: computeBoost(column.name, context.prefix), + })); +} + +function buildKeywordItems(prefix: string): SqlCompletionItem[] { + return SQL_KEYWORDS + .filter((keyword) => matchesPrefix(keyword, prefix)) + .map((keyword) => ({ + label: keyword, + type: "keyword" as const, + boost: computeBoost(keyword, prefix), + })); +} + +function matchesPrefix(candidate: string, prefix: string): boolean { + if (!prefix) return true; + return candidate.toLowerCase().includes(prefix.toLowerCase()); +} + +function computeBoost(candidate: string, prefix: string): number { + if (!prefix) return 1; + const startsWith = candidate.toLowerCase().startsWith(prefix.toLowerCase()); + return (startsWith ? 1000 : 100) - candidate.length; +} + +function dedupeAndSort(items: SqlCompletionItem[]): SqlCompletionItem[] { + const seen = new Set(); + return items + .sort((left, right) => right.boost - left.boost) + .filter((item) => { + const key = `${item.type}:${item.label}`; + if (seen.has(key)) return false; + seen.add(key); + return true; + }); +} diff --git a/src/stores/connectionStore.ts b/src/stores/connectionStore.ts index 9aa369c446..d560ff7a8a 100644 --- a/src/stores/connectionStore.ts +++ b/src/stores/connectionStore.ts @@ -1,7 +1,8 @@ import { defineStore } from "pinia"; import { ref } from "vue"; -import type { ConnectionConfig, TreeNode } from "@/types/database"; +import type { ColumnInfo, ConnectionConfig, TreeNode } from "@/types/database"; import { orderPinnedFirst } from "@/lib/pinnedItems"; +import type { SqlCompletionColumn, SqlCompletionTable } from "@/lib/sqlCompletion"; import * as api from "@/lib/tauri"; const PINNED_TREE_NODES_STORAGE_KEY = "dbx-pinned-tree-nodes"; @@ -13,6 +14,8 @@ export const useConnectionStore = defineStore("connection", () => { const pinnedTreeNodeIds = ref>(loadPinnedTreeNodeIds()); const connectedIds = ref>(new Set()); const editingConnectionId = ref(null); + const completionTablesCache = ref>({}); + const completionColumnsCache = ref>({}); function startEditing(id: string) { editingConnectionId.value = id; @@ -133,12 +136,23 @@ export const useConnectionStore = defineStore("connection", () => { persistConnections(); } + function invalidateCompletionCache(connectionId: string) { + const cachePrefix = `${connectionId}:`; + completionTablesCache.value = Object.fromEntries( + Object.entries(completionTablesCache.value).filter(([key]) => !key.startsWith(cachePrefix)), + ); + completionColumnsCache.value = Object.fromEntries( + Object.entries(completionColumnsCache.value).filter(([key]) => !key.startsWith(cachePrefix)), + ); + } + function removeConnection(id: string) { connections.value = connections.value.filter((c) => c.id !== id); treeNodes.value = treeNodes.value.filter((n) => n.id !== id); if (activeConnectionId.value === id) { activeConnectionId.value = null; } + invalidateCompletionCache(id); persistConnections(); } @@ -153,6 +167,7 @@ export const useConnectionStore = defineStore("connection", () => { node.children = []; } connectedIds.value.delete(config.id); + invalidateCompletionCache(config.id); persistConnections(); } @@ -197,6 +212,7 @@ export const useConnectionStore = defineStore("connection", () => { if (activeConnectionId.value === connectionId) { activeConnectionId.value = null; } + invalidateCompletionCache(connectionId); } async function ensureConnected(connectionId: string) { @@ -470,6 +486,64 @@ export const useConnectionStore = defineStore("connection", () => { } } + function isSchemaAwareDatabase(connectionId: string): boolean { + const dbType = getConfig(connectionId)?.db_type; + return dbType === "postgres" || dbType === "sqlserver" || dbType === "oracle"; + } + + async function listCompletionTables(connectionId: string, database: string): Promise { + const cacheKey = `${connectionId}:${database}`; + if (completionTablesCache.value[cacheKey]) { + return completionTablesCache.value[cacheKey]; + } + + await ensureConnected(connectionId); + + if (isSchemaAwareDatabase(connectionId)) { + const schemas = await api.listSchemas(connectionId, database); + const tableGroups = await Promise.all( + schemas.map(async (schema) => { + const tables = await api.listTables(connectionId, database, schema); + return tables.map((table) => ({ + name: table.name, + schema, + type: table.table_type === "VIEW" ? "view" as const : "table" as const, + })); + }), + ); + completionTablesCache.value[cacheKey] = tableGroups.flat(); + return completionTablesCache.value[cacheKey]; + } + + const tables = await api.listTables(connectionId, database, database); + completionTablesCache.value[cacheKey] = tables.map((table) => ({ + name: table.name, + type: table.table_type === "VIEW" ? "view" as const : "table" as const, + })); + return completionTablesCache.value[cacheKey]; + } + + async function listCompletionColumns( + connectionId: string, + database: string, + table: string, + schema?: string, + ): Promise { + const cacheKey = `${connectionId}:${database}:${schema || ""}:${table}`; + if (!completionColumnsCache.value[cacheKey]) { + await ensureConnected(connectionId); + const querySchema = schema || database; + completionColumnsCache.value[cacheKey] = await api.getColumns(connectionId, database, querySchema, table); + } + + return completionColumnsCache.value[cacheKey].map((column) => ({ + name: column.name, + table, + schema, + dataType: column.data_type, + })); + } + function findNode(nodes: TreeNode[], id: string): TreeNode | null { for (const node of nodes) { if (node.id === id) return node; @@ -552,6 +626,8 @@ export const useConnectionStore = defineStore("connection", () => { loadIndexes, loadForeignKeys, loadTriggers, + listCompletionTables, + listCompletionColumns, exportConnectionsToFile, importConnectionsFromFile, }; diff --git a/tests/sqlCompletion.test.ts b/tests/sqlCompletion.test.ts new file mode 100644 index 0000000000..139d6aa1de --- /dev/null +++ b/tests/sqlCompletion.test.ts @@ -0,0 +1,76 @@ +import { strict as assert } from "node:assert"; +import test from "node:test"; +import { + buildSqlCompletionItems, + type SqlCompletionColumn, + type SqlCompletionTable, +} from "../src/lib/sqlCompletion.ts"; + +const tables: SqlCompletionTable[] = [ + { name: "users", schema: "public", type: "table" }, + { name: "user_profiles", schema: "public", type: "table" }, + { name: "orders", schema: "public", type: "table" }, +]; + +const columnsByTable = new Map([ + ["public.users", [ + { name: "id", table: "users", schema: "public", dataType: "bigint" }, + { name: "name", table: "users", schema: "public", dataType: "varchar" }, + { name: "email", table: "users", schema: "public", dataType: "varchar" }, + ]], + ["public.orders", [ + { name: "id", table: "orders", schema: "public", dataType: "bigint" }, + { name: "user_id", table: "orders", schema: "public", dataType: "bigint" }, + { name: "status", table: "orders", schema: "public", dataType: "varchar" }, + ]], +]); + +test("suggests SQL keywords for generic keyword input", () => { + const items = buildSqlCompletionItems("sel", 3, { + tables, + columnsByTable, + }); + + assert.equal(items[0]?.label, "SELECT"); + assert.equal(items[0]?.type, "keyword"); +}); + +test("suggests matching table names after FROM", () => { + const sql = "select * from us"; + const items = buildSqlCompletionItems(sql, sql.length, { + tables, + columnsByTable, + }); + + assert.deepEqual( + items.slice(0, 2).map((item) => item.label), + ["users", "user_profiles"], + ); +}); + +test("suggests columns for an explicit alias qualifier", () => { + const sql = "select u. from public.users u"; + const cursor = "select u.".length; + const items = buildSqlCompletionItems(sql, cursor, { + tables, + columnsByTable, + }); + + assert.deepEqual( + items.map((item) => item.label), + ["id", "name", "email"], + ); + assert.ok(items.every((item) => item.type === "column")); +}); + +test("suggests columns from referenced tables in select list", () => { + const sql = "select na from public.users u join public.orders o on u.id = o.user_id"; + const cursor = "select na".length; + const items = buildSqlCompletionItems(sql, cursor, { + tables, + columnsByTable, + }); + + assert.equal(items[0]?.label, "name"); + assert.equal(items[0]?.type, "column"); +});