Add search features API

src/client/api/search.ts

@@ -0,0 +1,45 @@
+import { from, Observable, Subscription } from 'rxjs'
+import { ExtSearch } from 'src/extension/api/search'
+import { Connection } from 'src/protocol/jsonrpc2/connection'
+import { createProxyAndHandleRequests } from '../../common/proxy'
+import { TransformQuerySignature } from '../providers/queryTransformer'
+import { FeatureProviderRegistry } from '../providers/registry'
+import { SubscriptionMap } from './common'
+
+export interface SearchAPI {
+    $registerQueryTransformer(id: number): void
+    $unregister(id: number): void
+}
+
+export class Search implements SearchAPI {
+    private subscriptions = new Subscription()
+    private registrations = new SubscriptionMap()
+    private proxy: ExtSearch
+
+    constructor(
+        connection: Connection,
+        private queryTransformerRegistry: FeatureProviderRegistry<{}, TransformQuerySignature>
+    ) {
+        this.subscriptions.add(this.registrations)
+
+        this.proxy = createProxyAndHandleRequests('searchFeatures', connection, this)
+    }
+
+    public $registerQueryTransformer(id: number): void {
+        this.registrations.add(
+            id,
+            this.queryTransformerRegistry.registerProvider(
+                {},
+                (query: string): Observable<string | null | undefined> => from(this.proxy.$transformQuery(id, query))
+            )
+        )
+    }
+
+    public $unregister(id: number): void {
+        this.registrations.remove(id)
+    }
+
+    public unsubscribe(): void {
+        this.subscriptions.unsubscribe()
+    }
+}

src/client/controller.ts

@@ -19,6 +19,7 @@ import { ClientConfiguration } from './api/configuration'
 import { ClientContext } from './api/context'
 import { ClientDocuments } from './api/documents'
 import { ClientLanguageFeatures } from './api/languageFeatures'
+import { Search } from './api/search'
 import { ClientWindows } from './api/windows'
 import { applyContextUpdate, EMPTY_CONTEXT } from './context/context'
 import { EMPTY_ENVIRONMENT, Environment } from './environment'
@@ -268,6 +269,7 @@ export class Controller<X extends Extension, C extends ConfigurationCascade> imp
                 this.registries.textDocumentReferences
             )
         )
+        subscription.add(new Search(client, this.registries.queryTransformer))
         subscription.add(new ClientCommands(client, this.registries.commands))
     }
 

src/client/providers/queryTransformer.test.ts

@@ -0,0 +1,66 @@
+import * as assert from 'assert'
+import { of } from 'rxjs'
+import { TestScheduler } from 'rxjs/testing'
+import { transformQuery, TransformQuerySignature } from './queryTransformer'
+
+const scheduler = () => new TestScheduler((a, b) => assert.deepStrictEqual(a, b))
+
+const FIXTURE_INPUT = 'foo'
+const FIXTURE_RESULT_MERGED = 'foo bar'
+const FIXTURE_RESULT = 'bar'
+
+describe('transformQuery', () => {
+    describe('0 providers', () => {
+        it('returns original query', () =>
+            scheduler().run(({ cold, expectObservable }) =>
+                expectObservable(
+                    transformQuery(cold<TransformQuerySignature[]>('-a-|', { a: [] }), FIXTURE_INPUT)
+                ).toBe('-a-|', {
+                    a: FIXTURE_INPUT,
+                })
+            ))
+    })
+
+    describe('1 provider', () => {
+        it('returns result from provider', () =>
+            scheduler().run(({ cold, expectObservable }) =>
+                expectObservable(
+                    transformQuery(
+                        cold<TransformQuerySignature[]>('-a-|', {
+                            a: [q => of(FIXTURE_RESULT)],
+                        }),
+                        FIXTURE_INPUT
+                    )
+                ).toBe('-a-|', { a: FIXTURE_RESULT })
+            ))
+    })
+
+    describe('2 providers', () => {
+        it('returns a single query transformed by both providers', () =>
+            scheduler().run(({ cold, expectObservable }) =>
+                expectObservable(
+                    transformQuery(
+                        cold<TransformQuerySignature[]>('-a-|', {
+                            a: [q => of(q), q => of(`${q} ${FIXTURE_RESULT}`)],
+                        }),
+                        FIXTURE_INPUT
+                    )
+                ).toBe('-a-|', { a: FIXTURE_RESULT_MERGED })
+            ))
+    })
+
+    describe('Multiple emissions', () => {
+        it('returns stream of results', () =>
+            scheduler().run(({ cold, expectObservable }) =>
+                expectObservable(
+                    transformQuery(
+                        cold<TransformQuerySignature[]>('-a-b-|', {
+                            a: [q => of(q)],
+                            b: [q => of(FIXTURE_RESULT)],
+                        }),
+                        FIXTURE_INPUT
+                    )
+                ).toBe('-a-b-|', { a: FIXTURE_INPUT, b: FIXTURE_RESULT })
+            ))
+    })
+})

