From e3e625223010276d20724d6cc941002931acaa56 Mon Sep 17 00:00:00 2001 From: Will Scullin Date: Thu, 24 Oct 2024 10:58:18 -0700 Subject: [PATCH] Add support for top values --- src/common/schema.ts | 24 ++- src/common/types/message_types.ts | 9 +- src/extension/commands/open_composer.ts | 87 ++-------- .../utils/composer_message_manager.ts | 151 ++++++++++++++++++ src/extension/webviews/composer_page/App.tsx | 26 ++- .../webviews/composer_page/Composer.tsx | 4 +- 6 files changed, 213 insertions(+), 88 deletions(-) create mode 100644 src/extension/commands/utils/composer_message_manager.ts diff --git a/src/common/schema.ts b/src/common/schema.ts index cd4f81eb..e3dcbe66 100644 --- a/src/common/schema.ts +++ b/src/common/schema.ts @@ -21,7 +21,13 @@ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. */ -import {Explore, Field, JoinRelationship} from '@malloydata/malloy'; +import { + Explore, + Field, + JoinRelationship, + ModelDef, + isSourceDef, +} from '@malloydata/malloy'; export function isFieldAggregate(field: Field) { return field.isAtomicField() && field.isCalculation(); @@ -114,6 +120,22 @@ export const quoteIfNecessary = (element: string) => { return element; }; +/** + * Retrieve a source from a model safely + * + * @param modelDef Model definition + * @param sourceName Source name + * @returns SourceDef for given name, or throws if not a source + */ + +export const getSourceDef = (modelDef: ModelDef, sourceName: string) => { + const result = modelDef.contents[sourceName]; + if (isSourceDef(result)) { + return result; + } + throw new Error(`Not a source: ${sourceName}`); +}; + const RESERVED: string[] = [ 'ALL', 'AND', diff --git a/src/common/types/message_types.ts b/src/common/types/message_types.ts index 8b4cb2db..5f47d88c 100755 --- a/src/common/types/message_types.ts +++ b/src/common/types/message_types.ts @@ -287,6 +287,7 @@ export enum ComposerMessageType { NewModel = 'new-model', ResultSuccess = 'result-success', ResultError = 'result-error', + SearchIndex = 'search-index', } export interface ComposerMessageNewModel { @@ -309,10 +310,16 @@ export interface ComposerMessageResultError { error: string; } +export interface ComposerMessageSearchIndex { + type: ComposerMessageType.SearchIndex; + result: ResultJSON; +} + export type ComposerMessage = | ComposerMessageNewModel | ComposerMessageResultSuccess - | ComposerMessageResultError; + | ComposerMessageResultError + | ComposerMessageSearchIndex; export enum ComposerPageMessageType { Ready = 'ready', diff --git a/src/extension/commands/open_composer.ts b/src/extension/commands/open_composer.ts index 837a2fda..6e2971db 100644 --- a/src/extension/commands/open_composer.ts +++ b/src/extension/commands/open_composer.ts @@ -7,21 +7,11 @@ import * as vscode from 'vscode'; import {getWebviewHtml} from '../webviews'; -import { - ComposerMessage, - ComposerMessageType, - ComposerPageMessage, - ComposerPageMessageType, -} from '../../common/types/message_types'; -import { - getActiveDocumentMetadata, - runMalloyQuery, -} from './utils/run_query_utils'; +import {getActiveDocumentMetadata} from './utils/run_query_utils'; import {Utils} from 'vscode-uri'; import {MALLOY_EXTENSION_STATE} from '../state'; import {WorkerConnection} from '../worker_connection'; -import {WebviewMessageManager} from '../webview_message_manager'; -import {QuerySpec} from '../../common/types/query_spec'; +import {ComposerMessageManager} from './utils/composer_message_manager'; const icon = 'turtle.svg'; @@ -59,74 +49,17 @@ export async function openComposer( sourceName ??= Object.keys(modelDef.contents)[0]; - const messages = new WebviewMessageManager< - ComposerMessage, - ComposerPageMessage - >(composerPanel); - messages.postMessage({ - type: ComposerMessageType.NewModel, + const messageManager = new ComposerMessageManager( + worker, + composerPanel, documentMeta, modelDef, sourceName, - viewName, - }); - messages.onReceiveMessage(message => { - switch (message.type) { - case ComposerPageMessageType.RunQuery: - { - const {id, query, queryName} = message; - vscode.window - .withProgress( - { - location: vscode.ProgressLocation.Notification, - title: `Running ${queryName}`, - cancellable: true, - }, - async (progress, cancellationToken) => { - const querySpec: QuerySpec = { - type: 'string', - text: query, - documentMeta, - }; - const result = await runMalloyQuery( - worker, - querySpec, - id, - queryName, - {withWebview: false}, - cancellationToken, - progress - ); - return result; - } - ) - .then( - result => { - if (result) { - messages.postMessage({ - type: ComposerMessageType.ResultSuccess, - id, - result, - }); - } else { - messages.postMessage({ - type: ComposerMessageType.ResultError, - id, - error: 'No results', - }); - } - }, - error => { - messages.postMessage({ - type: ComposerMessageType.ResultError, - id, - error: error instanceof Error ? error.message : `${error}`, - }); - } - ); - } - break; - } + viewName + ); + + composerPanel.onDidDispose(() => { + messageManager.dispose(); }); } } diff --git a/src/extension/commands/utils/composer_message_manager.ts b/src/extension/commands/utils/composer_message_manager.ts new file mode 100644 index 00000000..dc9bd4bd --- /dev/null +++ b/src/extension/commands/utils/composer_message_manager.ts @@ -0,0 +1,151 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +import * as vscode from 'vscode'; +import {v1 as uuid} from 'uuid'; +import {WebviewMessageManager} from '../../webview_message_manager'; +import { + ComposerMessage, + ComposerMessageType, + ComposerPageMessage, + ComposerPageMessageRunQuery, + ComposerPageMessageType, +} from '../../../common/types/message_types'; +import {DocumentMetadata, QuerySpec} from '../../../common/types/query_spec'; +import {ModelDef, SearchValueMapResult} from '@malloydata/malloy'; +import {runMalloyQuery} from './run_query_utils'; +import {WorkerConnection} from '../../worker_connection'; +import {getSourceDef} from '../../../common/schema'; + +export class ComposerMessageManager + extends WebviewMessageManager + implements vscode.Disposable +{ + constructor( + private worker: WorkerConnection, + composerPanel: vscode.WebviewPanel, + private documentMeta: DocumentMetadata, + private modelDef: ModelDef, + private sourceName: string, + viewName?: string + ) { + super(composerPanel); + + this.postMessage({ + type: ComposerMessageType.NewModel, + documentMeta, + modelDef, + sourceName, + viewName, + }); + + this.onReceiveMessage(async message => { + switch (message.type) { + case ComposerPageMessageType.RunQuery: + await this.runQuery(message); + break; + } + }); + + void this.initializeIndex(); + } + + async runQuery(message: ComposerPageMessageRunQuery) { + { + const {id, query, queryName} = message; + try { + const result = await this.runQueryWithProgress(id, queryName, query); + if (result) { + this.postMessage({ + type: ComposerMessageType.ResultSuccess, + id, + result, + }); + } else { + this.postMessage({ + type: ComposerMessageType.ResultError, + id, + error: 'No results', + }); + } + } catch (error) { + this.postMessage({ + type: ComposerMessageType.ResultError, + id, + error: error instanceof Error ? error.message : `${error}`, + }); + } + } + } + + async initializeIndex(): Promise { + const sourceDef = getSourceDef(this.modelDef, this.sourceName); + const indexQuery = sourceDef.fields.find( + ({name, as}) => (as || name) === 'search_index' + ); + const limit = 10; + + if (indexQuery) { + const searchMapMalloy = ` + run: ${this.sourceName} + -> ${indexQuery.as || indexQuery.name} + -> { + where: fieldType = 'string' + group_by: fieldName + aggregate: cardinality is count(fieldValue) + nest: values is { + select: fieldValue, weight + order_by: weight desc + limit: ${limit} + } + limit: 1000 + } + `; + const result = await this.runQueryWithProgress( + uuid(), + `Search index for ${this.sourceName}`, + searchMapMalloy + ); + if (result) { + this.postMessage({ + type: ComposerMessageType.SearchIndex, + result, + }); + } + } + return undefined; + } + + async runQueryWithProgress(id: string, queryName: string, query: string) { + return vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: queryName, + cancellable: true, + }, + async (progress, cancellationToken) => { + const querySpec: QuerySpec = { + type: 'string', + text: query, + documentMeta: this.documentMeta, + }; + const result = await runMalloyQuery( + this.worker, + querySpec, + id, + queryName, + {withWebview: false}, + cancellationToken, + progress + ); + return result; + } + ); + } + + dispose() {} +} diff --git a/src/extension/webviews/composer_page/App.tsx b/src/extension/webviews/composer_page/App.tsx index 8a8b94d4..a97acc20 100644 --- a/src/extension/webviews/composer_page/App.tsx +++ b/src/extension/webviews/composer_page/App.tsx @@ -16,7 +16,7 @@ import { } from '../../../common/types/message_types'; import {Composer} from './Composer'; import {DocumentMetadata} from '../../../common/types/query_spec'; -import {ModelDef, Result} from '@malloydata/malloy'; +import {ModelDef, Result, SearchValueMapResult} from '@malloydata/malloy'; import {RunQuery} from '@malloydata/query-composer'; export interface AppProps { @@ -35,6 +35,7 @@ export const App: React.FC = ({vscode}) => { const [modelDef, setModelDef] = React.useState(); const [sourceName, setSourceName] = React.useState(); const [viewName, setViewName] = React.useState(); + const [topValues, setTopValues] = React.useState(); React.useEffect(() => { const messageHandler = ({data}: MessageEvent) => { @@ -59,13 +60,23 @@ export const App: React.FC = ({vscode}) => { } } break; - case ComposerMessageType.ResultError: { - const {id, error} = data; - if (QueriesInFlight[id]) { - QueriesInFlight[id]?.reject(new Error(error)); - delete QueriesInFlight[id]; + case ComposerMessageType.ResultError: + { + const {id, error} = data; + if (QueriesInFlight[id]) { + QueriesInFlight[id]?.reject(new Error(error)); + delete QueriesInFlight[id]; + } } - } + break; + case ComposerMessageType.SearchIndex: + { + const result = Result.fromJSON(data.result); + setTopValues( + result._queryResult.result as unknown as SearchValueMapResult[] + ); + } + break; } }; window.addEventListener('message', messageHandler); @@ -105,6 +116,7 @@ export const App: React.FC = ({vscode}) => { sourceName={sourceName} viewName={viewName} runQuery={runQuery} + topValues={topValues} /> ); diff --git a/src/extension/webviews/composer_page/Composer.tsx b/src/extension/webviews/composer_page/Composer.tsx index 91540e16..bb3c8ec1 100644 --- a/src/extension/webviews/composer_page/Composer.tsx +++ b/src/extension/webviews/composer_page/Composer.tsx @@ -24,10 +24,9 @@ export interface ComposerProps { sourceName: string; viewName?: string; runQuery: RunQuery; + topValues: SearchValueMapResult[] | undefined; } -const topValues: SearchValueMapResult[] = []; - const nullUpdateQueryInUrl = () => {}; export const Composer: React.FC = ({ @@ -36,6 +35,7 @@ export const Composer: React.FC = ({ sourceName, viewName, runQuery: runQueryImp, + topValues, }) => { const {queryMalloy, queryName, queryModifiers, querySummary} = useQueryBuilder(