Merge pull request #34301 from backstage/freben/incremental-ingestion-any-array

incremental-ingestion: use ANY(array) instead of IN(...) on Postgres

.changeset/incremental-ingestion-any-array.md

@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-catalog-backend-module-incremental-ingestion': patch
+---
+
+On PostgreSQL, `WHERE ref IN ($1, $2, ..., $N)` queries on the `ingestion_mark_entities` table now use `= ANY($1)` with a single array parameter instead. This reduces prepared statement bloat in the query plan cache when the number of entity refs varies between calls.

plugins/catalog-backend-module-incremental-ingestion/src/database/IncrementalIngestionDatabaseManager.test.ts

@@ -121,5 +121,124 @@ describe.each(databases.eachSupportedId())(
       expect(result.total).toBe(3);
       expect(typeof result.total).toBe('number');
     });
+
+    it('createMarkEntities handles existing and new refs correctly', async () => {
+      const knex = await databases.init(databaseId);
+      await knex.migrate.latest({ directory: migrationsDir });
+
+      const manager = new IncrementalIngestionDatabaseManager({ client: knex });
+      const { ingestionId } = (await manager.createProviderIngestionRecord(
+        'testProvider',
+      ))!;
+
+      const markId1 = uuid();
+      await manager.createMark({
+        record: {
+          id: markId1,
+          ingestion_id: ingestionId,
+          sequence: 1,
+          cursor: { data: 1 },
+        },
+      });
+
+      const makeEntity = (name: string): DeferredEntity => ({
+        entity: {
+          apiVersion: 'backstage.io/v1alpha1',
+          kind: 'Component',
+          metadata: { namespace: 'default', name },
+        },
+      });
+
+      // First batch: create 3 entities
+      await manager.createMarkEntities(markId1, [
+        makeEntity('a'),
+        makeEntity('b'),
+        makeEntity('c'),
+      ]);
+
+      const rows1 = await knex('ingestion_mark_entities').select('ref');
+      expect(rows1).toHaveLength(3);
+
+      // Second batch with overlap: b and c already exist, d is new.
+      // Existing refs should be updated to the new mark, new refs inserted.
+      const markId2 = uuid();
+      await manager.createMark({
+        record: {
+          id: markId2,
+          ingestion_id: ingestionId,
+          sequence: 2,
+          cursor: { data: 2 },
+        },
+      });
+
+      await manager.createMarkEntities(markId2, [
+        makeEntity('b'),
+        makeEntity('c'),
+        makeEntity('d'),
+      ]);
+
+      const rows2 = await knex('ingestion_mark_entities')
+        .select('ref', 'ingestion_mark_id')
+        .orderBy('ref');
+      expect(rows2).toHaveLength(4);
+
+      // a stays on markId1, b and c moved to markId2, d is new on markId2
+      expect(
+        rows2.find(r => r.ref === 'component:default/a')?.ingestion_mark_id,
+      ).toBe(markId1);
+      expect(
+        rows2.find(r => r.ref === 'component:default/b')?.ingestion_mark_id,
+      ).toBe(markId2);
+      expect(
+        rows2.find(r => r.ref === 'component:default/c')?.ingestion_mark_id,
+      ).toBe(markId2);
+      expect(
+        rows2.find(r => r.ref === 'component:default/d')?.ingestion_mark_id,
+      ).toBe(markId2);
+    });
+
+    it('deleteEntityRecordsByRef removes matching refs', async () => {
+      const knex = await databases.init(databaseId);
+      await knex.migrate.latest({ directory: migrationsDir });
+
+      const manager = new IncrementalIngestionDatabaseManager({ client: knex });
+      const { ingestionId } = (await manager.createProviderIngestionRecord(
+        'testProvider',
+      ))!;
+
+      const markId = uuid();
+      await manager.createMark({
+        record: {
+          id: markId,
+          ingestion_id: ingestionId,
+          sequence: 1,
+          cursor: { data: 1 },
+        },
+      });
+
+      const makeEntity = (name: string): DeferredEntity => ({
+        entity: {
+          apiVersion: 'backstage.io/v1alpha1',
+          kind: 'Component',
+          metadata: { namespace: 'default', name },
+        },
+      });
+
+      await manager.createMarkEntities(markId, [
+        makeEntity('x'),
+        makeEntity('y'),
+        makeEntity('z'),
+      ]);
+
+      // Delete two of the three
+      await manager.deleteEntityRecordsByRef([
+        { entityRef: 'component:default/x' },
+        { entityRef: 'component:default/z' },
+      ]);
+
+      const remaining = await knex('ingestion_mark_entities').select('ref');
+      expect(remaining).toHaveLength(1);
+      expect(remaining[0].ref).toBe('component:default/y');
+    });
   },
 );

