fix: improve SQL table and JOIN completion (#2629)

apps/desktop/src/components/editor/QueryEditor.vue

@@ -1632,7 +1632,8 @@ async function provideSqlCompletions(context: CompletionContext) {
     if (!explicit && !shouldAutoOpenSqlCompletion(fullDoc, position)) return null;
 
     const legacyCompletionContext = getSqlCompletionContext(fullDoc, position);
-    const completionContext = SEMANTIC_SQL_COMPLETION_ENABLED ? sqlCompletionContextFromSemantic(buildSqlSemanticModel(fullDoc, position, { databaseType: props.databaseType, dialect: props.dialect }), legacyCompletionContext) : legacyCompletionContext;
+    const semanticModel = SEMANTIC_SQL_COMPLETION_ENABLED ? buildSqlSemanticModel(fullDoc, position, { databaseType: props.databaseType, dialect: props.dialect }) : null;
+    const completionContext = semanticModel ? sqlCompletionContextFromSemantic(semanticModel, legacyCompletionContext) : legacyCompletionContext;
 
     if (!hasDatabase) {
       const items = buildSqlCompletionItemsFromContext(completionContext, {
@@ -1674,8 +1675,9 @@ async function provideSqlCompletions(context: CompletionContext) {
       scheduleCompletionMetadataRefresh(completionContext);
       if (!explicit) return localResult;
     }
-    const shouldResolveAsyncColumnCompletion = completionContext.suggestColumns && completionContext.referencedTables.length > 0 && completionContext.prefix.length > 0;
-    if (!explicit && !shouldResolveAsyncColumnCompletion) {
+    const tableNameCompletion = isTableNameCompletionContext(completionContext);
+    const shouldResolveAsyncCompletion = tableNameCompletion || (completionContext.suggestColumns && completionContext.referencedTables.length > 0 && completionContext.prefix.length > 0);
+    if (!explicit && !shouldResolveAsyncCompletion) {
       scheduleCompletionMetadataRefresh(completionContext);
       return null;
     }
@@ -1729,6 +1731,22 @@ function flushImeComposition() {
   if (editorIsActive) emitEditorSelection(latestSelection);
 }
 
+function shouldStartSqlCompletionAfterInput(insertedText: string, currentView: EditorViewType): boolean {
+  const position = currentView.state.selection.main.head;
+  const fullDoc = currentView.state.doc.toString();
+  if (insertedText.endsWith(".")) return true;
+  if (/[,(]$/.test(insertedText)) {
+    const completionContext = getSqlCompletionContext(fullDoc, position);
+    return !!completionContext.insertTable;
+  }
+  if (/\s$/.test(insertedText)) {
+    return shouldAutoOpenSqlCompletion(fullDoc, position);
+  }
+  if (!/[\w$@]$/.test(insertedText)) return false;
+  const completionContext = getSqlCompletionContext(fullDoc, position);
+  return isTableNameCompletionContext(completionContext);
+}
+
 function buildLocalSqlCompletionResult(completionContext: ReturnType<typeof getSqlCompletionContext>, fullDoc: string, position: number) {
   if (!props.connectionId || props.database == null) return null;
   const databaseNames = supportsDatabaseQualifierCompletion() && completionContext.suggestTables && !completionContext.insertTable ? connectionStore.lookupLocalCompletionDatabases(props.connectionId, completionContext.qualifier || completionContext.prefix, MAX_COMPLETION_TABLES) : [];
@@ -1823,6 +1841,7 @@ function scheduleCompletionMetadataRefresh(completionContext: ReturnType<typeof
   if (!props.connectionId || props.database == null) return;
   const localOnlyMetadata = usesLocalOnlyCompletionMetadata();
   const onDemandOnlyColumns = usesOnDemandOnlyCompletionColumns();
+  const tableNameCompletion = isTableNameCompletionContext(completionContext);
   const connectionId = props.connectionId;
   const database = props.database;
   const schema = completionContext.qualifier && completionContext.suggestTables ? completionContext.qualifier : props.schema;
@@ -1874,7 +1893,7 @@ function scheduleCompletionMetadataRefresh(completionContext: ReturnType<typeof
         .catch(() => {});
     }
   }
-  if (!onDemandOnlyColumns) {
+  if (!onDemandOnlyColumns && !tableNameCompletion) {
     for (const refTable of completionContext.referencedTables) {
       if (refTable.columns && refTable.columns.length > 0) continue;
       const cacheKey = refTable.schema ? `${refTable.schema}.${refTable.name}` : refTable.name;
@@ -1890,7 +1909,7 @@ function scheduleCompletionMetadataRefresh(completionContext: ReturnType<typeof
         .catch(() => {});
     }
   }
-  if ((completionContext.suggestTables || completionContext.suggestJoinConditions) && completionContext.referencedTables.length > 0) {
+  if (!tableNameCompletion && (completionContext.suggestTables || completionContext.suggestJoinConditions) && completionContext.referencedTables.length > 0) {
     void ensureForeignKeysForTables(completionContext.referencedTables);
   }
 }
@@ -2054,8 +2073,8 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext
     }
   }
 
-  const isTableNameCompletionContext = completionContext.suggestTables || completionContext.exclusiveTableSuggestions;
-  const shouldFetchColumnsForCompletion = !onDemandOnlyColumns || ((completionContext.suggestColumns || completionContext.exclusiveColumnSuggestions) && !isTableNameCompletionContext) || !!completionContext.insertTable;
+  const tableNameCompletion = isTableNameCompletionContext(completionContext);
+  const shouldFetchColumnsForCompletion = !tableNameCompletion && (!onDemandOnlyColumns || completionContext.suggestColumns || completionContext.exclusiveColumnSuggestions || !!completionContext.insertTable);
   if (shouldFetchColumnsForCompletion) {
     await Promise.all(
       refs.map(async (refTable) => {
@@ -2077,7 +2096,7 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext
   }
   if (epoch !== completionEpoch) return null;
 
-  if ((completionContext.suggestTables || completionContext.suggestJoinConditions) && refs.length > 0) {
+  if (!tableNameCompletion && (completionContext.suggestTables || completionContext.suggestJoinConditions) && refs.length > 0) {
     const fkPrefetchTables = completionContext.suggestTables ? [...refs, ...tables.slice(0, MAX_JOIN_FK_PREFETCH_TABLES)] : refs;
     await ensureForeignKeysForTables(fkPrefetchTables.filter((table) => !("columns" in table) || !table.columns || table.columns.length === 0));
     if (epoch !== completionEpoch) return null;
@@ -2155,6 +2174,10 @@ function isReferencedTableQualifier(completionContext: ReturnType<typeof getSqlC
   return completionContext.referencedTables.some((table) => table.alias?.toLowerCase() === qualifier || table.name.toLowerCase() === qualifier || (!!qualifiedColumnTarget && completionTablesMatch(table, qualifiedColumnTarget)));
 }
 
+function isTableNameCompletionContext(completionContext: ReturnType<typeof getSqlCompletionContext>): boolean {
+  return completionContext.suggestTables || completionContext.exclusiveTableSuggestions;
+}
+
 function mergeCompletionObjects(existing: SqlCompletionObject[], incoming: SqlCompletionObject[]) {
   const merged = [...existing];
   const seen = new Set(existing.map((object) => `${object.type}:${object.schema ?? ""}:${object.name}:${object.parentName ?? ""}`.toLowerCase()));
@@ -2444,8 +2467,10 @@ onMounted(async () => {
             update.changes.iterChanges((_fromA, _toA, _fromB, _toB, inserted) => {
               insertedText += inserted.toString();
             });
-            if (insertedText.endsWith(".")) {
-              startCompletion(update.view);
+            if (shouldStartSqlCompletionAfterInput(insertedText, update.view)) {
+              window.setTimeout(() => {
+                if (!isEditorComposing(update.view)) startCompletion(update.view);
+              }, 0);
             }
           }
         }