Update LunrSearchEngine to be stream-based

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
This commit is contained in:
Eric Peterson
2022-02-26 18:25:05 +01:00
parent cef19ee966
commit c87be57d3a
5 changed files with 370 additions and 59 deletions
@@ -16,28 +16,61 @@
import { getVoidLogger } from '@backstage/backend-common';
import lunr from 'lunr';
import { SearchEngine } from '@backstage/search-common';
import { IndexableDocument, SearchEngine } from '@backstage/search-common';
import {
ConcreteLunrQuery,
LunrSearchEngine,
decodePageCursor,
encodePageCursor,
} from './LunrSearchEngine';
import { LunrSearchEngineIndexer } from './LunrSearchEngineIndexer';
import { TestPipeline } from '../test-utils';
/**
* Just used to test the default translator shipped with LunrSearchEngine.
*/
class LunrSearchEngineForTranslatorTests extends LunrSearchEngine {
class LunrSearchEngineForTests extends LunrSearchEngine {
getDocStore() {
return this.docStore;
}
setDocStore(docStore: Record<string, IndexableDocument>) {
this.docStore = docStore;
}
getLunrIndices() {
return this.lunrIndices;
}
getTranslator() {
return this.translator;
}
}
const indexerMock = {
on: jest.fn(),
buildIndex: jest.fn(),
getDocumentStore: jest.fn(),
};
jest.mock('./LunrSearchEngineIndexer', () => ({
LunrSearchEngineIndexer: jest.fn().mockImplementation(() => indexerMock),
}));
const getActualIndexer = (engine: SearchEngine, index: string) => {
(LunrSearchEngineIndexer as unknown as jest.Mock).mockImplementationOnce(
() => {
const ActualIndexer = jest.requireActual(
'./LunrSearchEngineIndexer',
).LunrSearchEngineIndexer;
return new ActualIndexer();
},
);
return engine.getIndexer(index);
};
describe('LunrSearchEngine', () => {
let testLunrSearchEngine: SearchEngine;
beforeEach(() => {
testLunrSearchEngine = new LunrSearchEngine({ logger: getVoidLogger() });
jest.clearAllMocks();
});
describe('translator', () => {
@@ -65,7 +98,7 @@ describe('LunrSearchEngine', () => {
});
it('should return translated query', async () => {
const inspectableSearchEngine = new LunrSearchEngineForTranslatorTests({
const inspectableSearchEngine = new LunrSearchEngineForTests({
logger: getVoidLogger(),
});
const translatorUnderTest = inspectableSearchEngine.getTranslator();
@@ -107,7 +140,7 @@ describe('LunrSearchEngine', () => {
});
it('should have default offset and limit', async () => {
const inspectableSearchEngine = new LunrSearchEngineForTranslatorTests({
const inspectableSearchEngine = new LunrSearchEngineForTests({
logger: getVoidLogger(),
});
const translatorUnderTest = inspectableSearchEngine.getTranslator();
@@ -148,7 +181,7 @@ describe('LunrSearchEngine', () => {
});
it('should return translated query with 1 filter', async () => {
const inspectableSearchEngine = new LunrSearchEngineForTranslatorTests({
const inspectableSearchEngine = new LunrSearchEngineForTests({
logger: getVoidLogger(),
});
const translatorUnderTest = inspectableSearchEngine.getTranslator();
@@ -193,7 +226,7 @@ describe('LunrSearchEngine', () => {
});
it('should handle single-item array filter as scalar value', async () => {
const inspectableSearchEngine = new LunrSearchEngineForTranslatorTests({
const inspectableSearchEngine = new LunrSearchEngineForTests({
logger: getVoidLogger(),
});
const translatorUnderTest = inspectableSearchEngine.getTranslator();
@@ -224,7 +257,7 @@ describe('LunrSearchEngine', () => {
});
it('should return translated query with multiple filters', async () => {
const inspectableSearchEngine = new LunrSearchEngineForTranslatorTests({
const inspectableSearchEngine = new LunrSearchEngineForTests({
logger: getVoidLogger(),
});
const translatorUnderTest = inspectableSearchEngine.getTranslator();
@@ -273,7 +306,7 @@ describe('LunrSearchEngine', () => {
});
it('should throw if translated query references missing field', async () => {
const inspectableSearchEngine = new LunrSearchEngineForTranslatorTests({
const inspectableSearchEngine = new LunrSearchEngineForTests({
logger: getVoidLogger(),
});
const translatorUnderTest = inspectableSearchEngine.getTranslator();
@@ -334,7 +367,13 @@ describe('LunrSearchEngine', () => {
];
// Mock indexing of 1 document
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
// Perform search query
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -359,7 +398,13 @@ describe('LunrSearchEngine', () => {
];
// Mock indexing of 1 document
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
// Perform search query
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -392,7 +437,13 @@ describe('LunrSearchEngine', () => {
];
// Mock indexing of 1 document
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
// Perform search query
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -424,7 +475,13 @@ describe('LunrSearchEngine', () => {
];
// Mock indexing of 1 document
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
// Perform search query
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -456,7 +513,13 @@ describe('LunrSearchEngine', () => {
];
// Mock indexing of 1 document
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
// Perform search query
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -489,7 +552,13 @@ describe('LunrSearchEngine', () => {
];
// Mock indexing of 1 document
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
// Perform search query
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -522,7 +591,13 @@ describe('LunrSearchEngine', () => {
];
// Mock indexing of 1 document
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
// Perform search query
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -560,7 +635,13 @@ describe('LunrSearchEngine', () => {
];
// Mock indexing of 2 documents
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
// Perform search query
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -604,8 +685,21 @@ describe('LunrSearchEngine', () => {
];
// Mock 2 indices with 1 document each
await testLunrSearchEngine.index('test-index', mockDocuments);
await testLunrSearchEngine.index('test-index-2', mockDocuments2);
const indexer1 = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
const indexer2 = await getActualIndexer(
testLunrSearchEngine,
'test-index-2',
);
await TestPipeline.withSubject(indexer1)
.withDocuments(mockDocuments)
.execute();
await TestPipeline.withSubject(indexer2)
.withDocuments(mockDocuments2)
.execute();
// Perform search query scoped to "test-index-2" with a filter on the field "extraField"
const mockedSearchResult = await testLunrSearchEngine.query({
term: 'testTitle',
@@ -642,7 +736,13 @@ describe('LunrSearchEngine', () => {
];
// Mock indexing of 2 documents
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
// Perform search query
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -695,8 +795,20 @@ describe('LunrSearchEngine', () => {
];
// Mock 2 indices with 2 documents each
await testLunrSearchEngine.index('test-index', mockDocuments);
await testLunrSearchEngine.index('test-index-2', mockDocuments2);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
const indexer2 = await getActualIndexer(
testLunrSearchEngine,
'test-index-2',
);
await TestPipeline.withSubject(indexer2)
.withDocuments(mockDocuments2)
.execute();
// Perform search query scoped to "test-index-2"
const mockedSearchResult = await testLunrSearchEngine.query({
@@ -734,7 +846,13 @@ describe('LunrSearchEngine', () => {
location: `test/location/${i}`,
}));
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(
testLunrSearchEngine,
'test-index',
);
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
const mockedSearchResult = await testLunrSearchEngine.query({
term: 'testTitle',
@@ -767,7 +885,10 @@ describe('LunrSearchEngine', () => {
location: `test/location/${i}`,
}));
await testLunrSearchEngine.index('test-index', mockDocuments);
const indexer = await getActualIndexer(testLunrSearchEngine, 'test-index');
await TestPipeline.withSubject(indexer)
.withDocuments(mockDocuments)
.execute();
const mockedSearchResult = await testLunrSearchEngine.query({
term: 'testTitle',
@@ -793,22 +914,46 @@ describe('LunrSearchEngine', () => {
});
describe('index', () => {
it('should index document', async () => {
const indexSpy = jest.spyOn(testLunrSearchEngine, 'index');
const mockDocuments = [
{
title: 'testTerm',
text: 'testText',
location: 'test/location',
},
];
it('should get indexer', async () => {
const indexer = await testLunrSearchEngine.getIndexer('test-index');
expect(LunrSearchEngineIndexer).toHaveBeenCalled();
expect(indexer.on).toHaveBeenCalledWith('close', expect.any(Function));
});
// call index func and ensure the index func was invoked.
await testLunrSearchEngine.index('test-index', mockDocuments);
expect(indexSpy).toHaveBeenCalled();
expect(indexSpy).toHaveBeenCalledWith('test-index', [
{ title: 'testTerm', text: 'testText', location: 'test/location' },
]);
it('should manage indices and docs on close', async () => {
const doc = { title: 'A doc', text: 'test', location: 'some-location' };
// Set up an inspectable search engine to pre-set some data.
const inspectableSearchEngine = new LunrSearchEngineForTests({
logger: getVoidLogger(),
});
inspectableSearchEngine.setDocStore({ 'existing-location': doc });
// Mock methds called by close handler.
indexerMock.buildIndex.mockReturnValueOnce('expected-index');
indexerMock.getDocumentStore.mockReturnValueOnce({
'new-location': doc,
});
// Get the indexer and invoke its close handler.
await inspectableSearchEngine.getIndexer('test-index');
const onClose = indexerMock.on.mock.calls[0][1] as Function;
onClose();
// Ensure mocked methods were called.
expect(indexerMock.buildIndex).toHaveBeenCalled();
expect(indexerMock.getDocumentStore).toHaveBeenCalled();
// Ensure the lunr index was written to the search engine.
expect(inspectableSearchEngine.getLunrIndices()).toStrictEqual({
'test-index': 'expected-index',
});
// Ensure documents are merged into the existing store.
expect(inspectableSearchEngine.getDocStore()).toStrictEqual({
'existing-location': doc,
'new-location': doc,
});
});
});
});
@@ -23,6 +23,7 @@ import {
} from '@backstage/search-common';
import lunr from 'lunr';
import { Logger } from 'winston';
import { LunrSearchEngineIndexer } from './LunrSearchEngineIndexer';
export type ConcreteLunrQuery = {
lunrQueryBuilder: lunr.Index.QueryBuilder;
@@ -124,30 +125,17 @@ export class LunrSearchEngine implements SearchEngine {
this.translator = translator;
}
async index(type: string, documents: IndexableDocument[]): Promise<void> {
const lunrBuilder = new lunr.Builder();
async getIndexer(type: string) {
const indexer = new LunrSearchEngineIndexer();
lunrBuilder.pipeline.add(lunr.trimmer, lunr.stopWordFilter, lunr.stemmer);
lunrBuilder.searchPipeline.add(lunr.stemmer);
// Make this lunr index aware of all relevant fields.
Object.keys(documents[0]).forEach(field => {
lunrBuilder.field(field);
indexer.on('close', () => {
// Once the stream is closed, build the index and store the documents in
// memory for later retrieval.
this.lunrIndices[type] = indexer.buildIndex();
this.docStore = { ...this.docStore, ...indexer.getDocumentStore() };
});
// Set "location" field as reference field
lunrBuilder.ref('location');
documents.forEach((document: IndexableDocument) => {
// Add document to Lunar index
lunrBuilder.add(document);
// Store documents in memory to be able to look up document using the ref during query time
// This is not how you should implement your SearchEngine implementation! Do not copy!
this.docStore[document.location] = document;
});
// "Rotate" the index by simply overwriting any existing index of the same name.
this.lunrIndices[type] = lunrBuilder.build();
return indexer;
}
async query(query: SearchQuery): Promise<SearchResultSet> {
@@ -0,0 +1,109 @@
/*
* Copyright 2022 The Backstage Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import lunr from 'lunr';
import { range } from 'lodash';
import { TestPipeline } from '../test-utils';
import { LunrSearchEngineIndexer } from './LunrSearchEngineIndexer';
const lunrBuilderAddSpy = jest.fn();
const lunrBuilderRefSpy = jest.fn();
const lunrBuilderFieldSpy = jest.fn();
const lunrBuilderPipelineAddSpy = jest.fn();
const lunrBuilderSearchPipelineAddSpy = jest.fn();
jest.mock('lunr', () => {
const actualLunr = jest.requireActual('lunr');
return {
...actualLunr,
Builder: jest.fn().mockImplementation(() => {
const actualBuilder = new actualLunr.Builder();
actualBuilder.add = lunrBuilderAddSpy;
actualBuilder.ref = lunrBuilderRefSpy;
actualBuilder.field = lunrBuilderFieldSpy;
actualBuilder.pipeline.add = lunrBuilderPipelineAddSpy;
actualBuilder.searchPipeline.add = lunrBuilderSearchPipelineAddSpy;
return actualBuilder;
}),
};
});
describe('LunrSearchEngineIndexer', () => {
let indexer: LunrSearchEngineIndexer;
beforeEach(() => {
jest.clearAllMocks();
indexer = new LunrSearchEngineIndexer();
});
it('should index documents', async () => {
const documents = [
{
title: 'testTerm',
text: 'testText',
location: 'test/location',
},
];
await TestPipeline.withSubject(indexer).withDocuments(documents).execute();
expect(lunrBuilderAddSpy).toHaveBeenCalledWith(documents[0]);
});
it('should index documents in bulk', async () => {
const documents = range(350).map(i => ({
title: `Hello World ${i}`,
text: 'Lorem Ipsum',
location: `location-${i}`,
}));
await TestPipeline.withSubject(indexer).withDocuments(documents).execute();
expect(lunrBuilderAddSpy).toHaveBeenCalledTimes(350);
});
it('should initialize schema', async () => {
const documents = [
{
title: 'testTerm',
text: 'testText',
location: 'test/location',
extra: 'field',
},
];
await TestPipeline.withSubject(indexer).withDocuments(documents).execute();
// Builder ref should be set to location (and only once).
expect(lunrBuilderRefSpy).toHaveBeenCalledTimes(1);
expect(lunrBuilderRefSpy).toHaveBeenLastCalledWith('location');
// Builder fields should be based on document fields.
expect(lunrBuilderFieldSpy).toHaveBeenCalledTimes(4);
expect(lunrBuilderFieldSpy).toHaveBeenCalledWith('title');
expect(lunrBuilderFieldSpy).toHaveBeenCalledWith('text');
expect(lunrBuilderFieldSpy).toHaveBeenCalledWith('location');
expect(lunrBuilderFieldSpy).toHaveBeenCalledWith('extra');
});
it('should configure lunr pipeline', async () => {
expect(lunrBuilderSearchPipelineAddSpy).toHaveBeenLastCalledWith(
lunr.stemmer,
);
expect(lunrBuilderPipelineAddSpy).toHaveBeenCalledWith(
...[lunr.trimmer, lunr.stopWordFilter, lunr.stemmer],
);
});
});
@@ -0,0 +1,68 @@
/*
* Copyright 2022 The Backstage Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { IndexableDocument } from '@backstage/search-common';
import lunr from 'lunr';
import { BatchSearchEngineIndexer } from '../indexing';
export class LunrSearchEngineIndexer extends BatchSearchEngineIndexer {
private schemaInitialized = false;
private builder: lunr.Builder;
private docStore: Record<string, IndexableDocument> = {};
constructor() {
super({ batchSize: 100 });
this.builder = new lunr.Builder();
this.builder.pipeline.add(lunr.trimmer, lunr.stopWordFilter, lunr.stemmer);
this.builder.searchPipeline.add(lunr.stemmer);
}
// No async initialization required.
async initialize(): Promise<void> {}
async finalize(): Promise<void> {}
async index(documents: IndexableDocument[]): Promise<void> {
if (!this.schemaInitialized) {
// Make this lunr index aware of all relevant fields.
Object.keys(documents[0]).forEach(field => {
this.builder.field(field);
});
// Set "location" field as reference field
this.builder.ref('location');
this.schemaInitialized = true;
}
documents.forEach(document => {
// Add document to Lunar index
this.builder.add(document);
// Store documents in memory to be able to look up document using the ref during query time
// This is not how you should implement your SearchEngine implementation! Do not copy!
this.docStore[document.location] = document;
});
}
buildIndex() {
return this.builder.build();
}
getDocumentStore() {
return this.docStore;
}
}
@@ -16,3 +16,4 @@
export { LunrSearchEngine } from './LunrSearchEngine';
export type { ConcreteLunrQuery } from './LunrSearchEngine';
export type { LunrSearchEngineIndexer } from './LunrSearchEngineIndexer';