Refactor StateManager into the ProcessingEngine

And add some tests

Co-authored-by: Ben Lambert <ben@blam.sh>
Signed-off-by: Johan Haals <johan.haals@gmail.com>

packages/backend/src/plugins/catalog.ts

@@ -42,7 +42,7 @@ export default async function createPlugin(
     } = await builder.build();
 
     // TODO(jhaals): run and manage in background.
-    processingEngine.start();
+    await processingEngine.start();
 
     return await createRouter({
       entitiesCatalog,

plugins/catalog-backend/src/next/DefaultCatalogProcessingEngine.test.ts

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2021 Spotify AB
+ *
+ * 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 { getVoidLogger } from '@backstage/backend-common';
+import { DefaultProcessingDatabase } from './database/DefaultProcessingDatabase';
+
+import { DefaultCatalogProcessingEngine } from './DefaultCatalogProcessingEngine';
+import { Stitcher } from './Stitcher';
+import { CatalogProcessingOrchestrator } from './types';
+
+describe('DefaultCatalogProcessingEngine', () => {
+  const db = ({
+    transaction: jest.fn(),
+    getProcessableEntities: jest.fn(),
+    updateProcessedEntity: jest.fn(),
+  } as unknown) as jest.Mocked<DefaultProcessingDatabase>;
+  const orchestrator: jest.Mocked<CatalogProcessingOrchestrator> = {
+    process: jest.fn(),
+  };
+  const stitcher = ({
+    stitch: jest.fn(),
+  } as unknown) as jest.Mocked<Stitcher>;
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+  });
+  it('should process stuff', async () => {
+    orchestrator.process.mockResolvedValue({
+      ok: true,
+      completedEntity: {
+        apiVersion: '1',
+        kind: 'Location',
+        metadata: { name: 'test' },
+      },
+      relations: [],
+      errors: [],
+      deferredEntities: [],
+      state: new Map(),
+    });
+    const engine = new DefaultCatalogProcessingEngine(
+      getVoidLogger(),
+      [],
+      db,
+      orchestrator,
+      stitcher,
+    );
+
+    db.transaction.mockImplementation(cb => cb((() => {}) as any));
+
+    db.getProcessableEntities
+      .mockImplementation(async () => {
+        await engine.stop();
+        return { items: [] };
+      })
+      .mockResolvedValueOnce({
+        items: [
+          {
+            entityRef: 'foo',
+            id: '1',
+            unprocessedEntity: {
+              apiVersion: '1',
+              kind: 'Location',
+              metadata: { name: 'test' },
+            },
+            state: new Map(),
+            nextUpdateAt: '',
+            lastDiscoveryAt: '',
+          },
+        ],
+      });
+
+    await engine.start();
+    await new Promise<void>(resolve => setTimeout(resolve, 1));
+    expect(orchestrator.process).toBeCalledTimes(1);
+    expect(orchestrator.process).toBeCalledWith({
+      entity: { apiVersion: '1', kind: 'Location', metadata: { name: 'test' } },
+      state: expect.anything(),
+    });
+
+    await engine.stop();
+  });
+});

plugins/catalog-backend/src/next/DefaultCatalogProcessingEngine.ts

@@ -16,6 +16,7 @@
 
 import { stringifyEntityRef } from '@backstage/catalog-model';
 import { Logger } from 'winston';