src/client/providers/queryTransformer.ts

@@ -0,0 +1,39 @@
+import { Observable, of } from 'rxjs'
+import { flatMap, map, switchMap } from 'rxjs/operators'
+import { FeatureProviderRegistry } from './registry'
+
+export type TransformQuerySignature = (query: string) => Observable<string | null | undefined>
+
+/** Provides transformed queries from the first search extension. */
+export class QueryTransformerRegistry extends FeatureProviderRegistry<{}, TransformQuerySignature> {
+    public transformQuery(query: string): Observable<string | null | undefined> {
+        return transformQuery(this.providers, query)
+    }
+}
+
+/**
+ * Returns an observable that emits a query transformed by all providers whenever any of the last-emitted set of providers emits
+ * a query.
+ *
+ * Most callers should use QueryTransformerRegistry's transformQuery method, which uses the registered search
+ * providers.
+ */
+export function transformQuery(
+    providers: Observable<TransformQuerySignature[]>,
+    query: string
+): Observable<string | null | undefined> {
+    return providers.pipe(
+        switchMap(providers => {
+            if (providers.length === 0) {
+                return [query]
+            }
+            return providers.reduce(
+                (currentQuery, transformQuery) =>
+                    currentQuery.pipe(
+                        flatMap(q => transformQuery(q).pipe(map(transformedQuery => transformedQuery || q)))
+                    ),
+                of(query)
+            )
+        })
+    )
+}

src/client/registries.ts

@@ -7,6 +7,7 @@ import { ContributionRegistry } from './providers/contribution'
 import { TextDocumentDecorationProviderRegistry } from './providers/decoration'
 import { TextDocumentHoverProviderRegistry } from './providers/hover'
 import { TextDocumentLocationProviderRegistry, TextDocumentReferencesProviderRegistry } from './providers/location'
+import { QueryTransformerRegistry } from './providers/queryTransformer'
 
 /**
  * Registries is a container for all provider registries.
@@ -25,4 +26,5 @@ export class Registries<X extends Extension, C extends ConfigurationCascade> {
     public readonly textDocumentTypeDefinition = new TextDocumentLocationProviderRegistry()
     public readonly textDocumentHover = new TextDocumentHoverProviderRegistry()
     public readonly textDocumentDecoration = new TextDocumentDecorationProviderRegistry()
+    public readonly queryTransformer = new QueryTransformerRegistry()
 }

src/extension/api/search.ts

@@ -0,0 +1,26 @@
+import { Unsubscribable } from 'rxjs'
+import { QueryTransformerProvider } from 'sourcegraph'
+import { SearchAPI } from 'src/client/api/search'
+import { ProviderMap } from './common'
+
+/** @internal */
+export interface ExtSearchAPI {
+    $transformQuery: (id: number, query: string) => Promise<string | null | undefined>
+}
+
+/** @internal */
+export class ExtSearch implements ExtSearchAPI {
+    private registrations = new ProviderMap<QueryTransformerProvider>(id => this.proxy.$unregister(id))
+    constructor(private proxy: SearchAPI) {}
+
+    public registerQueryTransformer(provider: QueryTransformerProvider): Unsubscribable {
+        const { id, subscription } = this.registrations.add(provider)
+        this.proxy.$registerQueryTransformer(id)
+        return subscription
+    }
+
+    public $transformQuery(id: number, query: string): Promise<string | null | undefined> {
+        const provider = this.registrations.get<QueryTransformerProvider>(id)
+        return Promise.resolve(provider.transformQuery(query))
+    }
+}

