diff --git a/.changeset/wicked-nails-taste.md b/.changeset/wicked-nails-taste.md new file mode 100644 index 0000000000..880bb9c578 --- /dev/null +++ b/.changeset/wicked-nails-taste.md @@ -0,0 +1,5 @@ +--- +'@backstage/plugin-search-backend-module-pg': patch +--- + +Fix a bug in large document indexing logic by using sub-transaction rollbacks diff --git a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts index fd9e4de34a..7a83a8d563 100644 --- a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts +++ b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts @@ -20,10 +20,18 @@ import { PgSearchEngineIndexer } from './PgSearchEngineIndexer'; import { DatabaseStore } from '../database'; describe('PgSearchEngineIndexer', () => { + const subTx = { + rollback: jest.fn(), + commit: jest.fn(), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any; + const tx = { rollback: jest.fn(), commit: jest.fn(), + transaction: jest.fn().mockResolvedValue(subTx), isCompleted: jest.fn(), + // eslint-disable-next-line @typescript-eslint/no-explicit-any } as any; let database: jest.Mocked; let indexer: PgSearchEngineIndexer; @@ -65,8 +73,14 @@ describe('PgSearchEngineIndexer', () => { 'my-type', documents, ); - expect(database.completeInsert).toHaveBeenCalledWith(tx, 'my-type', false); + expect(database.completeInsert).toHaveBeenCalledWith( + subTx, + 'my-type', + false, + ); + expect(subTx.commit).toHaveBeenCalled(); expect(tx.commit).toHaveBeenCalled(); + expect(subTx.rollback).not.toHaveBeenCalled(); expect(tx.rollback).not.toHaveBeenCalled(); }); @@ -95,18 +109,20 @@ describe('PgSearchEngineIndexer', () => { expect(database.completeInsert).toHaveBeenCalledTimes(2); expect(database.completeInsert).toHaveBeenNthCalledWith( 1, - tx, + subTx, 'my-type', false, ); expect(database.completeInsert).toHaveBeenNthCalledWith( 2, - tx, + subTx, 'my-type', true, ); + expect(subTx.commit).toHaveBeenCalled(); expect(tx.commit).toHaveBeenCalled(); - expect(tx.rollback).toHaveBeenCalledTimes(1); + expect(subTx.rollback).toHaveBeenCalledTimes(1); + expect(tx.rollback).not.toHaveBeenCalled(); }); it('should batch insert documents', async () => { @@ -121,7 +137,11 @@ describe('PgSearchEngineIndexer', () => { expect(database.getTransaction).toHaveBeenCalledTimes(1); expect(database.prepareInsert).toHaveBeenCalledTimes(1); expect(database.insertDocuments).toHaveBeenCalledTimes(4); - expect(database.completeInsert).toHaveBeenCalledWith(tx, 'my-type', false); + expect(database.completeInsert).toHaveBeenCalledWith( + subTx, + 'my-type', + false, + ); }); it('should rollback transaction if no documents indexed', async () => { diff --git a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts index ad8c4f093f..f601ed2668 100644 --- a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts +++ b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts @@ -89,15 +89,28 @@ export class PgSearchEngineIndexer extends BatchSearchEngineIndexer { let retryTruncated = false; do { + // Sub-transaction allows a partial rollback if a truncated retry is needed + let subTx: Knex.Transaction; + try { + subTx = await this.tx!.transaction(); + } catch (e) { + // If sub-transaction creation fails, rollback the outer transaction + // and re-throw the error so that the stream can be closed and + // destroyed properly. + this.tx!.rollback!(e); + throw e; + } + // Attempt to complete and commit the transaction. try { - await this.store.completeInsert(this.tx!, this.type, retryTruncated); - this.tx!.commit(); + await this.store.completeInsert(subTx, this.type, retryTruncated); + + // Commit both transactions if successful + await subTx.commit(); + await this.tx!.commit(); retryTruncated = false; } 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); + await subTx.rollback(e); // Rollback the first completeInsert attempt if ( e instanceof Error && @@ -108,6 +121,9 @@ export class PgSearchEngineIndexer extends BatchSearchEngineIndexer { // retry the operation with truncated text. retryTruncated = true; } else { + // Otherwise, rollback the outer transaction and re-throw the error so that the + // stream can be closed and destroyed properly. + this.tx!.rollback!(e); throw e; } }