Update PGSearchEngine to be stream-based

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
2022-02-26 18:50:50 +01:00
parent a151cf2a88
commit 0d995e0ddd
8 changed files with 288 additions and 55 deletions
@@ -3,6 +3,7 @@
 > Do not edit this file. It is a report generated by [API Extractor](https://api-extractor.com/).

 ```ts
+import { BatchSearchEngineIndexer } from '@backstage/plugin-search-backend-node';
 import { IndexableDocument } from '@backstage/search-common';
 import { Knex } from 'knex';
 import { PluginDatabaseManager } from '@backstage/backend-common';
@@ -28,6 +29,8 @@ export class DatabaseDocumentStore implements DatabaseStore {
  // (undocumented)
  static create(knex: Knex): Promise<DatabaseDocumentStore>;
  // (undocumented)
+  getTransaction(): Promise<Knex.Transaction>;
+  // (undocumented)
  insertDocuments(
    tx: Knex.Transaction,
    type: string,
@@ -55,6 +58,8 @@ export interface DatabaseStore {
  // (undocumented)
  completeInsert(tx: Knex.Transaction, type: string): Promise<void>;
  // (undocumented)
+  getTransaction(): Promise<Knex.Transaction>;
+  // (undocumented)
  insertDocuments(
    tx: Knex.Transaction,
    type: string,
@@ -81,7 +86,7 @@ export class PgSearchEngine implements SearchEngine {
    database: PluginDatabaseManager;
  }): Promise<PgSearchEngine>;
  // (undocumented)
-  index(type: string, documents: IndexableDocument[]): Promise<void>;
+  getIndexer(type: string): Promise<PgSearchEngineIndexer>;
  // (undocumented)
  query(query: SearchQuery): Promise<SearchResultSet>;
  // (undocumented)
@@ -94,6 +99,28 @@ export class PgSearchEngine implements SearchEngine {
  translator(query: SearchQuery): ConcretePgSearchQuery;
 }

+// Warning: (ae-missing-release-tag) "PgSearchEngineIndexer" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal)
+//
+// @public (undocumented)
+export class PgSearchEngineIndexer extends BatchSearchEngineIndexer {
+  constructor(options: PgSearchEngineIndexerOptions);
+  // (undocumented)
+  finalize(): Promise<void>;
+  // (undocumented)
+  index(documents: IndexableDocument[]): Promise<void>;
+  // (undocumented)
+  initialize(): Promise<void>;
+}
+
+// Warning: (ae-missing-release-tag) "PgSearchEngineIndexerOptions" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal)
+//
+// @public (undocumented)
+export type PgSearchEngineIndexerOptions = {
+  batchSize: number;
+  type: string;
+  databaseStore: DatabaseStore;
+};
+
 // Warning: (ae-missing-release-tag) "PgSearchQuery" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal)
 //
 // @public (undocumented)
@@ -13,7 +13,6 @@
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
-import { range } from 'lodash';
 import { DatabaseStore } from '../database';
 import {
  ConcretePgSearchQuery,
@@ -21,6 +20,13 @@ import {
  encodePageCursor,
  PgSearchEngine,
 } from './PgSearchEngine';
+import { PgSearchEngineIndexer } from './PgSearchEngineIndexer';
+
+jest.mock('./PgSearchEngineIndexer', () => ({
+  PgSearchEngineIndexer: jest
+    .fn()
+    .mockImplementation(async () => 'the-expected-indexer'),
+}));

 describe('PgSearchEngine', () => {
  const tx: any = {} as any;
@@ -30,6 +36,7 @@ describe('PgSearchEngine', () => {
  beforeEach(() => {
    database = {
      transaction: jest.fn(),
+      getTransaction: jest.fn(),
      insertDocuments: jest.fn(),
      query: jest.fn(),
      completeInsert: jest.fn(),
@@ -122,46 +129,21 @@ describe('PgSearchEngine', () => {
    });
  });

-  describe('insert', () => {
-    it('should insert documents', async () => {
-      database.transaction.mockImplementation(fn => fn(tx));
+  describe('index', () => {
+    it('should instantiate indexer', async () => {
+      const indexer = await searchEngine.getIndexer('my-type');

-      const documents = [
-        { title: 'Hello World', text: 'Lorem Ipsum', location: 'location-1' },
-        {
-          location: 'location-2',
-          text: 'Hello World',
-          title: 'Dolor sit amet',
-        },
-      ];
-
-      await searchEngine.index('my-type', documents);
-
-      expect(database.transaction).toHaveBeenCalledTimes(1);
-      expect(database.prepareInsert).toHaveBeenCalledTimes(1);
-      expect(database.insertDocuments).toHaveBeenCalledWith(
-        tx,
-        'my-type',
-        documents,
+      // Indexer instantiated with expected args.
+      expect(PgSearchEngineIndexer).toHaveBeenCalledWith(
+        expect.objectContaining({
+          batchSize: 100,
+          type: 'my-type',
+          databaseStore: database,
+        }),
      );
-      expect(database.completeInsert).toHaveBeenCalledWith(tx, 'my-type');
-    });

-    it('should batch insert documents', async () => {
-      database.transaction.mockImplementation(fn => fn(tx));
-
-      const documents = range(350).map(i => ({
-        title: `Hello World ${i}`,
-        text: 'Lorem Ipsum',
-        location: `location-${i}`,
-      }));
-
-      await searchEngine.index('my-type', documents);
-
-      expect(database.transaction).toHaveBeenCalledTimes(1);
-      expect(database.prepareInsert).toHaveBeenCalledTimes(1);
-      expect(database.insertDocuments).toBeCalledTimes(4);
-      expect(database.completeInsert).toHaveBeenCalledWith(tx, 'my-type');
+      // Indexer is as expected.
+      expect(indexer).toBe('the-expected-indexer');
    });
  });

@@ -15,12 +15,8 @@
 */
 import { PluginDatabaseManager } from '@backstage/backend-common';
 import { SearchEngine } from '@backstage/plugin-search-backend-node';
-import {
-  IndexableDocument,
-  SearchQuery,
-  SearchResultSet,
-} from '@backstage/search-common';
-import { chunk } from 'lodash';
+import { SearchQuery, SearchResultSet } from '@backstage/search-common';
+import { PgSearchEngineIndexer } from './PgSearchEngineIndexer';
 import {
  DatabaseDocumentStore,
  DatabaseStore,
@@ -77,16 +73,11 @@ export class PgSearchEngine implements SearchEngine {
    this.translator = translator;
  }

-  async index(type: string, documents: IndexableDocument[]): Promise<void> {
-    await this.databaseStore.transaction(async tx => {
-      await this.databaseStore.prepareInsert(tx);
-
-      const batchSize = 100;
-      for (const documentBatch of chunk(documents, batchSize)) {
-        await this.databaseStore.insertDocuments(tx, type, documentBatch);
-      }
-
-      await this.databaseStore.completeInsert(tx, type);
+  async getIndexer(type: string) {
+    return new PgSearchEngineIndexer({
+      batchSize: 100,
+      type,
+      databaseStore: this.databaseStore,
    });
  }

@@ -0,0 +1,150 @@
+/*
+ * 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 { TestPipeline } from '@backstage/plugin-search-backend-node';
+import { range } from 'lodash';
+import { PgSearchEngineIndexer } from './PgSearchEngineIndexer';
+import { DatabaseStore } from '../database';
+
+describe('PgSearchEngineIndexer', () => {
+  const tx = {
+    rollback: jest.fn(),
+    commit: jest.fn(),
+  } as any;
+  let database: jest.Mocked<DatabaseStore>;
+  let indexer: PgSearchEngineIndexer;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    database = {
+      transaction: jest.fn().mockImplementation(fn => fn(tx)),
+      getTransaction: jest.fn().mockReturnValue(tx),
+      insertDocuments: jest.fn(),
+      query: jest.fn(),
+      completeInsert: jest.fn(),
+      prepareInsert: jest.fn(),
+    };
+    indexer = new PgSearchEngineIndexer({
+      batchSize: 100,
+      type: 'my-type',
+      databaseStore: database,
+    });
+  });
+
+  it('should insert documents', async () => {
+    const documents = [
+      { title: 'Hello World', text: 'Lorem Ipsum', location: 'location-1' },
+      {
+        location: 'location-2',
+        text: 'Hello World',
+        title: 'Dolor sit amet',
+      },
+    ];
+
+    await TestPipeline.withSubject(indexer).withDocuments(documents).execute();
+
+    expect(database.getTransaction).toHaveBeenCalledTimes(1);
+    expect(database.prepareInsert).toHaveBeenCalledTimes(1);
+    expect(database.insertDocuments).toHaveBeenCalledWith(
+      tx,
+      'my-type',
+      documents,
+    );
+    expect(database.completeInsert).toHaveBeenCalledWith(tx, 'my-type');
+    expect(tx.commit).toHaveBeenCalled();
+  });
+
+  it('should batch insert documents', 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(database.getTransaction).toHaveBeenCalledTimes(1);
+    expect(database.prepareInsert).toHaveBeenCalledTimes(1);
+    expect(database.insertDocuments).toBeCalledTimes(4);
+    expect(database.completeInsert).toHaveBeenCalledWith(tx, 'my-type');
+  });
+
+  it('should close out stream and bubble up error on prepare', async () => {
+    const expectedError = new Error('Prepare error');
+    const documents = [
+      {
+        title: `Hello World`,
+        text: 'Lorem Ipsum',
+        location: `location`,
+      },
+    ];
+
+    database.prepareInsert.mockRejectedValueOnce(expectedError);
+    const result = await TestPipeline.withSubject(indexer)
+      .withDocuments(documents)
+      .execute();
+
+    expect(database.getTransaction).toHaveBeenCalledTimes(1);
+    expect(database.insertDocuments).not.toHaveBeenCalled();
+    expect(database.completeInsert).not.toHaveBeenCalled();
+    expect(result.error).toBe(expectedError);
+    expect(tx.rollback).toHaveBeenCalledWith(expectedError);
+  });
+
+  it('should close tx and bubble up error on insert', async () => {
+    const expectedError = new Error('Index error');
+    const documents = [
+      {
+        title: `Hello World`,
+        text: 'Lorem Ipsum',
+        location: `location`,
+      },
+    ];
+
+    database.insertDocuments.mockRejectedValueOnce(expectedError);
+    const result = await TestPipeline.withSubject(indexer)
+      .withDocuments(documents)
+      .execute();
+
+    expect(database.getTransaction).toHaveBeenCalledTimes(1);
+    expect(database.prepareInsert).toHaveBeenCalledTimes(1);
+    expect(database.completeInsert).not.toHaveBeenCalled();
+    expect(result.error).toBe(expectedError);
+    expect(tx.rollback).toHaveBeenCalledWith(expectedError);
+  });
+
+  it('should close tx and bubble up error on completion', async () => {
+    const expectedError = new Error('Completion error');
+    const documents = [
+      {
+        title: `Hello World`,
+        text: 'Lorem Ipsum',
+        location: `location`,
+      },
+    ];
+
+    database.completeInsert.mockRejectedValueOnce(expectedError);
+    const result = await TestPipeline.withSubject(indexer)
+      .withDocuments(documents)
+      .execute();
+
+    expect(database.getTransaction).toHaveBeenCalledTimes(1);
+    expect(database.prepareInsert).toHaveBeenCalledTimes(1);
+    expect(database.insertDocuments).toHaveBeenCalledTimes(1);
+    expect(database.completeInsert).toHaveBeenCalledTimes(1);
+    expect(result.error).toBe(expectedError);
+    expect(tx.rollback).toHaveBeenCalledWith(expectedError);
+  });
+});
@@ -0,0 +1,74 @@
+/*
+ * 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 { BatchSearchEngineIndexer } from '@backstage/plugin-search-backend-node';
+import { IndexableDocument } from '@backstage/search-common';
+import { Knex } from 'knex';
+import { DatabaseStore } from '../database';
+
+export type PgSearchEngineIndexerOptions = {
+  batchSize: number;
+  type: string;
+  databaseStore: DatabaseStore;
+};
+
+export class PgSearchEngineIndexer extends BatchSearchEngineIndexer {
+  private store: DatabaseStore;
+  private type: string;
+  private tx: Knex.Transaction | undefined;
+
+  constructor(options: PgSearchEngineIndexerOptions) {
+    super({ batchSize: options.batchSize });
+    this.store = options.databaseStore;
+    this.type = options.type;
+  }
+
+  async initialize(): Promise<void> {
+    this.tx = await this.store.getTransaction();
+    try {
+      await this.store.prepareInsert(this.tx);
+    } catch (e) {
+      // In case of error, rollback the transaction and re-throw the error so
+      // that the stream can be closed and destroyed properly.
+      this.tx.rollback(e);
+      throw e;
+    }
+  }
+
+  async index(documents: IndexableDocument[]): Promise<void> {
+    try {
+      await this.store.insertDocuments(this.tx!, this.type, documents);
+    } catch (e) {
+      // In case of error, rollback the transaction and re-throw the error so
+      // that the stream can be closed and destroyed properly.
+      this.tx!.rollback(e);
+      throw e;
+    }
+  }
+
+  async finalize(): Promise<void> {
+    // Attempt to complete and commit the transaction.
+    try {
+      await this.store.completeInsert(this.tx!, this.type);
+      this.tx!.commit();
+    } catch (e) {
+      // Otherwise, rollback the transaction and re-throw the error so that the
+      // stream can be closed and destroyed properly.
+      this.tx!.rollback!(e);
+      throw e;
+    }
+  }
+}
@@ -15,3 +15,7 @@
 */
 export { PgSearchEngine } from './PgSearchEngine';
 export type { ConcretePgSearchQuery } from './PgSearchEngine';
+export type {
+  PgSearchEngineIndexer,
+  PgSearchEngineIndexerOptions,
+} from './PgSearchEngineIndexer';
@@ -71,6 +71,10 @@ export class DatabaseDocumentStore implements DatabaseStore {
    return await this.db.transaction(fn);
  }

+  async getTransaction(): Promise<Knex.Transaction> {
+    return this.db.transaction();
+  }
+
  async prepareInsert(tx: Knex.Transaction): Promise<void> {
    // We create a temporary table to collect the hashes of the documents that
    // we expect to be in the documents table at the end. The table is deleted
@@ -26,6 +26,7 @@ export interface PgSearchQuery {

 export interface DatabaseStore {
  transaction<T>(fn: (tx: Knex.Transaction) => Promise<T>): Promise<T>;
+  getTransaction(): Promise<Knex.Transaction>;
  prepareInsert(tx: Knex.Transaction): Promise<void>;
  insertDocuments(
    tx: Knex.Transaction,