src/extension/extensionHost.ts

@@ -8,6 +8,7 @@ import { ExtConfiguration } from './api/configuration'
 import { ExtContext } from './api/context'
 import { ExtDocuments } from './api/documents'
 import { ExtLanguageFeatures } from './api/languageFeatures'
+import { ExtSearch } from './api/search'
 import { ExtWindows } from './api/windows'
 import { Location } from './types/location'
 import { Position } from './types/position'
@@ -82,6 +83,9 @@ function createExtensionHandle(initData: InitData, connection: Connection): type
     const languageFeatures = new ExtLanguageFeatures(proxy('languageFeatures'), documents)
     handleRequests(connection, 'languageFeatures', languageFeatures)
 
+    const searchFeatures = new ExtSearch(proxy('searchFeatures'))
+    handleRequests(connection, 'searchFeatures', searchFeatures)
+
     const commands = new ExtCommands(proxy('commands'))
     handleRequests(connection, 'commands', commands)
 
@@ -129,6 +133,10 @@ function createExtensionHandle(initData: InitData, connection: Connection): type
                 languageFeatures.registerReferenceProvider(selector, provider),
         },
 
+        search: {
+            registerQueryTransformer: provider => searchFeatures.registerQueryTransformer(provider),
+        },
+
         commands: {
             registerCommand: (command, callback) => commands.registerCommand({ command, callback }),
             executeCommand: (command, ...args) => commands.executeCommand(command, args),

src/integration-test/search.test.ts

@@ -0,0 +1,46 @@
+import * as assert from 'assert'
+import { take } from 'rxjs/operators'
+import { integrationTestContext } from './helpers.test'
+
+describe('search (integration)', () => {
+    it('registers a query transformer', async () => {
+        const { clientController, extensionHost, ready } = await integrationTestContext()
+
+        // Register the provider and call it
+        const unsubscribe = extensionHost.search.registerQueryTransformer({ transformQuery: () => 'bar' })
+        await ready
+        assert.deepStrictEqual(
+            await clientController.registries.queryTransformer
+                .transformQuery('foo')
+                .pipe(take(1))
+                .toPromise(),
+            'bar'
+        )
+
+        // Unregister the provider and ensure it's removed
+        unsubscribe.unsubscribe()
+        assert.deepStrictEqual(
+            await clientController.registries.queryTransformer
+                .transformQuery('foo')
+                .pipe(take(1))
+                .toPromise(),
+            'foo'
+        )
+    })
+
+    it('supports multiple query transformers', async () => {
+        const { clientController, extensionHost, ready } = await integrationTestContext()
+
+        // Register the provider and call it
+        extensionHost.search.registerQueryTransformer({ transformQuery: q => `${q} bar` })
+        extensionHost.search.registerQueryTransformer({ transformQuery: q => `${q} baz` })
+        await ready
+        assert.deepStrictEqual(
+            await clientController.registries.queryTransformer
+                .transformQuery('foo')
+                .pipe(take(1))
+                .toPromise(),
+            'foo bar baz'
+        )
+    })
+})

src/sourcegraph.d.ts

@@ -717,6 +717,14 @@ declare module 'sourcegraph' {
         ): ProviderResult<Location[]>
     }
 
+    export interface QueryTransformerProvider {
+        transformQuery(query: string): ProviderResult<string>
+    }
+
+    export namespace search {
+        export function registerQueryTransformer(provider: QueryTransformerProvider): Unsubscribable
+    }
+
     export namespace languages {
         export function registerHoverProvider(selector: DocumentSelector, provider: HoverProvider): Unsubscribable