diff --git a/plugins/search-backend-module-pg/api-report.md b/plugins/search-backend-module-pg/api-report.md index e2edfc291a..ba8fde5f17 100644 --- a/plugins/search-backend-module-pg/api-report.md +++ b/plugins/search-backend-module-pg/api-report.md @@ -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; // (undocumented) + getTransaction(): Promise; + // (undocumented) insertDocuments( tx: Knex.Transaction, type: string, @@ -55,6 +58,8 @@ export interface DatabaseStore { // (undocumented) completeInsert(tx: Knex.Transaction, type: string): Promise; // (undocumented) + getTransaction(): Promise; + // (undocumented) insertDocuments( tx: Knex.Transaction, type: string, @@ -81,7 +86,7 @@ export class PgSearchEngine implements SearchEngine { database: PluginDatabaseManager; }): Promise; // (undocumented) - index(type: string, documents: IndexableDocument[]): Promise; + getIndexer(type: string): Promise; // (undocumented) query(query: SearchQuery): Promise; // (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; + // (undocumented) + index(documents: IndexableDocument[]): Promise; + // (undocumented) + initialize(): Promise; +} + +// 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) diff --git a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngine.test.ts b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngine.test.ts index 56a3a6e643..8618bed4c8 100644 --- a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngine.test.ts +++ b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngine.test.ts @@ -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'); }); }); diff --git a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngine.ts b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngine.ts index 6fd6c571a8..ac42007401 100644 --- a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngine.ts +++ b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngine.ts @@ -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 { - 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, }); } diff --git a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts new file mode 100644 index 0000000000..1fc7e74fc4 --- /dev/null +++ b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts @@ -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; + 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); + }); +}); diff --git a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts new file mode 100644 index 0000000000..53d040cb87 --- /dev/null +++ b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts @@ -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 { + 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 { + 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 { + // 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; + } + } +} diff --git a/plugins/search-backend-module-pg/src/PgSearchEngine/index.ts b/plugins/search-backend-module-pg/src/PgSearchEngine/index.ts index 7994998baf..7f8e297648 100644 --- a/plugins/search-backend-module-pg/src/PgSearchEngine/index.ts +++ b/plugins/search-backend-module-pg/src/PgSearchEngine/index.ts @@ -15,3 +15,7 @@ */ export { PgSearchEngine } from './PgSearchEngine'; export type { ConcretePgSearchQuery } from './PgSearchEngine'; +export type { + PgSearchEngineIndexer, + PgSearchEngineIndexerOptions, +} from './PgSearchEngineIndexer'; diff --git a/plugins/search-backend-module-pg/src/database/DatabaseDocumentStore.ts b/plugins/search-backend-module-pg/src/database/DatabaseDocumentStore.ts index 8c180ea019..0d3ee63ca4 100644 --- a/plugins/search-backend-module-pg/src/database/DatabaseDocumentStore.ts +++ b/plugins/search-backend-module-pg/src/database/DatabaseDocumentStore.ts @@ -71,6 +71,10 @@ export class DatabaseDocumentStore implements DatabaseStore { return await this.db.transaction(fn); } + async getTransaction(): Promise { + return this.db.transaction(); + } + async prepareInsert(tx: Knex.Transaction): Promise { // 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 diff --git a/plugins/search-backend-module-pg/src/database/types.ts b/plugins/search-backend-module-pg/src/database/types.ts index 0c0596160e..0a0dc35682 100644 --- a/plugins/search-backend-module-pg/src/database/types.ts +++ b/plugins/search-backend-module-pg/src/database/types.ts @@ -26,6 +26,7 @@ export interface PgSearchQuery { export interface DatabaseStore { transaction(fn: (tx: Knex.Transaction) => Promise): Promise; + getTransaction(): Promise; prepareInsert(tx: Knex.Transaction): Promise; insertDocuments( tx: Knex.Transaction,