plugins/catalog-backend-module-incremental-ingestion/src/database/IncrementalIngestionDatabaseManager.ts

@@ -34,6 +34,17 @@ export class IncrementalIngestionDatabaseManager {
     this.client = options.client;
   }
 
+  private whereInArray(column: string, values: string[]) {
+    const isPg = this.client.client.config.client === 'pg';
+    return (qb: Knex.QueryBuilder) => {
+      if (isPg) {
+        qb.whereRaw('?? = ANY(?)', [column, values]);
+      } else {
+        qb.whereIn(column, values);
+      }
+    };
+  }
+
   /**
    * Performs an update to the ingestion record with matching `id`.
    * @param options - IngestionRecordUpdate
@@ -76,24 +87,25 @@ export class IncrementalIngestionDatabaseManager {
     tx: Knex.Transaction,
     ids: { id: string }[],
   ) {
-    const chunks: { id: string }[][] = [];
-    for (let i = 0; i < ids.length; i += 100) {
-      const chunk = ids.slice(i, i + 100);
-      chunks.push(chunk);
+    if (ids.length === 0) {
+      return 0;
+    }
+
+    const allIds = ids.map(entry => entry.id);
+
+    if (this.client.client.config.client === 'pg') {
+      return await tx('ingestion_mark_entities')
+        .delete()
+        .modify(this.whereInArray('id', allIds));
     }
 
     let deleted = 0;
-
-    for (const chunk of chunks) {
-      const chunkDeleted = await tx('ingestion_mark_entities')
+    for (let i = 0; i < allIds.length; i += 100) {
+      const chunk = allIds.slice(i, i + 100);
+      deleted += await tx('ingestion_mark_entities')
         .delete()
-        .whereIn(
-          'id',
-          chunk.map(entry => entry.id),
-        );
-      deleted += chunkDeleted;
+        .whereIn('id', chunk);
     }
-
     return deleted;
   }
 
@@ -276,7 +288,9 @@ export class IncrementalIngestionDatabaseManager {
   async deleteEntityRecordsByRef(entities: { entityRef: string }[]) {
     const refs = entities.map(e => e.entityRef);
     await this.client.transaction(async tx => {
-      await tx('ingestion_mark_entities').delete().whereIn('ref', refs);
+      await tx('ingestion_mark_entities')
+        .delete()
+        .modify(this.whereInArray('ref', refs));
     });
   }
 
@@ -603,16 +617,18 @@ export class IncrementalIngestionDatabaseManager {
       const existingRefsArray = (
         await tx<{ ref: string }>('ingestion_mark_entities')
           .select('ref')
-          .whereIn('ref', refs)
-      ).map(e => e.ref);
+          .modify(this.whereInArray('ref', refs))
+      ).map((e: { ref: string }) => e.ref);
 
       const existingRefsSet = new Set(existingRefsArray);
 
       const newRefs = refs.filter(e => !existingRefsSet.has(e));
 
-      await tx('ingestion_mark_entities')
-        .update('ingestion_mark_id', markId)
-        .whereIn('ref', existingRefsArray);
+      if (existingRefsArray.length > 0) {
+        await tx('ingestion_mark_entities')
+          .update('ingestion_mark_id', markId)
+          .modify(this.whereInArray('ref', existingRefsArray));
+      }
 
       if (newRefs.length > 0) {
         await tx('ingestion_mark_entities').insert(