diff --git a/plugins/search-backend-node/src/indexing/BatchSearchEngineIndexer.test.ts b/plugins/search-backend-node/src/indexing/BatchSearchEngineIndexer.test.ts new file mode 100644 index 0000000000..b692ce3aff --- /dev/null +++ b/plugins/search-backend-node/src/indexing/BatchSearchEngineIndexer.test.ts @@ -0,0 +1,156 @@ +/* + * 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 { BatchSearchEngineIndexer } from './BatchSearchEngineIndexer'; +import { TestPipeline } from '../test-utils'; + +const indexSpy = jest.fn().mockResolvedValue(undefined); +const initializeSpy = jest.fn().mockResolvedValue(undefined); +const finalizeSpy = jest.fn().mockResolvedValue(undefined); + +class ConcreteBatchIndexer extends BatchSearchEngineIndexer { + async index(documents: IndexableDocument[]): Promise { + return indexSpy(documents); + } + async initialize(): Promise { + return initializeSpy(); + } + async finalize(): Promise { + return finalizeSpy(); + } +} + +describe('BatchSearchEngineIndexer', () => { + const document = { + title: 'Some Document', + text: 'Some document text.', + location: '/some/location', + }; + beforeEach(() => { + jest.resetAllMocks(); + }); + + it('should work end-to-end', async () => { + const indexer = new ConcreteBatchIndexer({ batchSize: 1 }); + await TestPipeline.withSubject(indexer) + .withDocuments([document, document, document]) + .execute(); + expect(indexSpy).toHaveBeenCalledTimes(3); + }); + + it('should call initialize at construction', () => { + // @ts-expect-error + const _indexer = new ConcreteBatchIndexer({ batchSize: 1 }); + + return new Promise(done => { + // Allow initialization to complete. + setImmediate(() => { + expect(initializeSpy).toHaveBeenCalled(); + done(); + }); + }); + }); + + it('should emit error if initialization throws', () => { + // Cause the initializer to throw. + const expectedError = new Error('some error'); + initializeSpy.mockRejectedValue(expectedError); + const indexer = new ConcreteBatchIndexer({ batchSize: 1 }); + + return new Promise(done => { + // Listen for the error and assert it's what was thrown. + indexer.on('error', error => { + expect(error).toStrictEqual(expectedError); + done(); + }); + + // Write a document to force the error state to become known. + indexer.write(document); + }); + }); + + it('should call index according to batchSize', () => { + const indexer = new ConcreteBatchIndexer({ batchSize: 2 }); + + return new Promise(done => { + // Listen for it to finish and assert the batches. + indexer.on('finish', () => { + expect(indexSpy).toHaveBeenCalledTimes(2); + expect(indexSpy).toHaveBeenNthCalledWith(1, [document, document]); + expect(indexSpy).toHaveBeenNthCalledWith(2, [document]); + done(); + }); + + // Write batchSize + 1 documents and end the stream. + indexer.write(document); + indexer.write(document); + indexer.write(document); + indexer.end(); + }); + }); + + it('should call index without exceeding batchSize', () => { + const indexer = new ConcreteBatchIndexer({ batchSize: 2 }); + + return new Promise(done => { + // Listen for it to finish and assert that it still wrote. + indexer.on('finish', () => { + expect(indexSpy).toHaveBeenCalledTimes(1); + expect(indexSpy).toHaveBeenNthCalledWith(1, [document]); + done(); + }); + + // Write batchSize - 1 documents and end the stream. + indexer.write(document); + indexer.end(); + }); + }); + + it('should emit error if index throws', () => { + // Cause the indexer to throw. + const expectedError = new Error('index error'); + indexSpy.mockRejectedValue(expectedError); + const indexer = new ConcreteBatchIndexer({ batchSize: 1 }); + + return new Promise(done => { + // Listen for the error and assert it's what was thrown. + indexer.on('error', error => { + expect(error).toStrictEqual(expectedError); + done(); + }); + + indexer.write(document); + }); + }); + + it('should emit error if finalize throws', () => { + // Cause the indexer to throw. + const expectedError = new Error('finalize error'); + finalizeSpy.mockRejectedValue(expectedError); + const indexer = new ConcreteBatchIndexer({ batchSize: 1 }); + + return new Promise(done => { + // Listen for the error and assert it's what was thrown. + indexer.on('error', error => { + expect(error).toStrictEqual(expectedError); + done(); + }); + + indexer.end(); + }); + }); +}); diff --git a/plugins/search-backend-node/src/indexing/BatchSearchEngineIndexer.ts b/plugins/search-backend-node/src/indexing/BatchSearchEngineIndexer.ts new file mode 100644 index 0000000000..4c29d0b573 --- /dev/null +++ b/plugins/search-backend-node/src/indexing/BatchSearchEngineIndexer.ts @@ -0,0 +1,121 @@ +/* + * 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 { assertError } from '@backstage/errors'; +import { IndexableDocument } from '@backstage/search-common'; +import { Writable } from 'stream'; + +export type BatchSearchEngineOptions = { + batchSize: number; +}; + +/** + * Base class encapsulating batch-based stream processing. Useful as a base + * class for search engine indexers. + */ +export abstract class BatchSearchEngineIndexer extends Writable { + private batchSize: number; + private currentBatch: IndexableDocument[] = []; + private initialized: Promise; + + constructor(options: BatchSearchEngineOptions) { + super({ objectMode: true }); + this.batchSize = options.batchSize; + + // @todo Once node v15 is minimum, convert to _construct implementation. + this.initialized = new Promise(done => { + // Necessary to allow concrete implementation classes to construct + // themselves before calling their initialize() methods. + setImmediate(async () => { + try { + await this.initialize(); + done(undefined); + } catch (e) { + assertError(e); + done(e); + } + }); + }); + } + + /** + * Receives an array of indexable documents (of size this.batchSize) which + * should be written to the search engine. This method won't be called again + * at least until it resolves. + */ + public abstract index(documents: IndexableDocument[]): Promise; + + /** + * Any asynchronous setup tasks can be performed here. + */ + public abstract initialize(): Promise; + + /** + * Any asynchronous teardown tasks can be performed here. + */ + public abstract finalize(): Promise; + + /** + * Encapsulates batch stream write logic. + * @internal + */ + async _write( + doc: IndexableDocument, + _e: any, + done: (error?: Error | null) => void, + ) { + // Wait for init before proceeding. Throw error if initialization failed. + const maybeError = await this.initialized; + if (maybeError) { + done(maybeError); + return; + } + + this.currentBatch.push(doc); + if (this.currentBatch.length < this.batchSize) { + done(); + return; + } + + try { + await this.index(this.currentBatch); + this.currentBatch = []; + done(); + } catch (e) { + assertError(e); + done(e); + } + } + + /** + * Encapsulates finalization and final error handling logic. + * @internal + */ + async _final(done: (error?: Error | null) => void) { + try { + // Index any remaining documents. + if (this.currentBatch.length) { + await this.index(this.currentBatch); + this.currentBatch = []; + } + await this.finalize(); + done(); + } catch (e) { + assertError(e); + done(e); + } + } +} diff --git a/plugins/search-backend-node/src/indexing/DecoratorBase.test.ts b/plugins/search-backend-node/src/indexing/DecoratorBase.test.ts new file mode 100644 index 0000000000..3b045dcddf --- /dev/null +++ b/plugins/search-backend-node/src/indexing/DecoratorBase.test.ts @@ -0,0 +1,157 @@ +/* + * 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 { DecoratorBase } from './DecoratorBase'; +import { TestPipeline } from '../test-utils'; + +const decorateSpy = jest.fn().mockResolvedValue(undefined); +const initializeSpy = jest.fn().mockResolvedValue(undefined); +const finalizeSpy = jest.fn().mockResolvedValue(undefined); + +class ConcreteDecorator extends DecoratorBase { + public initialize(): Promise { + return initializeSpy(); + } + public decorate( + document: IndexableDocument, + ): Promise { + return decorateSpy(document); + } + public finalize(): Promise { + return finalizeSpy(); + } +} + +describe('DecoratorBase', () => { + const document = { + title: 'Some Document', + text: 'Some document text.', + location: '/some/location', + }; + + afterEach(() => { + jest.resetAllMocks(); + }); + + it('should work end-to-end', async () => { + decorateSpy.mockImplementation(doc => ({ + ...doc, + transformed: true, + })); + + const decorator = new ConcreteDecorator(); + const { documents } = await TestPipeline.withSubject(decorator) + .withDocuments([document, document, document]) + .execute(); + + expect(documents.length).toBe(3); + expect((documents[0] as unknown as any).transformed).toBe(true); + expect((documents[1] as unknown as any).transformed).toBe(true); + expect((documents[2] as unknown as any).transformed).toBe(true); + }); + + it('should allow filtering', async () => { + decorateSpy.mockResolvedValue(undefined); + + const decorator = new ConcreteDecorator(); + const { documents } = await TestPipeline.withSubject(decorator) + .withDocuments([document, document, document]) + .execute(); + + expect(decorateSpy).toHaveBeenCalledTimes(3); + expect(documents.length).toBe(0); + }); + + it('should allow fanning', async () => { + decorateSpy.mockImplementation(doc => { + return [doc, doc]; + }); + + const decorator = new ConcreteDecorator(); + const { documents } = await TestPipeline.withSubject(decorator) + .withDocuments([document, document, document]) + .execute(); + + expect(decorateSpy).toHaveBeenCalledTimes(3); + expect(documents.length).toBe(6); + }); + + it('should call initialize at construction', () => { + // @ts-expect-error + const _indexer = new ConcreteDecorator(); + + return new Promise(done => { + // Allow initialization to complete. + setImmediate(() => { + expect(initializeSpy).toHaveBeenCalled(); + done(); + }); + }); + }); + + it('should emit error if initialization throws', () => { + // Cause the initializer to throw. + const expectedError = new Error('some error'); + initializeSpy.mockRejectedValue(expectedError); + const decorator = new ConcreteDecorator(); + + return new Promise(done => { + // Listen for the error and assert it's what was thrown. + decorator.on('error', error => { + expect(error).toStrictEqual(expectedError); + done(); + }); + + // Write a document to force the error state to become known. + decorator.write(document); + }); + }); + + it('should emit error if index throws', () => { + // Cause the indexer to throw. + const expectedError = new Error('decorate error'); + decorateSpy.mockRejectedValue(expectedError); + const decorator = new ConcreteDecorator(); + + return new Promise(done => { + // Listen for the error and assert it's what was thrown. + decorator.on('error', error => { + expect(error).toStrictEqual(expectedError); + done(); + }); + + decorator.write(document); + }); + }); + + it('should emit error if finalize throws', () => { + // Cause the indexer to throw. + const expectedError = new Error('finalize error'); + finalizeSpy.mockRejectedValue(expectedError); + const decorator = new ConcreteDecorator(); + + return new Promise(done => { + // Listen for the error and assert it's what was thrown. + decorator.on('error', error => { + expect(error).toStrictEqual(expectedError); + done(); + }); + + decorator.end(); + }); + }); +}); diff --git a/plugins/search-backend-node/src/indexing/DecoratorBase.ts b/plugins/search-backend-node/src/indexing/DecoratorBase.ts new file mode 100644 index 0000000000..a28d652fd7 --- /dev/null +++ b/plugins/search-backend-node/src/indexing/DecoratorBase.ts @@ -0,0 +1,126 @@ +/* + * 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 { assertError } from '@backstage/errors'; +import { IndexableDocument } from '@backstage/search-common'; +import { Transform } from 'stream'; + +/** + * Base class encapsulating simple async transformations. Useful as a base + * class for Backstage search decorators. + */ +export abstract class DecoratorBase extends Transform { + private initialized: Promise; + + constructor() { + super({ objectMode: true }); + + // @todo Once node v15 is minimum, convert to _construct implementation. + this.initialized = new Promise(done => { + // Necessary to allow concrete implementation classes to construct + // themselves before calling their initialize() methods. + setImmediate(async () => { + try { + await this.initialize(); + done(undefined); + } catch (e) { + assertError(e); + done(e); + } + }); + }); + } + + /** + * Any asynchronous setup tasks can be performed here. + */ + public abstract initialize(): Promise; + + /** + * Receives a single indexable document. In your decorate method, you can: + * + * - Resolve `undefined` to indicate the record should be omitted. + * - Resolve a single modified document, which could contain new fields, + * edited fields, or removed fields. + * - Resolve an array of indexable documents, if the purpose if the decorator + * is to convert one document into multiple derivative documents. + */ + public abstract decorate( + document: IndexableDocument, + ): Promise; + + /** + * Any asynchronous teardown tasks can be performed here. + */ + public abstract finalize(): Promise; + + /** + * Encapsulates simple transform stream logic. + * @internal + */ + async _transform( + document: IndexableDocument, + _: any, + done: (error?: Error | null) => void, + ) { + // Wait for init before proceeding. Throw error if initialization failed. + const maybeError = await this.initialized; + if (maybeError) { + done(maybeError); + return; + } + + try { + const decorated = await this.decorate(document); + + // If undefined was returned, omit the record and move on. + if (decorated === undefined) { + done(); + return; + } + + // If an array of documents was given, push them all. + if (Array.isArray(decorated)) { + decorated.forEach(doc => { + this.push(doc); + }); + done(); + return; + } + + // Otherwise, just push the decorated document. + this.push(decorated); + done(); + } catch (e) { + assertError(e); + done(e); + } + } + + /** + * Encapsulates finalization and final error handling logic. + * @internal + */ + async _final(done: (error?: Error | null) => void) { + try { + await this.finalize(); + done(); + } catch (e) { + assertError(e); + done(e); + } + } +} diff --git a/plugins/search-backend-node/src/indexing/index.ts b/plugins/search-backend-node/src/indexing/index.ts new file mode 100644 index 0000000000..e86235df16 --- /dev/null +++ b/plugins/search-backend-node/src/indexing/index.ts @@ -0,0 +1,19 @@ +/* + * 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. + */ + +export { BatchSearchEngineIndexer } from './BatchSearchEngineIndexer'; +export { DecoratorBase } from './DecoratorBase'; +export type { BatchSearchEngineOptions } from './BatchSearchEngineIndexer';