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