+import { ProcessingDatabase } from './database/types';
 import { Stitcher } from './Stitcher';
 import {
   CatalogProcessingEngine,
@@ -23,44 +24,46 @@ import {
   EntityProvider,
   EntityProviderConnection,
   EntityProviderMutation,
-  ProcessingStateManager,
 } from './types';
 
 class Connection implements EntityProviderConnection {
   constructor(
     private readonly config: {
-      stateManager: ProcessingStateManager;
+      processingDatabase: ProcessingDatabase;
       id: string;
     },
   ) {}
 
   async applyMutation(mutation: EntityProviderMutation): Promise<void> {
+    const db = this.config.processingDatabase;
     if (mutation.type === 'full') {
-      await this.config.stateManager.replaceProcessingItems({
-        sourceKey: this.config.id,
-        type: 'full',
-        items: mutation.entities,
+      db.transaction(async tx => {
+        await db.replaceUnprocessedEntities(tx, {
+          sourceKey: this.config.id,
+          type: 'full',
+          items: mutation.entities,
+        });
       });
-
       return;
     }
-
-    await this.config.stateManager.replaceProcessingItems({
-      sourceKey: this.config.id,
-      type: 'delta',
-      added: mutation.added,
-      removed: mutation.removed,
+    db.transaction(async tx => {
+      await db.replaceUnprocessedEntities(tx, {
+        sourceKey: this.config.id,
+        type: 'delta',
+        added: mutation.added,
+        removed: mutation.removed,
+      });
     });
   }
 }
 
 export class DefaultCatalogProcessingEngine implements CatalogProcessingEngine {
-  private running: boolean = false;
+  private running = false;
 
   constructor(
     private readonly logger: Logger,
     private readonly entityProviders: EntityProvider[],
-    private readonly stateManager: ProcessingStateManager,
+    private readonly processingDatabase: ProcessingDatabase,
     private readonly orchestrator: CatalogProcessingOrchestrator,
     private readonly stitcher: Stitcher,
   ) {}
@@ -69,51 +72,60 @@ export class DefaultCatalogProcessingEngine implements CatalogProcessingEngine {
     for (const provider of this.entityProviders) {
       await provider.connect(
         new Connection({
-          stateManager: this.stateManager,
+          processingDatabase: this.processingDatabase,
           id: provider.getProviderName(),
         }),
       );
     }
-
     this.running = true;
+    this.run();
+  }
 
+  private async run() {
     while (this.running) {
-      const {
-        id,
-        entity,
-        state: initialState,
-      } = await this.stateManager.getNextProcessingItem();
-
-      const result = await this.orchestrator.process({
-        entity,
-        state: initialState,
+      const { items } = await this.processingDatabase.transaction(async tx => {
+        return this.processingDatabase.getProcessableEntities(tx, {
+          processBatchSize: 1,
+        });
       });
 
-      for (const error of result.errors) {
-        this.logger.warn(error.message);
+      if (items.length) {
+        const { id, state, unprocessedEntity } = items[0];
+
+        const result = await this.orchestrator.process({
+          entity: unprocessedEntity,
+          state,
+        });
+        for (const error of result.errors) {
+          this.logger.warn(error.message);
+        }
+        if (!result.ok) {
+          return;
+        }
+
+        result.completedEntity.metadata.uid = id;
+        await this.processingDatabase.transaction(async tx => {
+          await this.processingDatabase.updateProcessedEntity(tx, {
+            id,
+            processedEntity: result.completedEntity,
+            state: result.state,
+            errors: JSON.stringify(result.errors),
+            relations: result.relations,
+            deferredEntities: result.deferredEntities,
+          });
+        });
+
+        const setOfThingsToStitch = new Set<string>([
+          stringifyEntityRef(result.completedEntity),
+          ...result.relations.map(relation =>
+            stringifyEntityRef(relation.source),
+          ),
+        ]);
+        await this.stitcher.stitch(setOfThingsToStitch);
+      } else {
+        // Wait one second before fetching next item.
+        await new Promise<void>(resolve => setTimeout(resolve, 1000));
       }
-
-      if (!result.ok) {
-        return;
-      }
-
-      result.completedEntity.metadata.uid = id;
-      await this.stateManager.setProcessingItemResult({
-        id,
-        entity: result.completedEntity,
-        state: result.state,
-        errors: result.errors,
-        relations: result.relations,
-        deferredEntities: result.deferredEntities,
-      });
-
-      const setOfThingsToStitch = new Set<string>([
-        stringifyEntityRef(result.completedEntity),
-        ...result.relations.map(relation =>
-          stringifyEntityRef(relation.source),
-        ),
-      ]);
-      await this.stitcher.stitch(setOfThingsToStitch);
     }
   }
 

plugins/catalog-backend/src/next/DefaultProcessingStateManager.ts

@@ -1,69 +0,0 @@
-/*
- * Copyright 2021 Spotify AB
- *
- * 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 { ProcessingDatabase } from './database/types';
-import {
-  ProcessingItem,
-  ProcessingItemResult,
-  ProcessingStateManager,
-  ReplaceProcessingItemsRequest,
-} from './types';
-
-export class DefaultProcessingStateManager implements ProcessingStateManager {
-  constructor(private readonly db: ProcessingDatabase) {}
-
-  replaceProcessingItems(
-    request: ReplaceProcessingItemsRequest,
-  ): Promise<void> {
-    return this.db.transaction(async tx => {
-      await this.db.replaceUnprocessedEntities(tx, request);
-    });
-  }
-
-  async setProcessingItemResult(result: ProcessingItemResult) {
-    return this.db.transaction(async tx => {
-      await this.db.updateProcessedEntity(tx, {
-        id: result.id,
-        processedEntity: result.entity,
-        errors: JSON.stringify(result.errors),
-        state: result.state,
-        relations: result.relations,
-        deferredEntities: result.deferredEntities,
-      });
-    });
-  }
-
-  async getNextProcessingItem(): Promise<ProcessingItem> {
-    for (;;) {
-      const { items } = await this.db.transaction(async tx => {
-        return this.db.getProcessableEntities(tx, {
-          processBatchSize: 1,
-        });
-      });
-
-      if (items.length) {
-        const { id, state, unprocessedEntity } = items[0];
-        return {
-          id,
-          entity: unprocessedEntity,
-          state,
-        };
-      }
-
-      await new Promise<void>(resolve => setTimeout(resolve, 1000));
-    }
-  }
-}

plugins/catalog-backend/src/next/NextCatalogBuilder.ts

@@ -69,7 +69,6 @@ import { DefaultProcessingDatabase } from './database/DefaultProcessingDatabase'
 import { DefaultCatalogProcessingEngine } from './DefaultCatalogProcessingEngine';
 import { DefaultCatalogProcessingOrchestrator } from './DefaultCatalogProcessingOrchestrator';
 import { DefaultLocationStore } from './DefaultLocationStore';
-import { DefaultProcessingStateManager } from './DefaultProcessingStateManager';
 import { NextEntitiesCatalog } from './NextEntitiesCatalog';
 import { Stitcher } from './Stitcher';
 
@@ -252,7 +251,7 @@ export class NextCatalogBuilder {
     const db = new CommonDatabase(dbClient, logger);
 
     const processingDatabase = new DefaultProcessingDatabase(dbClient, logger);
-    const stateManager = new DefaultProcessingStateManager(processingDatabase);
+    // const stateManager = new DefaultProcessingStateManager(processingDatabase);
     const integrations = ScmIntegrations.fromConfig(config);
     const orchestrator = new DefaultCatalogProcessingOrchestrator({
       processors,
@@ -269,7 +268,7 @@ export class NextCatalogBuilder {
     const processingEngine = new DefaultCatalogProcessingEngine(
       logger,
       [locationStore, configLocationProvider],
-      stateManager,
+      processingDatabase,
       orchestrator,
       stitcher,
     );

plugins/catalog-backend/src/next/types.ts

@@ -80,37 +80,3 @@ export type EntityProcessingResult =
 export interface CatalogProcessingOrchestrator {
   process(request: EntityProcessingRequest): Promise<EntityProcessingResult>;
 }
-
-export type ProcessingItemResult = {
-  id: string;
-  entity: Entity;
-  state: Map<string, JsonObject>;
-  errors: Error[];
-  relations: EntityRelationSpec[];
-  deferredEntities: Entity[];
-};
-
-export type ProcessingItem = {
-  id: string;
-  entity: Entity;
-  state: Map<string, JsonObject>;
-};
-
-export type ReplaceProcessingItemsRequest =
-  | {
-      sourceKey: string;
-      items: Entity[];
-      type: 'full';
-    }
-  | {
-      sourceKey: string;
-      added: Entity[];
-      removed: Entity[];
-      type: 'delta';
-    };
-
-export interface ProcessingStateManager {
-  setProcessingItemResult(result: ProcessingItemResult): Promise<void>;
-  getNextProcessingItem(): Promise<ProcessingItem>;
-  replaceProcessingItems(request: ReplaceProcessingItemsRequest): Promise<void>;
-}