Skip to content

Commit 16b4371

Browse files
takemyoxygentakemyoxygen
committed
Allow drivers to provide completions from the raw query text and position
1 parent cb22702 commit 16b4371

File tree

3 files changed

+81
-52
lines changed

3 files changed

+81
-52
lines changed

‎packages/language-server/src/connection.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { NSDatabase, IConnectionDriver, IConnection, MConnectionExplorer, ContextValue, InternalID, IQueryOptions } from '@sqltools/types';
1+
import { NSDatabase, IConnectionDriver, IConnection, MConnectionExplorer, ContextValue, InternalID, IQueryOptions, IRawCompletionItem } from '@sqltools/types';
22
import decorateLSException from '@sqltools/util/decorators/ls-decorate-exception';
33
import { getConnectionId } from '@sqltools/util/connection';
44
import ConfigRO from '@sqltools/util/config-manager';
@@ -177,4 +177,9 @@ export default class Connection {
177177
if (typeof this.conn.getStaticCompletions !== 'function') return Promise.resolve({} as any);
178178
return this.conn.getStaticCompletions();
179179
}
180+
181+
public getCompletionsForRawQuery(text: string, currentOffset: number): Promise<IRawCompletionItem[] | null> {
182+
if (typeof this.conn.getCompletionsForRawQuery !== 'function') return Promise.resolve(null);
183+
return this.conn.getCompletionsForRawQuery(text, currentOffset);
184+
}
180185
}

‎packages/plugins/intellisense/language-server.ts

Lines changed: 70 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -107,13 +107,70 @@ export default class IntellisensePlugin<T extends ILanguageServer> implements IL
107107
return [];
108108
}
109109

