Introduce a base classes to simplify stream-based implementations

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>

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<void> {
+    return indexSpy(documents);
+  }
+  async initialize(): Promise<void> {
+    return initializeSpy();
+  }
+  async finalize(): Promise<void> {
+    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<void>(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<void>(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<void>(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<void>(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<void>(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<void>(done => {
+      // Listen for the error and assert it's what was thrown.
+      indexer.on('error', error => {
+        expect(error).toStrictEqual(expectedError);
+        done();
+      });
+
+      indexer.end();
+    });
+  });
+});

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<undefined | Error>;
+
+  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<void>;
+
+  /**
+   * Any asynchronous setup tasks can be performed here.
+   */
+  public abstract initialize(): Promise<void>;
+
+  /**
+   * Any asynchronous teardown tasks can be performed here.
+   */
+  public abstract finalize(): Promise<void>;
+
+  /**
+   * 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);
+    }
+  }
+}

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<void> {
+    return initializeSpy();
+  }
+  public decorate(
+    document: IndexableDocument,
+  ): Promise<IndexableDocument | IndexableDocument[] | undefined> {
+    return decorateSpy(document);
+  }
+  public finalize(): Promise<void> {
+    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<void>(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<void>(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<void>(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<void>(done => {
+      // Listen for the error and assert it's what was thrown.
+      decorator.on('error', error => {
+        expect(error).toStrictEqual(expectedError);
+        done();
+      });
+
+      decorator.end();
+    });
+  });
+});

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<undefined | Error>;
+
+  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<void>;
+
+  /**
+   * 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<IndexableDocument | IndexableDocument[] | undefined>;
+
+  /**
+   * Any asynchronous teardown tasks can be performed here.
+   */
+  public abstract finalize(): Promise<void>;
+
+  /**
+   * 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);
+    }
+  }
+}

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';