feat: ensure that refresh_state_references are cleaned

Signed-off-by: blam <ben@blam.sh>

plugins/catalog-backend/src/database/DefaultProviderDatabase.ts

@@ -165,7 +165,6 @@ export class DefaultProviderDatabase implements ProviderDatabase {
 
           await tx<DbRefreshStateReferencesRow>('refresh_state_references')
             .where('target_entity_ref', entityRef)
-            .andWhere({ source_key: options.sourceKey })
             .delete();
 
           if (ok) {

plugins/catalog-backend/src/tests/integration.test.ts

@@ -31,7 +31,7 @@ import {
 } from '@backstage/plugin-catalog-node';
 import { PermissionEvaluator } from '@backstage/plugin-permission-common';
 import { createHash } from 'crypto';
-import { Knex } from 'knex';
+import knexFactory, { Knex } from 'knex';
 import { EntitiesCatalog } from '../catalog/types';
 import { DefaultCatalogDatabase } from '../database/DefaultCatalogDatabase';
 import { DefaultProcessingDatabase } from '../database/DefaultProcessingDatabase';
@@ -55,6 +55,7 @@ import { mockServices } from '@backstage/backend-test-utils';
 import { LoggerService } from '@backstage/backend-plugin-api';
 import { DatabaseManager } from '@backstage/backend-common';
 import { entitiesResponseToObjects } from '../service/response';
+import { DbRefreshStateReferencesRow } from '../database/tables';
 
 const voidLogger = mockServices.logger.mock();
 
@@ -194,7 +195,7 @@ class TestHarness {
   readonly #catalog: EntitiesCatalog;
   readonly #engine: CatalogProcessingEngine;
   readonly #refresh: RefreshService;
-  readonly #provider: TestProvider;
+  readonly #providers: TestProvider[];
   readonly #proxyProgressTracker: ProxyProgressTracker;
 
   static async create(options?: {
@@ -202,6 +203,7 @@ class TestHarness {
     logger?: LoggerService;
     db?: Knex;
     permissions?: PermissionEvaluator;
+    providers?: TestProvider[];
     processEntity?(
       entity: Entity,
       location: LocationSpec,
@@ -300,9 +302,8 @@ class TestHarness {
 
     const refresh = new DefaultRefreshService({ database: catalogDatabase });
 
-    const provider = new TestProvider();
-
-    await connectEntityProviders(providerDatabase, [provider]);
+    const providers = options?.providers ?? [new TestProvider()];
+    await connectEntityProviders(providerDatabase, providers);
 
     return new TestHarness(
       catalog,
@@ -317,7 +318,7 @@ class TestHarness {
         },
       },
       refresh,
-      provider,
+      providers,
       proxyProgressTracker,
     );
   }
@@ -326,13 +327,13 @@ class TestHarness {
     catalog: EntitiesCatalog,
     engine: CatalogProcessingEngine,
     refresh: RefreshService,
-    provider: TestProvider,
+    providers: TestProvider[],
     proxyProgressTracker: ProxyProgressTracker,
   ) {
     this.#catalog = catalog;
     this.#engine = engine;
     this.#refresh = refresh;
-    this.#provider = provider;
+    this.#providers = providers;
     this.#proxyProgressTracker = proxyProgressTracker;
   }
 
@@ -353,13 +354,15 @@ class TestHarness {
   }
 
   async setInputEntities(entities: (Entity & { locationKey?: string })[]) {
-    return this.#provider.getConnection().applyMutation({
-      type: 'full',
-      entities: entities.map(({ locationKey, ...entity }) => ({
-        entity,
-        locationKey,
-      })),
-    });
+    for (const provider of this.#providers) {
+      await provider.getConnection().applyMutation({
+        type: 'full',
+        entities: entities.map(({ locationKey, ...entity }) => ({
+          entity,
+          locationKey,
+        })),
+      });
+    }
   }
 
   async getOutputEntities(): Promise<Record<string, Entity>> {
@@ -838,4 +841,135 @@ describe('Catalog Backend Integration', () => {
       },
     });
   });
+
+  it('should replace any refresh_state_references that are dangling after claiming an entityRef with locationKey', async () => {
+    const db = knexFactory({
+      client: 'better-sqlite3',
+      connection: ':memory:',
+      useNullAsDefault: true,
+    });
+
+    const firstProvider = new TestProvider();
+    firstProvider.getProviderName = () => 'first';
+
+    const secondProvider = new TestProvider();
+    secondProvider.getProviderName = () => 'second';
+
+    const harness = await TestHarness.create({
+      providers: [firstProvider, secondProvider],
+      db,
+    });
+
+    await firstProvider.getConnection().applyMutation({
+      type: 'full',
+      entities: [
+        {
+          entity: {
+            apiVersion: 'backstage.io/v1alpha1',
+            kind: 'Component',
+            metadata: {
+              name: 'component-1',
+              annotations: {
+                'backstage.io/managed-by-location': 'url:.',
+                'backstage.io/managed-by-origin-location': 'url:.',
+              },
+            },
+            spec: {
+              type: 'service',
+              owner: 'no-location-key',
+            },
+          },
+        },
+      ],
+    });
+
+    await expect(harness.process()).resolves.toEqual({});
+
+    await expect(harness.getOutputEntities()).resolves.toEqual({
+      'component:default/component-1': expect.objectContaining({
+        spec: {
+          type: 'service',
+          owner: 'no-location-key',
+        },
+      }),
+    });
+
+    await secondProvider.getConnection().applyMutation({
+      type: 'full',
+      entities: [
+        {
+          locationKey: 'takeover',
+          entity: {
+            apiVersion: 'backstage.io/v1alpha1',
+            kind: 'Component',
+            metadata: {
+              name: 'component-1',
+              annotations: {
+                'backstage.io/managed-by-location': 'url:.',
+                'backstage.io/managed-by-origin-location': 'url:.',
+              },
+            },
+            spec: {
+              type: 'service',
+              owner: 'location-key',
+            },
+          },
+        },
+      ],
+    });
+
+    await expect(harness.process()).resolves.toEqual({});
+
+    await expect(harness.getOutputEntities()).resolves.toEqual({
+      'component:default/component-1': expect.objectContaining({
+        spec: {
+          type: 'service',
+          owner: 'location-key',
+        },
+      }),
+    });
+
+    await expect(
+      db<DbRefreshStateReferencesRow>('refresh_state_references')
+        .where({ target_entity_ref: 'component:default/component-1' })
+        .select(),
+    ).resolves.toEqual([
+      expect.objectContaining({
+        source_key: 'second',
+        target_entity_ref: 'component:default/component-1',
+      }),
+    ]);
+
+    await secondProvider.getConnection().applyMutation({
+      type: 'full',
+      entities: [
+        {
+          locationKey: 'takeover',
+          entity: {
+            apiVersion: 'backstage.io/v1alpha1',
+            kind: 'Component',
+            metadata: {
+              name: 'component-2',
+              annotations: {
+                'backstage.io/managed-by-location': 'url:.',
+                'backstage.io/managed-by-origin-location': 'url:.',
+              },
+            },
+            spec: {
+              type: 'service',
+              owner: 'location-key',
+            },
+          },
+        },
+      ],
+    });
+
+    await expect(harness.process()).resolves.toEqual({});
+
+    await expect(
+      db<DbRefreshStateReferencesRow>('refresh_state_references')
+        .where({ target_entity_ref: 'component:default/component-1' })
+        .select(),
+    ).resolves.toEqual([]);
+  });
 });