110-
private onCompletion: Arg0<ILanguageServer['onCompletion']> = async params => {
110+
private getCompletionsFromHueAst = async ({ currentWord, conn, text, currentOffset }: { currentWord: string; conn: Connection | null; text: string; currentOffset: number }) => {
111111
let completionsMap = {
112112
query: [],
113113
tables: [],
114114
columns: [],
115115
dbs: []
116116
};
117+
118+
const hueAst = sqlAutocompleteParser.parseSql(text.substring(0, currentOffset), text.substring(currentOffset));
119+
120+
completionsMap.query = (hueAst.suggestKeywords || []).filter(kw => kw.value.startsWith(currentWord)).map(kw => <CompletionItem>{
121+
label: kw.value,
122+
detail: kw.value,
123+
filterText: kw.value,
124+
// weights provided by hue are in reversed order
125+
sortText: `${String(10000 - kw.weight).padStart(5, '0')}:${kw.value}`,
126+
kind: CompletionItemKind.Keyword,
127+
documentation: {
128+
value: `\`\`\`yaml\nWORD: ${kw.value}\n\`\`\``,
129+
kind: 'markdown'
130+
}
131+
})
132+
const visitedKeywords: [string] = (hueAst.suggestKeywords || []).map(kw => kw.value)
133+
if (!conn) {
134+
log.info('no active connection completions count: %d', completionsMap.query.length);
135+
return completionsMap.query;
136+
};
137+
// Can't distinguish functions types, so put all other keywords
138+
if (hueAst.suggestFunctions || hueAst.suggestAggregateFunctions || hueAst.suggestAnalyticFunctions) {
139+
const staticCompletions = await conn.getStaticCompletions();
140+
for (let keyword in staticCompletions) {
141+
if (visitedKeywords.includes(keyword)) {
142+
continue;
143+
}
144+
if (!keyword.startsWith(currentWord)) {
145+
continue;
146+
}
147+
visitedKeywords.push(keyword);
148+
const value: NSDatabase.IStaticCompletion = staticCompletions[keyword]
149+
completionsMap.query.push({
150+
...value, sortText: `4:${value.label}`,
151+
kind: CompletionItemKind.Function
152+
});
153+
}
154+
}
155+
156+
const [tableCompletions, columnCompletions, dbCompletions] = await Promise.all([
157+
(hueAst.suggestTables != undefined) ? this.getTableCompletions({ currentWord, conn, suggestTables: hueAst.suggestTables }) : [],
158+
(hueAst.suggestColumns != undefined) ? this.getColumnCompletions({ currentWord, conn, suggestColumns: hueAst.suggestColumns }) : [],
159+
(hueAst.suggestDatabases != undefined) ? this.getDatabasesCompletions({ currentWord, conn, suggestDatabases: hueAst.suggestDatabases }) : [],
160+
]);
161+
completionsMap.tables = tableCompletions;
162+
completionsMap.columns = columnCompletions;
163+
completionsMap.dbs = dbCompletions;
164+
165+
const completions = completionsMap.columns
166+
.concat(completionsMap.tables)
167+
.concat(completionsMap.dbs)
168+
.concat(completionsMap.query);
169+
170+
return completions;
171+
}
172+
173+
private onCompletion: Arg0<ILanguageServer['onCompletion']> = async params => {
117174
try {
118175
const {
119176
currentWord,
@@ -122,61 +179,23 @@ export default class IntellisensePlugin<T extends ILanguageServer> implements IL
122179
currentOffset
123180
} = await this.getQueryData(params);
124181

125-
const hueAst = sqlAutocompleteParser.parseSql(text.substring(0, currentOffset), text.substring(currentOffset));
126-
127-
completionsMap.query = (hueAst.suggestKeywords || []).filter(kw => kw.value.startsWith(currentWord)).map(kw => <CompletionItem>{
128-
label: kw.value,
129-
detail: kw.value,
130-
filterText: kw.value,
131-
// weights provided by hue are in reversed order
132-
sortText: `${String(10000 - kw.weight).padStart(5, '0')}:${kw.value}`,
133-
kind: CompletionItemKind.Keyword,
134-
documentation: {
135-
value: `\`\`\`yaml\nWORD: ${kw.value}\n\`\`\``,
136-
kind: 'markdown'
137-
}
138-
})
139-
const visitedKeywords: [string] = (hueAst.suggestKeywords || []).map(kw => kw.value)
140-
if (!conn) {
141-
log.info('no active connection completions count: %d', completionsMap.query.length);
142-
return completionsMap.query;
143-
};
144-
// Can't distinguish functions types, so put all other keywords
145-
if (hueAst.suggestFunctions || hueAst.suggestAggregateFunctions || hueAst.suggestAnalyticFunctions) {
146-
const staticCompletions = await conn.getStaticCompletions();
147-
for (let keyword in staticCompletions) {
148-
if (visitedKeywords.includes(keyword)) {
149-
continue;
150-
}
151-
if (!keyword.startsWith(currentWord)) {
152-
continue;
153-
}
154-
visitedKeywords.push(keyword);
155-
const value: NSDatabase.IStaticCompletion = staticCompletions[keyword]
156-
completionsMap.query.push({
157-
...value, sortText: `4:${value.label}`,
158-
kind: CompletionItemKind.Function
159-
});
160-
}
182+
// First try to get completions from connection's getCompletionsForRawQuery method
183+
const connectionCompletions = await conn.getCompletionsForRawQuery(text, currentOffset);
184+
if (connectionCompletions !== null) {
185+
log.debug('using connection completions, count: %d', connectionCompletions.length);
186+
return connectionCompletions;
161187
}
188+
162189

163-
const [tableCompletions, columnCompletions, dbCompletions] = await Promise.all([
164-
(hueAst.suggestTables != undefined) ? this.getTableCompletions({ currentWord, conn, suggestTables: hueAst.suggestTables }) : [],
165-
(hueAst.suggestColumns != undefined) ? this.getColumnCompletions({ currentWord, conn, suggestColumns: hueAst.suggestColumns }) : [],
166-
(hueAst.suggestDatabases != undefined) ? this.getDatabasesCompletions({ currentWord, conn, suggestDatabases: hueAst.suggestDatabases }) : [],
167-
]);
168-
completionsMap.tables = tableCompletions;
169-
completionsMap.columns = columnCompletions;
170-
completionsMap.dbs = dbCompletions;
190+
// Fallback to hue AST-based completions
191+
log.debug('falling back to hue AST completions');
192+
const completions = await this.getCompletionsFromHueAst({ currentWord, conn, text, currentOffset });
193+
log.debug('total completions %d', completions.length);
194+
return completions;
171195
} catch (error) {
172196
log.error('got an error:\n %O', error);
197+
return [];
173198
}
174-
const completions = completionsMap.columns
175-
.concat(completionsMap.tables)
176-
.concat(completionsMap.dbs)
177-
.concat(completionsMap.query);
178-
log.debug('total completions %d', completions.length);
179-
return completions;
180199
}
181200

182201
public register(server: T) {

‎packages/types/index.d.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,10 @@ export interface IQueryOptions {
284284
export interface IConnectionDriverConstructor {
285285
new(credentials: IConnection<any>, getWorkspaceFolders?: LSIConnection['workspace']['getWorkspaceFolders']): IConnectionDriver;
286286
}
287+
export interface IRawCompletionItem {
288+
// TODO add more fields or maybe reuse existing types
289+
label: string;
290+
}
287291
export interface IConnectionDriver {
288292
connection: any;
289293
credentials: IConnection<any>;
@@ -311,6 +315,7 @@ export interface IConnectionDriver {
311315
port: number;
312316
}
313317
): Promise<{ port: number }>;
318+
getCompletionsForRawQuery?(text: string, currentOffset: number): Promise<IRawCompletionItem[]>;
314319
}
315320

316321
export declare enum ContextValue {

0 commit comments

Comments
 (0)