From 047c018c91f2a113896aaf8737abe7da67478842 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fredrik=20Adel=C3=B6w?= <freben@gmail.com>
Date: Sun, 6 Dec 2020 15:19:33 +0100
Subject: [PATCH] catalog-backend: read all relations at once

---
 .changeset/fluffy-pillows-smile.md            |  5 ++
 .../src/database/CommonDatabase.ts            | 75 +++++++++++++------
 .../src/ingestion/HigherOrderOperations.ts    |  2 +-
 3 files changed, 58 insertions(+), 24 deletions(-)
 create mode 100644 .changeset/fluffy-pillows-smile.md

diff --git a/.changeset/fluffy-pillows-smile.md b/.changeset/fluffy-pillows-smile.md
new file mode 100644
index 0000000000..5d15974035
--- /dev/null
+++ b/.changeset/fluffy-pillows-smile.md
@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-catalog-backend': patch
+---
+
+Batch the fetching of relations
diff --git a/plugins/catalog-backend/src/database/CommonDatabase.ts b/plugins/catalog-backend/src/database/CommonDatabase.ts
index 287756b0a8..e38d3032b4 100644
--- a/plugins/catalog-backend/src/database/CommonDatabase.ts
+++ b/plugins/catalog-backend/src/database/CommonDatabase.ts
@@ -271,7 +271,7 @@ export class CommonDatabase implements Database {
       .select('entities.*')
       .orderBy('full_name', 'asc');
 
-    return Promise.all(rows.map(row => this.toEntityResponse(tx, row)));
+    return this.toEntityResponses(tx, rows);
   }
 
   async entityByName(
@@ -290,7 +290,7 @@ export class CommonDatabase implements Database {
       return undefined;
     }
 
-    return this.toEntityResponse(tx, rows[0]);
+    return this.toEntityResponses(tx, rows).then(r => r[0]);
   }
 
   async entityByUid(
@@ -307,7 +307,7 @@ export class CommonDatabase implements Database {
       return undefined;
     }
 
-    return this.toEntityResponse(tx, rows[0]);
+    return this.toEntityResponses(tx, rows).then(r => r[0]);
   }
 
   async removeEntityByUid(txOpaque: Transaction, uid: string): Promise<void> {
@@ -494,31 +494,60 @@ export class CommonDatabase implements Database {
     };
   }
 
-  private async toEntityResponse(
+  private async toEntityResponses(
     tx: Knex.Transaction<any, any>,
-    row: DbEntitiesRow,
-  ): Promise<DbEntityResponse> {
-    const entity = JSON.parse(row.data) as Entity;
-    entity.metadata.uid = row.id;
-    entity.metadata.etag = row.etag;
-    entity.metadata.generation = Number(row.generation); // cast due to sqlite
-
+    rows: DbEntitiesRow[],
+  ): Promise<DbEntityResponse[]> {
     // TODO(Rugvip): This is here because it's simple for now, but we likely
     //               need to refactor this to be more efficient or introduce pagination.
-    const relations = await tx<DbEntitiesRelationsRow>('entities_relations')
-      .where({ source_full_name: row.full_name })
-      .orderBy(['type', 'target_full_name'])
-      .select();
+    const relations = await this.getRelationsPerFullName(
+      tx,
+      rows.map(r => r.full_name),
+    );
 
-    entity.relations = deduplicateRelations(relations).map(r => ({
-      target: parseEntityName(r.target_full_name),
-      type: r.type,
-    }));
+    const result = new Array<DbEntityResponse>();
+    for (const row of rows) {
+      const entity = JSON.parse(row.data) as Entity;
+      entity.metadata.uid = row.id;
+      entity.metadata.etag = row.etag;
+      entity.metadata.generation = Number(row.generation); // cast due to sqlite
 
-    return {
-      locationId: row.location_id || undefined,
-      entity,
-    };
+      entity.relations = (relations[row.full_name] ?? []).map(r => ({
+        target: parseEntityName(r.target_full_name),
+        type: r.type,
+      }));
+
+      result.push({
+        locationId: row.location_id || undefined,
+        entity,
+      });
+    }
+
+    return result;
+  }
+
+  // Returns a mapping from e.g. component:default/foo to the relations whose
+  // source_full_name matches that.
+  private async getRelationsPerFullName(
+    tx: Knex.Transaction<any, any>,
+    sourceFullNames: string[],
+  ): Promise<Record<string, DbEntitiesRelationsRow[]>> {
+    const batches = lodash.chunk(lodash.uniq(sourceFullNames), 500);
+
+    const relations = new Array<DbEntitiesRelationsRow>();
+    for (const batch of batches) {
+      relations.push(
+        ...(await tx<DbEntitiesRelationsRow>('entities_relations')
+          .whereIn('source_full_name', batch)
+          .orderBy(['type', 'target_full_name'])
+          .select()),
+      );
+    }
+
+    return lodash.groupBy(
+      deduplicateRelations(relations),
+      r => r.source_full_name,
+    );
   }
 }
 
diff --git a/plugins/catalog-backend/src/ingestion/HigherOrderOperations.ts b/plugins/catalog-backend/src/ingestion/HigherOrderOperations.ts
index bc277e1767..ca19a1d1c5 100644
--- a/plugins/catalog-backend/src/ingestion/HigherOrderOperations.ts
+++ b/plugins/catalog-backend/src/ingestion/HigherOrderOperations.ts
@@ -85,7 +85,7 @@ export class HigherOrderOperations implements HigherOrderOperation {
     // Write
     if (!previousLocation && !dryRun) {
       // TODO: We do not include location operations in the dryRun. We might perform
-      // this operation as a seperate dry run.
+      // this operation as a separate dry run.
       await this.locationsCatalog.addLocation(location);
     }
     if (readerOutput.entities.length === 0) {