Update Catalog Collator to be stream-based

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>

plugins/catalog-backend/api-report.md

@@ -10,7 +10,7 @@ import { CatalogApi } from '@backstage/catalog-client';
 import { ConditionalPolicyDecision } from '@backstage/plugin-permission-node';
 import { Conditions } from '@backstage/plugin-permission-node';
 import { Config } from '@backstage/config';
-import { DocumentCollator } from '@backstage/search-common';
+import { DocumentCollatorFactory } from '@backstage/search-common';
 import { Entity } from '@backstage/catalog-model';
 import { EntityName } from '@backstage/catalog-model';
 import { EntityPolicy } from '@backstage/catalog-model';
@@ -30,6 +30,7 @@ import { PermissionCriteria } from '@backstage/plugin-permission-common';
 import { PermissionRule } from '@backstage/plugin-permission-node';
 import { PluginDatabaseManager } from '@backstage/backend-common';
 import { PluginEndpointDiscovery } from '@backstage/backend-common';
+import { Readable } from 'stream';
 import { Router } from 'express';
 import { ScmIntegrationRegistry } from '@backstage/integration';
 import { TokenManager } from '@backstage/backend-common';
@@ -415,8 +416,8 @@ export function createRandomRefreshInterval(options: {
 // @public
 export function createRouter(options: RouterOptions): Promise<express.Router>;
 
-// @public (undocumented)
-export class DefaultCatalogCollator implements DocumentCollator {
+// @public @deprecated (undocumented)
+export class DefaultCatalogCollator {
   constructor(options: {
     discovery: PluginEndpointDiscovery;
     tokenManager: TokenManager;
@@ -456,6 +457,31 @@ export class DefaultCatalogCollator implements DocumentCollator {
   readonly visibilityPermission: Permission;
 }
 
+// @public (undocumented)
+export class DefaultCatalogCollatorFactory implements DocumentCollatorFactory {
+  // (undocumented)
+  static fromConfig(
+    _config: Config,
+    options: DefaultCatalogCollatorFactoryOptions,
+  ): DefaultCatalogCollatorFactory;
+  // (undocumented)
+  getCollator(): Promise<Readable>;
+  // (undocumented)
+  readonly type: string;
+  // (undocumented)
+  readonly visibilityPermission: Permission;
+}
+
+// @public (undocumented)
+export type DefaultCatalogCollatorFactoryOptions = {
+  discovery: PluginEndpointDiscovery;
+  tokenManager: TokenManager;
+  locationTemplate?: string;
+  filter?: GetEntitiesRequest['filter'];
+  batchSize?: number;
+  catalogClient?: CatalogApi;
+};
+
 // @public (undocumented)
 export class DefaultCatalogProcessingOrchestrator
   implements CatalogProcessingOrchestrator

plugins/catalog-backend/package.json

@@ -73,6 +73,7 @@
     "@backstage/backend-test-utils": "^0.1.19",
     "@backstage/cli": "^0.14.1",
     "@backstage/plugin-permission-common": "^0.5.1",
+    "@backstage/plugin-search-backend-node": "0.4.7",
     "@backstage/test-utils": "^0.2.6",
     "@types/core-js": "^2.5.4",
     "@types/git-url-parse": "^9.0.0",

plugins/catalog-backend/src/search/DefaultCatalogCollator.ts

@@ -23,7 +23,6 @@ import {
   stringifyEntityRef,
   UserEntity,
 } from '@backstage/catalog-model';
-import { IndexableDocument, DocumentCollator } from '@backstage/search-common';
 import { Config } from '@backstage/config';
 import {
   CatalogApi,
@@ -31,18 +30,14 @@ import {
   GetEntitiesRequest,
 } from '@backstage/catalog-client';
 import { catalogEntityReadPermission } from '@backstage/plugin-catalog-common';
+import { CatalogEntityDocument } from './DefaultCatalogCollatorFactory';
 
-/** @public */
-export interface CatalogEntityDocument extends IndexableDocument {
-  componentType: string;
-  namespace: string;
-  kind: string;
-  lifecycle: string;
-  owner: string;
-}
-
-/** @public */
-export class DefaultCatalogCollator implements DocumentCollator {
+/**
+ * @public
+ * @deprecated Upgrade to a more recent `@backstage/search-backend-node` and
+ * use `DefaultCatalogCollatorFactory` instead.
+ */
+export class DefaultCatalogCollator {
   protected discovery: PluginEndpointDiscovery;
   protected locationTemplate: string;
   protected filter?: GetEntitiesRequest['filter'];

plugins/catalog-backend/src/search/DefaultCatalogCollatorFactory.test.ts

@@ -0,0 +1,213 @@
+/*
+ * 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 {
+  PluginEndpointDiscovery,
+  TokenManager,
+} from '@backstage/backend-common';
+import { Entity } from '@backstage/catalog-model';
+import { ConfigReader } from '@backstage/config';
+import { TestPipeline } from '@backstage/plugin-search-backend-node';
+import { rest } from 'msw';
+import { setupServer } from 'msw/node';
+import { Readable } from 'stream';
+import { DefaultCatalogCollatorFactory } from './DefaultCatalogCollatorFactory';
+
+const server = setupServer();
+
+const expectedEntities: Entity[] = [
+  {
+    apiVersion: 'backstage.io/v1alpha1',
+    kind: 'Component',
+    metadata: {
+      name: 'test-entity',
+      description: 'The expected description',
+    },
+    spec: {
+      type: 'some-type',
+      lifecycle: 'experimental',
+      owner: 'someone',
+    },
+  },
+  {
+    apiVersion: 'backstage.io/v1alpha1',
+    kind: 'Component',
+    metadata: {
+      title: 'Test Entity',
+      name: 'test-entity-2',
+      description: 'The expected description 2',
+    },
+    spec: {
+      type: 'some-type',
+      lifecycle: 'experimental',
+      owner: 'someone',
+    },
+  },
+];
+
+describe('DefaultCatalogCollatorFactory', () => {
+  const config = new ConfigReader({});
+  const mockDiscoveryApi: jest.Mocked<PluginEndpointDiscovery> = {
+    getBaseUrl: jest.fn().mockResolvedValue('http://localhost:7007'),
+    getExternalBaseUrl: jest.fn(),
+  };
+  const mockTokenManager: jest.Mocked<TokenManager> = {
+    getToken: jest.fn().mockResolvedValue({ token: '' }),
+    authenticate: jest.fn(),
+  };
+  const options = {
+    discovery: mockDiscoveryApi,
+    tokenManager: mockTokenManager,
+  };
+
+  beforeAll(() => {
+    server.listen();
+  });
+
+  beforeEach(() => {
+    server.use(
+      rest.get('http://localhost:7007/entities', (req, res, ctx) => {
+        if (req.url.searchParams.has('filter')) {
+          const filter = req.url.searchParams.get('filter');
+          if (filter === 'kind=Foo,kind=Bar') {
+            // When filtering on the 'Foo,Bar' kinds we simply return no items, to simulate a filter
+            return res(ctx.json([]));
+          }
+          throw new Error('Unexpected filter parameter');
+        }
+
+        // Imitate offset/limit pagination.
+        const offset = parseInt(req.url.searchParams.get('offset') || '0', 10);
+        const limit = parseInt(req.url.searchParams.get('limit') || '500', 10);
+        return res(ctx.json(expectedEntities.slice(offset, limit + offset)));
+      }),
+    );
+  });
+
+  afterAll(() => {
+    server.close();
+  });
+
+  afterEach(() => server.resetHandlers());
+
+  it('has expected type', () => {
+    const factory = DefaultCatalogCollatorFactory.fromConfig(config, options);
+    expect(factory.type).toBe('software-catalog');
+  });
+
+  describe('getCollator', () => {
+    let factory: DefaultCatalogCollatorFactory;
+    let collator: Readable;
+
+    beforeEach(async () => {
+      factory = DefaultCatalogCollatorFactory.fromConfig(config, options);
+      collator = await factory.getCollator();
+    });
+
+    it('returns a readable stream', async () => {
+      expect(collator).toBeInstanceOf(Readable);
+    });
+
+    it('fetches from the configured catalog service', async () => {
+      const pipeline = TestPipeline.withSubject(collator);
+      const { documents } = await pipeline.execute();
+      expect(mockDiscoveryApi.getBaseUrl).toHaveBeenCalledWith('catalog');
+      expect(documents).toHaveLength(expectedEntities.length);
+    });
+
+    it('maps a returned entity to an expected CatalogEntityDocument', async () => {
+      const pipeline = TestPipeline.withSubject(collator);
+      const { documents } = await pipeline.execute();
+
+      expect(documents[0]).toMatchObject({
+        title: expectedEntities[0].metadata.name,
+        location: '/catalog/default/component/test-entity',
+        text: expectedEntities[0].metadata.description,
+        namespace: 'default',
+        componentType: expectedEntities[0]!.spec!.type,
+        lifecycle: expectedEntities[0]!.spec!.lifecycle,
+        owner: expectedEntities[0]!.spec!.owner,
+        authorization: {
+          resourceRef: 'component:default/test-entity',
+        },
+      });
+      expect(documents[1]).toMatchObject({
+        title: expectedEntities[1].metadata.title,
+        location: '/catalog/default/component/test-entity-2',
+        text: expectedEntities[1].metadata.description,
+        namespace: 'default',
+        componentType: expectedEntities[1]!.spec!.type,
+        lifecycle: expectedEntities[1]!.spec!.lifecycle,
+        owner: expectedEntities[1]!.spec!.owner,
+        authorization: {
+          resourceRef: 'component:default/test-entity-2',
+        },
+      });
+    });
+
+    it('maps a returned entity with a custom locationTemplate', async () => {
+      // Provide an alternate location template.
+      factory = DefaultCatalogCollatorFactory.fromConfig(new ConfigReader({}), {
+        discovery: mockDiscoveryApi,
+        tokenManager: mockTokenManager,
+        locationTemplate: '/software/:name',
+      });
+      collator = await factory.getCollator();
+
+      const pipeline = TestPipeline.withSubject(collator);
+      const { documents } = await pipeline.execute();
+      expect(documents[0]).toMatchObject({
+        location: '/software/test-entity',
+      });
+    });
+
+    it('allows filtering of the retrieved catalog entities', async () => {
+      // Provide a custom filter.
+      factory = DefaultCatalogCollatorFactory.fromConfig(new ConfigReader({}), {
+        discovery: mockDiscoveryApi,
+        tokenManager: mockTokenManager,
+        filter: {
+          kind: ['Foo', 'Bar'],
+        },
+      });
+      collator = await factory.getCollator();
+
+      const pipeline = TestPipeline.withSubject(collator);
+      const { documents } = await pipeline.execute();
+
+      // The simulated 'Foo,Bar' filter should return in an empty list
+      expect(documents).toHaveLength(0);
+    });
+
+    it('paginates through catalog entities using batchSize', async () => {
+      factory = DefaultCatalogCollatorFactory.fromConfig(config, {
+        ...options,
+        batchSize: 1,
+      });
+      collator = await factory.getCollator();
+
+      const pipeline = TestPipeline.withSubject(collator);
+      const { documents } = await pipeline.execute();
+
+      expect(documents).toHaveLength(expectedEntities.length);
+      expect(documents[0].location).toBe(
+        '/catalog/default/component/test-entity',
+      );
+      expect(documents[1].location).toBe(
+        '/catalog/default/component/test-entity-2',
+      );
+    });
+  });
+});

plugins/catalog-backend/src/search/DefaultCatalogCollatorFactory.ts

@@ -0,0 +1,173 @@
+/*
+ * 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 {
+  PluginEndpointDiscovery,
+  TokenManager,
+} from '@backstage/backend-common';
+import {
+  CatalogApi,
+  CatalogClient,
+  GetEntitiesRequest,
+} from '@backstage/catalog-client';
+import {
+  Entity,
+  stringifyEntityRef,
+  UserEntity,
+} from '@backstage/catalog-model';
+import { Config } from '@backstage/config';
+import {
+  DocumentCollatorFactory,
+  IndexableDocument,
+} from '@backstage/search-common';
+import { catalogEntityReadPermission } from '@backstage/plugin-catalog-common';
+import { Readable } from 'stream';
+
+/** @public */
+export interface CatalogEntityDocument extends IndexableDocument {
+  componentType: string;
+  namespace: string;
+  kind: string;
+  lifecycle: string;
+  owner: string;
+}
+
+/** @public */
+export type DefaultCatalogCollatorFactoryOptions = {
+  discovery: PluginEndpointDiscovery;
+  tokenManager: TokenManager;
+  locationTemplate?: string;
+  filter?: GetEntitiesRequest['filter'];
+  batchSize?: number;
+  catalogClient?: CatalogApi;
+};
+
+/** @public */
+export class DefaultCatalogCollatorFactory implements DocumentCollatorFactory {
+  public readonly type: string = 'software-catalog';
+  public readonly visibilityPermission = catalogEntityReadPermission;
+
+  private locationTemplate: string;
+  private filter?: GetEntitiesRequest['filter'];
+  private batchSize: number;
+  private readonly catalogClient: CatalogApi;
+  private tokenManager: TokenManager;
+
+  static fromConfig(
+    _config: Config,
+    options: DefaultCatalogCollatorFactoryOptions,
+  ) {
+    return new DefaultCatalogCollatorFactory(options);
+  }
+
+  private constructor(options: DefaultCatalogCollatorFactoryOptions) {
+    const {
+      batchSize,
+      discovery,
+      locationTemplate,
+      filter,
+      catalogClient,
+      tokenManager,
+    } = options;
+
+    this.locationTemplate =
+      locationTemplate || '/catalog/:namespace/:kind/:name';
+    this.filter = filter;
+    this.batchSize = batchSize || 500;
+    this.catalogClient =
+      catalogClient || new CatalogClient({ discoveryApi: discovery });
+    this.tokenManager = tokenManager;
+  }
+
+  async getCollator(): Promise<Readable> {
+    return Readable.from(this.execute());
+  }
+
+  private applyArgsToFormat(
+    format: string,
+    args: Record<string, string>,
+  ): string {
+    let formatted = format;
+    for (const [key, value] of Object.entries(args)) {
+      formatted = formatted.replace(`:${key}`, value);
+    }
+    return formatted.toLowerCase();
+  }
+
+  private isUserEntity(entity: Entity): entity is UserEntity {
+    return entity.kind.toLocaleUpperCase('en-US') === 'USER';
+  }
+
+  private getDocumentText(entity: Entity): string {
+    let documentText = entity.metadata.description || '';
+    if (this.isUserEntity(entity)) {
+      if (entity.spec?.profile?.displayName && documentText) {
+        // combine displayName and description
+        const displayName = entity.spec?.profile?.displayName;
+        documentText = displayName.concat(' : ', documentText);
+      } else {
+        documentText = entity.spec?.profile?.displayName || documentText;
+      }
+    }
+    return documentText;
+  }
+
+  private async *execute(): AsyncGenerator<CatalogEntityDocument> {
+    const { token } = await this.tokenManager.getToken();
+    let entitiesRetrieved = 0;
+    let moreEntitiesToGet = true;
+
+    // Offset/limit pagination is used on the Catalog Client in order to
+    // limit (and allow some control over) memory used by the search backend
+    // at index-time.
+    while (moreEntitiesToGet) {
+      const entities = (
+        await this.catalogClient.getEntities(
+          {
+            filter: this.filter,
+            limit: this.batchSize,
+            offset: entitiesRetrieved,
+          },
+          { token },
+        )
+      ).items;
+
+      // Control looping through entity batches.
+      moreEntitiesToGet = entities.length === this.batchSize;
+      entitiesRetrieved += entities.length;
+
+      for (const entity of entities) {
+        yield {
+          title: entity.metadata.title ?? entity.metadata.name,
+          location: this.applyArgsToFormat(this.locationTemplate, {
+            namespace: entity.metadata.namespace || 'default',
+            kind: entity.kind,
+            name: entity.metadata.name,
+          }),
+          text: this.getDocumentText(entity),
+          componentType: entity.spec?.type?.toString() || 'other',
+          namespace: entity.metadata.namespace || 'default',
+          kind: entity.kind,
+          lifecycle: (entity.spec?.lifecycle as string) || '',
+          owner: (entity.spec?.owner as string) || '',
+          authorization: {
+            resourceRef: stringifyEntityRef(entity),
+          },
+        };
+      }
+    }
+  }
+}

plugins/catalog-backend/src/search/index.ts

@@ -14,5 +14,11 @@
  * limitations under the License.
  */
 
+export { DefaultCatalogCollatorFactory } from './DefaultCatalogCollatorFactory';
+export type { DefaultCatalogCollatorFactoryOptions } from './DefaultCatalogCollatorFactory';
+export type { CatalogEntityDocument } from './DefaultCatalogCollatorFactory';
+
+/**
+ * todo(backstage/techdocs-core): stop exporting this in a future release.
+ */
 export { DefaultCatalogCollator } from './DefaultCatalogCollator';
-export type { CatalogEntityDocument } from './DefaultCatalogCollator';