From 2f5162c4b783598e4ff7a1df494a1617edcc302a Mon Sep 17 00:00:00 2001 From: Hellgren Heikki Date: Fri, 29 Aug 2025 08:30:24 +0300 Subject: [PATCH] feat(catalog): batch size + mode config for entity streaming adds parameters to control the batch size and mode for the streamEntities method. with array mode, the entities are returned as batches while the default single mode returns entities one by one. Signed-off-by: Hellgren Heikki --- .../catalog-client/report-testUtils.api.md | 9 ++++++- packages/catalog-client/report.api.md | 19 +++++++++----- .../catalog-client/src/CatalogClient.test.ts | 13 ++++++++-- packages/catalog-client/src/CatalogClient.ts | 25 +++++++++++-------- .../testUtils/InMemoryCatalogClient.test.ts | 12 ++++++++- .../src/testUtils/InMemoryCatalogClient.ts | 17 +++++++++---- packages/catalog-client/src/types/api.ts | 13 ++++++++-- plugins/catalog-node/report-testUtils.api.md | 2 +- plugins/catalog-node/report.api.md | 2 +- plugins/catalog-node/src/catalogService.ts | 4 +-- plugins/catalog-node/src/testUtils/types.ts | 2 +- 11 files changed, 86 insertions(+), 32 deletions(-) diff --git a/packages/catalog-client/report-testUtils.api.md b/packages/catalog-client/report-testUtils.api.md index 26e79c9f96..cadf58aeaf 100644 --- a/packages/catalog-client/report-testUtils.api.md +++ b/packages/catalog-client/report-testUtils.api.md @@ -71,7 +71,14 @@ export class InMemoryCatalogClient implements CatalogApi { // (undocumented) removeLocationById(_id: string): Promise; // (undocumented) - streamEntities(request?: StreamEntitiesRequest): AsyncIterable; + streamEntities< + T extends StreamEntitiesRequest, + R = T extends { + mode: 'array'; + } & StreamEntitiesRequest + ? Entity[] + : Entity, + >(request?: T): AsyncIterable; // (undocumented) validateEntity( _entity: Entity, diff --git a/packages/catalog-client/report.api.md b/packages/catalog-client/report.api.md index cda0d0c99d..6000686fbd 100644 --- a/packages/catalog-client/report.api.md +++ b/packages/catalog-client/report.api.md @@ -91,7 +91,7 @@ export interface CatalogApi { streamEntities( request?: StreamEntitiesRequest, options?: CatalogRequestOptions, - ): AsyncIterable; + ): AsyncIterable; validateEntity( entity: Entity, locationRef: string, @@ -174,10 +174,14 @@ export class CatalogClient implements CatalogApi { id: string, options?: CatalogRequestOptions, ): Promise; - streamEntities( - request?: StreamEntitiesRequest, - options?: CatalogRequestOptions, - ): AsyncIterable; + streamEntities< + T extends StreamEntitiesRequest, + R = T extends { + mode: 'array'; + } & StreamEntitiesRequest + ? Entity[] + : Entity, + >(request?: T, options?: CatalogRequestOptions): AsyncIterable; validateEntity( entity: Entity, locationRef: string, @@ -329,7 +333,10 @@ export type QueryEntitiesResponse = { export type StreamEntitiesRequest = Omit< QueryEntitiesRequest, 'limit' | 'offset' ->; +> & { + batchSize?: number; + mode?: 'single' | 'array'; +}; // @public export type ValidateEntityResponse = diff --git a/packages/catalog-client/src/CatalogClient.test.ts b/packages/catalog-client/src/CatalogClient.test.ts index 325f60cc13..c5e2c8cc21 100644 --- a/packages/catalog-client/src/CatalogClient.test.ts +++ b/packages/catalog-client/src/CatalogClient.test.ts @@ -579,7 +579,7 @@ describe('CatalogClient', () => { ); }); - it('should stream entities', async () => { + it('should stream entities one by one', async () => { const stream = client.streamEntities({}, { token }); const results: Entity[] = []; for await (const entity of stream) { @@ -588,6 +588,15 @@ describe('CatalogClient', () => { expect(results).toEqual(defaultResponse.items); }); + it('should stream entities in batches', async () => { + const stream = client.streamEntities({ mode: 'array' }, { token }); + const results: Entity[] = []; + for await (const entityBatch of stream) { + results.push(...entityBatch); + } + expect(results).toEqual(defaultResponse.items); + }); + it('should handle errors', async () => { const mockedEndpoint = jest .fn() @@ -597,7 +606,7 @@ describe('CatalogClient', () => { const stream = client.streamEntities({}, { token }); await expect(async () => { - const results: Entity[] = []; + const results = []; for await (const entity of stream) { results.push(entity); } diff --git a/packages/catalog-client/src/CatalogClient.ts b/packages/catalog-client/src/CatalogClient.ts index a412c790f6..b090966807 100644 --- a/packages/catalog-client/src/CatalogClient.ts +++ b/packages/catalog-client/src/CatalogClient.ts @@ -51,7 +51,7 @@ import type { } from '@backstage/plugin-catalog-common'; // Number of entities to return in a single streamEntities request -const STREAM_ENTITIES_LIMIT = 500; +const DEFAULT_STREAM_ENTITIES_LIMIT = 500; /** * A frontend and backend compatible client for communicating with the Backstage @@ -463,20 +463,25 @@ export class CatalogClient implements CatalogApi { /** * {@inheritdoc CatalogApi.streamEntities} */ - async *streamEntities( - request?: StreamEntitiesRequest, - options?: CatalogRequestOptions, - ): AsyncIterable { + async *streamEntities< + T extends StreamEntitiesRequest, + R = T extends { mode: 'array' } & StreamEntitiesRequest ? Entity[] : Entity, + >(request?: T, options?: CatalogRequestOptions): AsyncIterable { let cursor: string | undefined = undefined; + const limit = request?.batchSize ?? DEFAULT_STREAM_ENTITIES_LIMIT; + const mode = request?.mode ?? 'single'; do { const res = await this.queryEntities( - cursor - ? { ...request, cursor, limit: STREAM_ENTITIES_LIMIT } - : { ...request, limit: STREAM_ENTITIES_LIMIT }, + cursor ? { ...request, cursor, limit } : { ...request, limit }, options, ); - for (const entity of res.items) { - yield entity; + + if (mode === 'single') { + for (const entity of res.items) { + yield entity as R; + } + } else { + yield res.items as R; } cursor = res.pageInfo.nextCursor; diff --git a/packages/catalog-client/src/testUtils/InMemoryCatalogClient.test.ts b/packages/catalog-client/src/testUtils/InMemoryCatalogClient.test.ts index 5481c2ca3c..afc26346a7 100644 --- a/packages/catalog-client/src/testUtils/InMemoryCatalogClient.test.ts +++ b/packages/catalog-client/src/testUtils/InMemoryCatalogClient.test.ts @@ -123,7 +123,7 @@ describe('InMemoryCatalogClient', () => { }); }); - it('streamEntities', async () => { + it('streamEntities single', async () => { const client = new InMemoryCatalogClient({ entities }); const stream = client.streamEntities(); const results: Entity[] = []; @@ -133,6 +133,16 @@ describe('InMemoryCatalogClient', () => { expect(results).toEqual(entities); }); + it('streamEntities batch', async () => { + const client = new InMemoryCatalogClient({ entities }); + const stream = client.streamEntities({ mode: 'array' }); + const results: Entity[] = []; + for await (const entity of stream) { + results.push(...entity); + } + expect(results).toEqual(entities); + }); + it('getEntityAncestors', async () => { const client = new InMemoryCatalogClient({ entities }); await expect( diff --git a/packages/catalog-client/src/testUtils/InMemoryCatalogClient.ts b/packages/catalog-client/src/testUtils/InMemoryCatalogClient.ts index a21709d798..fe23e7ceb1 100644 --- a/packages/catalog-client/src/testUtils/InMemoryCatalogClient.ts +++ b/packages/catalog-client/src/testUtils/InMemoryCatalogClient.ts @@ -280,16 +280,23 @@ export class InMemoryCatalogClient implements CatalogApi { throw new NotImplementedError('Method not implemented.'); } - async *streamEntities( - request?: StreamEntitiesRequest, - ): AsyncIterable { + async *streamEntities< + T extends StreamEntitiesRequest, + R = T extends { mode: 'array' } & StreamEntitiesRequest ? Entity[] : Entity, + >(request?: T): AsyncIterable { let cursor: string | undefined = undefined; + const mode = request?.mode ?? 'single'; do { const res = await this.queryEntities( cursor ? { ...request, cursor } : request, ); - for (const entity of res.items) { - yield entity; + + if (mode === 'single') { + for (const entity of res.items) { + yield entity as R; + } + } else { + yield res.items as R; } cursor = res.pageInfo.nextCursor; diff --git a/packages/catalog-client/src/types/api.ts b/packages/catalog-client/src/types/api.ts index 9eae527a33..fbf3d82a83 100644 --- a/packages/catalog-client/src/types/api.ts +++ b/packages/catalog-client/src/types/api.ts @@ -473,7 +473,16 @@ export type QueryEntitiesResponse = { export type StreamEntitiesRequest = Omit< QueryEntitiesRequest, 'limit' | 'offset' ->; +> & { + /** + * The number of entities to fetch in each batch. Defaults to 500. + */ + batchSize?: number; + /** + * The mode in which entities should be yielded. Defaults to 'single'. + */ + mode?: 'single' | 'array'; +}; /** * A client for interacting with the Backstage software catalog through its API. @@ -715,5 +724,5 @@ export interface CatalogApi { streamEntities( request?: StreamEntitiesRequest, options?: CatalogRequestOptions, - ): AsyncIterable; + ): AsyncIterable; } diff --git a/plugins/catalog-node/report-testUtils.api.md b/plugins/catalog-node/report-testUtils.api.md index 00f96ac347..4d0f7fdd48 100644 --- a/plugins/catalog-node/report-testUtils.api.md +++ b/plugins/catalog-node/report-testUtils.api.md @@ -111,7 +111,7 @@ export interface CatalogServiceMock extends CatalogService, CatalogApi { streamEntities( request?: StreamEntitiesRequest, options?: CatalogServiceRequestOptions | CatalogRequestOptions, - ): AsyncIterable; + ): AsyncIterable; // (undocumented) validateEntity( entity: Entity, diff --git a/plugins/catalog-node/report.api.md b/plugins/catalog-node/report.api.md index 5723666093..6e476f88f6 100644 --- a/plugins/catalog-node/report.api.md +++ b/plugins/catalog-node/report.api.md @@ -199,7 +199,7 @@ export interface CatalogService { streamEntities( request: StreamEntitiesRequest | undefined, options: CatalogServiceRequestOptions, - ): AsyncIterable; + ): AsyncIterable; // (undocumented) validateEntity( entity: Entity, diff --git a/plugins/catalog-node/src/catalogService.ts b/plugins/catalog-node/src/catalogService.ts index 666a97274e..d8509ddb20 100644 --- a/plugins/catalog-node/src/catalogService.ts +++ b/plugins/catalog-node/src/catalogService.ts @@ -146,7 +146,7 @@ export interface CatalogService { streamEntities( request: StreamEntitiesRequest | undefined, options: CatalogServiceRequestOptions, - ): AsyncIterable; + ): AsyncIterable; } class DefaultCatalogService implements CatalogService { @@ -329,7 +329,7 @@ class DefaultCatalogService implements CatalogService { async *streamEntities( request: StreamEntitiesRequest | undefined, options: CatalogServiceRequestOptions, - ): AsyncIterable { + ): AsyncIterable { yield* this.#catalogApi.streamEntities( request, await this.#getOptions(options), diff --git a/plugins/catalog-node/src/testUtils/types.ts b/plugins/catalog-node/src/testUtils/types.ts index 526527dcc9..fafdcb82fa 100644 --- a/plugins/catalog-node/src/testUtils/types.ts +++ b/plugins/catalog-node/src/testUtils/types.ts @@ -140,5 +140,5 @@ export interface CatalogServiceMock extends CatalogService, CatalogApi { streamEntities( request?: StreamEntitiesRequest, options?: CatalogServiceRequestOptions | CatalogRequestOptions, - ): AsyncIterable; + ): AsyncIterable; }