Update the entity unregister dialog behavior, to support both unregistration as well as plain deletion

Signed-off-by: Fredrik Adelöw <freben@gmail.com>

.changeset/flat-toes-enjoy.md

@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-catalog-backend': minor
+---
+
+DELETE on an entity now just deletes the entity, rather than removing all related entities and the location

.changeset/silver-starfishes-refuse.md

@@ -0,0 +1,12 @@
+---
+'@backstage/catalog-client': patch
+'@backstage/plugin-auth-backend': patch
+'@backstage/plugin-catalog-import': patch
+'@backstage/plugin-catalog': patch
+'@backstage/plugin-explore': patch
+'@backstage/plugin-register-component': patch
+'@backstage/plugin-scaffolder': patch
+'@backstage/plugin-todo-backend': patch
+---
+
+Added the `getOriginLocationByEntity` and `removeLocationById` methods to the catalog client

.changeset/young-humans-destroy.md

@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-catalog': patch
+---
+
+Improve the unregister dialog, to support both unregistration and plain deletion

.github/styles/vocab.txt

@@ -283,6 +283,7 @@ transpiled
 ui
 unmanaged
 unregister
+unregistration
 untracked
 upvote
 url

packages/catalog-client/src/CatalogClient.ts

@@ -19,6 +19,7 @@ import {
   EntityName,
   Location,
   LOCATION_ANNOTATION,
+  ORIGIN_LOCATION_ANNOTATION,
   stringifyLocationReference,
 } from '@backstage/catalog-model';
 import { ResponseError } from '@backstage/errors';
@@ -44,7 +45,7 @@ export class CatalogClient implements CatalogApi {
     id: String,
     options?: CatalogRequestOptions,
   ): Promise<Location | undefined> {
-    return await this.getOptional(`/locations/${id}`, options);
+    return await this.requestOptional('GET', `/locations/${id}`, options);
   }
 
   async getEntities(
@@ -60,6 +61,7 @@ export class CatalogClient implements CatalogApi {
         filterParts.push(`${encodeURIComponent(key)}=${encodeURIComponent(v)}`);
       }
     }
+
     if (filterParts.length) {
       params.push(`filter=${filterParts.join(',')}`);
     }
@@ -69,7 +71,8 @@ export class CatalogClient implements CatalogApi {
     }
 
     const query = params.length ? `?${params.join('&')}` : '';
-    const entities: Entity[] = await this.getRequired(
+    const entities: Entity[] = await this.requestRequired(
+      'GET',
       `/entities${query}`,
       options,
     );
@@ -81,7 +84,8 @@ export class CatalogClient implements CatalogApi {
     options?: CatalogRequestOptions,
   ): Promise<Entity | undefined> {
     const { kind, namespace = 'default', name } = compoundName;
-    return this.getOptional(
+    return this.requestOptional(
+      'GET',
       `/entities/by-name/${kind}/${namespace}/${name}`,
       options,
     );
@@ -126,12 +130,17 @@ export class CatalogClient implements CatalogApi {
     };
   }
 
-  async getLocationByEntity(
+  async getOriginLocationByEntity(
     entity: Entity,
     options?: CatalogRequestOptions,
   ): Promise<Location | undefined> {
-    const locationCompound = entity.metadata.annotations?.[LOCATION_ANNOTATION];
-    const all: { data: Location }[] = await this.getRequired(
+    const locationCompound =
+      entity.metadata.annotations?.[ORIGIN_LOCATION_ANNOTATION];
+    if (!locationCompound) {
+      return undefined;
+    }
+    const all: { data: Location }[] = await this.requestRequired(
+      'GET',
       '/locations',
       options,
     );
@@ -140,39 +149,68 @@ export class CatalogClient implements CatalogApi {
       .find(l => locationCompound === stringifyLocationReference(l));
   }
 
+  async getLocationByEntity(
+    entity: Entity,
+    options?: CatalogRequestOptions,
+  ): Promise<Location | undefined> {
+    const locationCompound = entity.metadata.annotations?.[LOCATION_ANNOTATION];
+    if (!locationCompound) {
+      return undefined;
+    }
+    const all: { data: Location }[] = await this.requestRequired(
+      'GET',
+      '/locations',
+      options,
+    );
+    return all
+      .map(r => r.data)
+      .find(l => locationCompound === stringifyLocationReference(l));
+  }
+
+  async removeLocationById(
+    id: string,
+    options?: CatalogRequestOptions,
+  ): Promise<void> {
+    await this.requestIgnored('DELETE', `/locations/${id}`, options);
+  }
+
   async removeEntityByUid(
     uid: string,
     options?: CatalogRequestOptions,
   ): Promise<void> {
-    const response = await fetch(
-      `${await this.discoveryApi.getBaseUrl('catalog')}/entities/by-uid/${uid}`,
-      {
-        headers: options?.token
-          ? { Authorization: `Bearer ${options.token}` }
-          : {},
-        method: 'DELETE',
-      },
-    );
-    if (!response.ok) {
-      throw await ResponseError.fromResponse(response);
-    }
-    return undefined;
+    await this.requestIgnored('DELETE', `/entities/by-uid/${uid}`, options);
   }
 
   //
   // Private methods
   //
 
-  private async getRequired(
+  private async requestIgnored(
+    method: string,
+    path: string,
+    options?: CatalogRequestOptions,
+  ): Promise<void> {
+    const url = `${await this.discoveryApi.getBaseUrl('catalog')}${path}`;
+    const headers = new Headers(
+      options?.token ? { Authorization: `Bearer ${options.token}` } : {},
+    );
+    const response = await fetch(url, { method, headers });
+
+    if (!response.ok) {
+      throw await ResponseError.fromResponse(response);
+    }
+  }
+
+  private async requestRequired(
+    method: string,
     path: string,
     options?: CatalogRequestOptions,
   ): Promise<any> {
     const url = `${await this.discoveryApi.getBaseUrl('catalog')}${path}`;
-    const response = await fetch(url, {
-      headers: options?.token
-        ? { Authorization: `Bearer ${options.token}` }
-        : {},
-    });
+    const headers = new Headers(
+      options?.token ? { Authorization: `Bearer ${options.token}` } : {},
+    );
+    const response = await fetch(url, { method, headers });
 
     if (!response.ok) {
       throw await ResponseError.fromResponse(response);
@@ -181,16 +219,16 @@ export class CatalogClient implements CatalogApi {
     return await response.json();
   }
 
-  private async getOptional(
+  private async requestOptional(
+    method: string,
     path: string,
     options?: CatalogRequestOptions,
   ): Promise<any | undefined> {
     const url = `${await this.discoveryApi.getBaseUrl('catalog')}${path}`;
-    const response = await fetch(url, {
-      headers: options?.token
-        ? { Authorization: `Bearer ${options.token}` }
-        : {},
-    });
+    const headers = new Headers(
+      options?.token ? { Authorization: `Bearer ${options.token}` } : {},
+    );
+    const response = await fetch(url, { method, headers });
 
     if (!response.ok) {
       if (response.status === 404) {

packages/catalog-client/src/types.ts

@@ -30,28 +30,39 @@ export type CatalogRequestOptions = {
 };
 
 export interface CatalogApi {
-  getLocationById(
-    id: String,
-    options?: CatalogRequestOptions,
-  ): Promise<Location | undefined>;
-  getEntityByName(
-    name: EntityName,
-    options?: CatalogRequestOptions,
-  ): Promise<Entity | undefined>;
+  // Entities
   getEntities(
     request?: CatalogEntitiesRequest,
     options?: CatalogRequestOptions,
   ): Promise<CatalogListResponse<Entity>>;
-  addLocation(
-    location: AddLocationRequest,
+  getEntityByName(
+    name: EntityName,
     options?: CatalogRequestOptions,
-  ): Promise<AddLocationResponse>;
+  ): Promise<Entity | undefined>;
+  removeEntityByUid(
+    uid: string,
+    options?: CatalogRequestOptions,
+  ): Promise<void>;
+
+  // Locations
+  getLocationById(
+    id: String,
+    options?: CatalogRequestOptions,
+  ): Promise<Location | undefined>;
+  getOriginLocationByEntity(
+    entity: Entity,
+    options?: CatalogRequestOptions,
+  ): Promise<Location | undefined>;
   getLocationByEntity(
     entity: Entity,
     options?: CatalogRequestOptions,
   ): Promise<Location | undefined>;
-  removeEntityByUid(
-    uid: string,
+  addLocation(
+    location: AddLocationRequest,
+    options?: CatalogRequestOptions,
+  ): Promise<AddLocationResponse>;
+  removeLocationById(
+    id: string,
     options?: CatalogRequestOptions,
   ): Promise<void>;
 }

plugins/auth-backend/src/lib/catalog/CatalogIdentityClient.test.ts

@@ -24,6 +24,8 @@ describe('CatalogIdentityClient', () => {
     getEntityByName: jest.fn(),
     getEntities: jest.fn(),
     addLocation: jest.fn(),
+    removeLocationById: jest.fn(),
+    getOriginLocationByEntity: jest.fn(),
     getLocationByEntity: jest.fn(),
     removeEntityByUid: jest.fn(),
   };

plugins/auth-backend/src/providers/aws-alb/provider.test.ts

@@ -70,7 +70,9 @@ describe('AwsALBAuthProvider', () => {
   const catalogApi = {
     /* eslint-disable-next-line @typescript-eslint/no-unused-vars */
     addLocation: jest.fn(),
+    removeLocationById: jest.fn(),
     getEntities: jest.fn(),
+    getOriginLocationByEntity: jest.fn(),
     getLocationByEntity: jest.fn(),
     getLocationById: jest.fn(),
     removeEntityByUid: jest.fn(),

plugins/badges-backend/src/service/router.test.ts

@@ -61,8 +61,10 @@ describe('createRouter', () => {
       addLocation: jest.fn(),
       getEntities: jest.fn(),
       getEntityByName: jest.fn(),
+      getOriginLocationByEntity: jest.fn(),
       getLocationByEntity: jest.fn(),
       getLocationById: jest.fn(),
+      removeLocationById: jest.fn(),
       removeEntityByUid: jest.fn(),
     };
 

plugins/catalog-backend/src/catalog/DatabaseEntitiesCatalog.ts

@@ -22,7 +22,7 @@ import {
   LOCATION_ANNOTATION,
   serializeEntityRef,
 } from '@backstage/catalog-model';
-import { ConflictError, NotFoundError } from '@backstage/errors';
+import { ConflictError } from '@backstage/errors';
 import { chunk, groupBy } from 'lodash';
 import limiterFactory from 'p-limit';
 import { Logger } from 'winston';
@@ -84,37 +84,8 @@ export class DatabaseEntitiesCatalog implements EntitiesCatalog {
   }
 
   async removeEntityByUid(uid: string): Promise<void> {
-    return await this.database.transaction(async tx => {
-      const entityResponse = await this.database.entityByUid(tx, uid);
-      if (!entityResponse) {
-        throw new NotFoundError(`Entity with ID ${uid} was not found`);
-      }
-
-      const location =
-        entityResponse.entity.metadata.annotations?.[LOCATION_ANNOTATION];
-
-      const colocatedEntities = location
-        ? (
-            await this.database.entities(tx, {
-              filter: basicEntityFilter({
-                [`metadata.annotations.${LOCATION_ANNOTATION}`]: location,
-              }),
-            })
-          ).entities
-        : [entityResponse];
-
-      for (const dbResponse of colocatedEntities) {
-        await this.database.removeEntityByUid(
-          tx,
-          dbResponse?.entity.metadata.uid!,
-        );
-      }
-
-      if (entityResponse.locationId) {
-        await this.database.removeLocation(tx, entityResponse?.locationId!);
-      }
-
-      return undefined;
+    await this.database.transaction(async tx => {
+      await this.database.removeEntityByUid(tx, uid);
     });
   }
 

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

@@ -317,7 +317,6 @@ export class CommonDatabase implements Database {
     const tx = txOpaque as Knex.Transaction;
 
     const result = await tx<DbEntitiesRow>('entities').where({ id: uid }).del();
-
     if (!result) {
       throw new NotFoundError(`Found no entity with ID ${uid}`);
     }

plugins/catalog-import/src/api/CatalogImportClient.test.ts

@@ -95,7 +95,9 @@ describe('CatalogImportClient', () => {
   const catalogApi: jest.Mocked<typeof catalogApiRef.T> = {
     getEntities: jest.fn(),
     addLocation: jest.fn(),
+    removeLocationById: jest.fn(),
     getEntityByName: jest.fn(),
+    getOriginLocationByEntity: jest.fn(),
     getLocationByEntity: jest.fn(),
     getLocationById: jest.fn(),
     removeEntityByUid: jest.fn(),

plugins/catalog-import/src/components/StepPrepareCreatePullRequest/StepPrepareCreatePullRequest.test.tsx

@@ -36,8 +36,10 @@ describe('<StepPrepareCreatePullRequest />', () => {
     getEntities: jest.fn(),
     addLocation: jest.fn(),
     getEntityByName: jest.fn(),
+    getOriginLocationByEntity: jest.fn(),
     getLocationByEntity: jest.fn(),
     getLocationById: jest.fn(),
+    removeLocationById: jest.fn(),
     removeEntityByUid: jest.fn(),
   };
 

plugins/catalog/package.json

@@ -33,6 +33,7 @@
     "@backstage/catalog-client": "^0.3.8",
     "@backstage/catalog-model": "^0.7.4",
     "@backstage/core": "^0.7.3",
+    "@backstage/errors": "^0.1.1",
     "@backstage/integration": "^0.5.1",
     "@backstage/integration-react": "^0.1.1",
     "@backstage/plugin-catalog-react": "^0.1.3",

plugins/catalog/src/CatalogClientWrapper.ts

@@ -77,6 +77,15 @@ export class CatalogClientWrapper implements CatalogApi {
     });
   }
 
+  async getOriginLocationByEntity(
+    entity: Entity,
+    options?: CatalogRequestOptions,
+  ): Promise<Location | undefined> {
+    return await this.client.getOriginLocationByEntity(entity, {
+      token: options?.token ?? (await this.identityApi.getIdToken()),
+    });
+  }
+
   async getLocationByEntity(
     entity: Entity,
     options?: CatalogRequestOptions,
@@ -86,6 +95,15 @@ export class CatalogClientWrapper implements CatalogApi {
     });
   }
 
+  async removeLocationById(
+    id: string,
+    options?: CatalogRequestOptions,
+  ): Promise<void> {
+    return await this.client.removeLocationById(id, {
+      token: options?.token ?? (await this.identityApi.getIdToken()),
+    });
+  }
+
   async removeEntityByUid(
     uid: string,
     options?: CatalogRequestOptions,

plugins/catalog/src/components/UnregisterEntityDialog/UnregisterEntityDialog.test.tsx

@@ -0,0 +1,310 @@
+/*
+ * 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.
+ */
+
+jest.mock('./useUnregisterEntityDialogState');
+
+import userEvent from '@testing-library/user-event';
+import React from 'react';
+import { UnregisterEntityDialog } from './UnregisterEntityDialog';
+import { ORIGIN_LOCATION_ANNOTATION } from '@backstage/catalog-model';
+import {
+  AlertApi,
+  alertApiRef,
+  ApiProvider,
+  ApiRegistry,
+  DiscoveryApi,
+} from '@backstage/core';
+import { CatalogClient } from '@backstage/catalog-client';
+import { catalogApiRef } from '@backstage/plugin-catalog-react';
+import { screen, waitFor } from '@testing-library/react';
+import { renderInTestApp } from '@backstage/test-utils';
+import * as state from './useUnregisterEntityDialogState';
+
+describe('UnregisterEntityDialog', () => {
+  const discoveryApi: DiscoveryApi = {
+    async getBaseUrl(pluginId) {
+      return `http://example.com/${pluginId}`;
+    },
+  };
+  const alertApi: AlertApi = {
+    post() {
+      return undefined;
+    },
+    alert$() {
+      throw new Error('not implemented');
+    },
+  };
+
+  const apis = ApiRegistry.with(
+    catalogApiRef,
+    new CatalogClient({ discoveryApi }),
+  ).with(alertApiRef, alertApi);
+
+  const entity = {
+    apiVersion: 'backstage.io/v1alpha1',
+    kind: 'Component',
+    metadata: {
+      name: 'n',
+      namespace: 'ns',
+      annotations: {
+        [ORIGIN_LOCATION_ANNOTATION]: 'url:http://example.com',
+      },
+    },
+    spec: {},
+  };
+
+  const Wrapper = ({ children }: { children?: React.ReactNode }) => (
+    <ApiProvider apis={apis}>{children}</ApiProvider>
+  );
+
+  const stateSpy = jest.spyOn(state, 'useUnregisterEntityDialogState');
+
+  afterEach(() => {
+    jest.resetAllMocks();
+  });
+
+  it('can cancel', async () => {
+    const onClose = jest.fn();
+    stateSpy.mockImplementation(() => ({ type: 'loading' }));
+
+    await renderInTestApp(
+      <Wrapper>
+        <UnregisterEntityDialog
+          open
+          onClose={onClose}
+          onConfirm={() => {}}
+          entity={entity}
+        />
+      </Wrapper>,
+    );
+
+    userEvent.click(screen.getByText('Cancel'));
+
+    await waitFor(() => {
+      expect(onClose).toBeCalled();
+    });
+  });
+
+  it('handles the loading state', async () => {
+    stateSpy.mockImplementation(() => ({ type: 'loading' }));
+
+    await renderInTestApp(
+      <Wrapper>
+        <UnregisterEntityDialog
+          open
+          onClose={() => {}}
+          onConfirm={() => {}}
+          entity={entity}
+        />
+      </Wrapper>,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByTestId('progress')).toBeInTheDocument();
+    });
+  });
+
+  it('handles the error state', async () => {
+    stateSpy.mockImplementation(() => ({
+      type: 'error',
+      error: new TypeError('eek!'),
+    }));
+
+    await renderInTestApp(
+      <Wrapper>
+        <UnregisterEntityDialog
+          open
+          onClose={() => {}}
+          onConfirm={() => {}}
+          entity={entity}
+        />
+      </Wrapper>,
+    );
+
+    await waitFor(() => {
+      expect(screen.getAllByText('eek!').length).toBeGreaterThan(0);
+      expect(screen.getAllByText('TypeError').length).toBeGreaterThan(0);
+    });
+  });
+
+  it('handles the bootstrap state', async () => {
+    const deleteEntity = jest.fn();
+    const onConfirm = jest.fn();
+
+    stateSpy.mockImplementation(() => ({
+      type: 'bootstrap',
+      location: 'bootstrap:bootstrap',
+      deleteEntity,
+    }));
+
+    await renderInTestApp(
+      <Wrapper>
+        <UnregisterEntityDialog
+          open
+          onClose={() => {}}
+          onConfirm={onConfirm}
+          entity={entity}
+        />
+      </Wrapper>,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText(/You cannot unregister/)).toBeInTheDocument();
+    });
+
+    userEvent.click(screen.getByText('Advanced Options'));
+
+    await waitFor(() => {
+      expect(screen.getByText(/option to delete/)).toBeInTheDocument();
+    });
+
+    userEvent.click(screen.getByText('Delete Entity'));
+
+    await waitFor(() => {
+      expect(deleteEntity).toBeCalled();
+      expect(onConfirm).toBeCalled();
+    });
+  });
+
+  it('handles the only-delete state', async () => {
+    const deleteEntity = jest.fn();
+    const onConfirm = jest.fn();
+
+    stateSpy.mockImplementation(() => ({
+      type: 'only-delete',
+      location: 'url:http://example.com',
+      deleteEntity,
+    }));
+
+    await renderInTestApp(
+      <Wrapper>
+        <UnregisterEntityDialog
+          open
+          onClose={() => {}}
+          onConfirm={onConfirm}
+          entity={entity}
+        />
+      </Wrapper>,
+    );
+
+    await waitFor(() => {
+      expect(
+        screen.getByText(/You therefore only have the option to delete it/),
+      ).toBeInTheDocument();
+    });
+
+    userEvent.click(screen.getByText('Delete Entity'));
+
+    await waitFor(() => {
+      expect(deleteEntity).toBeCalled();
+      expect(onConfirm).toBeCalled();
+    });
+  });
+
+  it('handles the unregister state, choosing to unregister', async () => {
+    const unregisterLocation = jest.fn();
+    const deleteEntity = jest.fn();
+    const onConfirm = jest.fn();
+
+    stateSpy.mockImplementation(() => ({
+      type: 'unregister',
+      location: 'url:http://example.com',
+      colocatedEntities: [
+        { kind: 'k1', namespace: 'ns1', name: 'n1' },
+        { kind: 'k2', namespace: 'ns2', name: 'n2' },
+      ],
+      unregisterLocation,
+      deleteEntity,
+    }));
+
+    await renderInTestApp(
+      <Wrapper>
+        <UnregisterEntityDialog
+          open
+          onClose={() => {}}
+          onConfirm={onConfirm}
+          entity={entity}
+        />
+      </Wrapper>,
+    );
+
+    await waitFor(() => {
+      expect(
+        screen.getByText(/will unregister the following entities/),
+      ).toBeInTheDocument();
+      expect(screen.getByText(/k1:ns1\/n1/)).toBeInTheDocument();
+      expect(screen.getByText(/k2:ns2\/n2/)).toBeInTheDocument();
+    });
+
+    userEvent.click(screen.getByText('Unregister Location'));
+
+    await waitFor(() => {
+      expect(unregisterLocation).toBeCalled();
+      expect(onConfirm).toBeCalled();
+    });
+  });
+
+  it('handles the unregister state, choosing to delete', async () => {
+    const unregisterLocation = jest.fn();
+    const deleteEntity = jest.fn();
+    const onConfirm = jest.fn();
+
+    stateSpy.mockImplementation(() => ({
+      type: 'unregister',
+      location: 'url:http://example.com',
+      colocatedEntities: [
+        { kind: 'k1', namespace: 'ns1', name: 'n1' },
+        { kind: 'k2', namespace: 'ns2', name: 'n2' },
+      ],
+      unregisterLocation,
+      deleteEntity,
+    }));
+
+    await renderInTestApp(
+      <Wrapper>
+        <UnregisterEntityDialog
+          open
+          onClose={() => {}}
+          onConfirm={onConfirm}
+          entity={entity}
+        />
+      </Wrapper>,
+    );
+
+    await waitFor(() => {
+      expect(
+        screen.getByText(/will unregister the following entities/),
+      ).toBeInTheDocument();
+      expect(screen.getByText(/k1:ns1\/n1/)).toBeInTheDocument();
+      expect(screen.getByText(/k2:ns2\/n2/)).toBeInTheDocument();
+    });
+
+    userEvent.click(screen.getByText('Advanced Options'));
+
+    await waitFor(() => {
+      expect(
+        screen.getByText(/You also have the option to delete/),
+      ).toBeInTheDocument();
+    });
+
+    userEvent.click(screen.getByText('Delete Entity'));
+
+    await waitFor(() => {
+      expect(deleteEntity).toBeCalled();
+      expect(onConfirm).toBeCalled();
+    });
+  });
+});

plugins/catalog/src/components/UnregisterEntityDialog/UnregisterEntityDialog.tsx

@@ -14,25 +14,35 @@
  * limitations under the License.
  */
 
-import { Entity, ORIGIN_LOCATION_ANNOTATION } from '@backstage/catalog-model';
-import { alertApiRef, configApiRef, Progress, useApi } from '@backstage/core';
+import { Entity } from '@backstage/catalog-model';
 import {
-  catalogApiRef,
-  formatEntityRefTitle,
-} from '@backstage/plugin-catalog-react';
+  alertApiRef,
+  configApiRef,
+  Progress,
+  ResponseErrorPanel,
+  useApi,
+} from '@backstage/core';
+import { EntityRefLink } from '@backstage/plugin-catalog-react';
 import {
+  Box,
   Button,
   Dialog,
   DialogActions,
   DialogContent,
   DialogContentText,
   DialogTitle,
-  Typography,
+  Divider,
+  makeStyles,
 } from '@material-ui/core';
 import Alert from '@material-ui/lab/Alert';
-import React from 'react';
-import { useAsync } from 'react-use';
-import { AsyncState } from 'react-use/lib/useAsync';
+import React, { useCallback, useState } from 'react';
+import { useUnregisterEntityDialogState } from './useUnregisterEntityDialogState';
+
+const useStyles = makeStyles({
+  advancedButton: {
+    fontSize: '0.7em',
+  },
+});
 
 type Props = {
   open: boolean;
@@ -41,129 +51,221 @@ type Props = {
   entity: Entity;
 };
 
-class DeniedLocationException extends Error {
-  constructor(public readonly locationName: string) {
-    super(`You may not remove the location ${locationName}`);
-    this.name = 'DeniedLocationException';
+const Contents = ({
+  entity,
+  onConfirm,
+}: {
+  entity: Entity;
+  onConfirm: () => any;
+}) => {
+  const alertApi = useApi(alertApiRef);
+  const configApi = useApi(configApiRef);
+  const classes = useStyles();
+  const state = useUnregisterEntityDialogState(entity);
+  const [showDelete, setShowDelete] = useState(false);
+  const [busy, setBusy] = useState(false);
+  const appTitle = configApi.getOptionalString('app.title') ?? 'Backstage';
+
+  const onUnregister = useCallback(
+    async function onUnregisterFn() {
+      if ('unregisterLocation' in state) {
+        setBusy(true);
+        try {
+          await state.unregisterLocation();
+          onConfirm();
+        } catch (err) {
+          alertApi.post({ message: err.message });
+        } finally {
+          setBusy(false);
+        }
+      }
+    },
+    [alertApi, onConfirm, state],
+  );
+
+  const onDelete = useCallback(
+    async function onDeleteFn() {
+      if ('deleteEntity' in state) {
+        setBusy(true);
+        try {
+          await state.deleteEntity();
+          onConfirm();
+        } catch (err) {
+          alertApi.post({ message: err.message });
+        } finally {
+          setBusy(false);
+        }
+      }
+    },
+    [alertApi, onConfirm, state],
+  );
+
+  if (state.type === 'loading') {
+    return <Progress />;
   }
-}
 
-function useColocatedEntities(entity: Entity): AsyncState<Entity[]> {
-  const catalogApi = useApi(catalogApiRef);
-  return useAsync(async () => {
-    const myLocation =
-      entity.metadata.annotations?.[ORIGIN_LOCATION_ANNOTATION];
-    if (!myLocation) {
-      return [];
-    }
+  if (state.type === 'error') {
+    return <ResponseErrorPanel error={state.error} />;
+  }
 
-    if (myLocation === 'bootstrap:bootstrap') {
-      throw new DeniedLocationException(myLocation);
-    }
+  if (state.type === 'bootstrap') {
+    return (
+      <>
+        <Alert severity="info">
+          You cannot unregister this entity, since it originates from a
+          protected Backstage configuration (location "{state.location}"). If
+          you believe this is in error, please contact the {appTitle}{' '}
+          integrator.
+        </Alert>
 
-    const response = await catalogApi.getEntities({
-      filter: {
-        [`metadata.annotations.${ORIGIN_LOCATION_ANNOTATION}`]: myLocation,
-      },
-    });
-    return response.items;
-  }, [catalogApi, entity]);
-}
+        <Box marginTop={2}>
+          {!showDelete && (
+            <Button
+              variant="text"
+              size="small"
+              color="primary"
+              className={classes.advancedButton}
+              onClick={() => setShowDelete(true)}
+            >
+              Advanced Options
+            </Button>
+          )}
+
+          {showDelete && (
+            <>
+              <DialogContentText>
+                You have the option to delete the entity itself from the
+                catalog. Note that this should only be done if you know that the
+                catalog file has been deleted at, or moved from, its origin
+                location. If that is not the case, the entity will reappear
+                shortly as the next refresh round is performed by the catalog.
+              </DialogContentText>
+              <Button
+                variant="contained"
+                color="secondary"
+                disabled={busy}
+                onClick={onDelete}
+              >
+                Delete Entity
+              </Button>
+            </>
+          )}
+        </Box>
+      </>
+    );
+  }
+
+  if (state.type === 'only-delete') {
+    return (
+      <>
+        <DialogContentText>
+          This entity does not seem to originate from a location. You therefore
+          only have the option to delete it outright from the catalog.
+        </DialogContentText>
+        <Button
+          variant="contained"
+          color="secondary"
+          disabled={busy}
+          onClick={onDelete}
+        >
+          Delete Entity
+        </Button>
+      </>
+    );
+  }
+
+  if (state.type === 'unregister') {
+    return (
+      <>
+        <DialogContentText>
+          This action will unregister the following entities:
+        </DialogContentText>
+        <DialogContentText component="ul">
+          {state.colocatedEntities.map(e => (
+            <li key={`${e.kind}:${e.namespace}/${e.name}`}>
+              <EntityRefLink entityRef={e} />
+            </li>
+          ))}
+        </DialogContentText>
+        <DialogContentText>
+          Located at the following location:
+        </DialogContentText>
+        <DialogContentText component="ul">
+          <li>{state.location}</li>
+        </DialogContentText>
+        <DialogContentText>
+          To undo, just re-register the entity in {appTitle}.
+        </DialogContentText>
+        <Box marginTop={2}>
+          <Button
+            variant="contained"
+            color="secondary"
+            disabled={busy}
+            onClick={onUnregister}
+          >
+            Unregister Location
+          </Button>
+          {!showDelete && (
+            <Box component="span" marginLeft={2}>
+              <Button
+                variant="text"
+                size="small"
+                color="primary"
+                className={classes.advancedButton}
+                onClick={() => setShowDelete(true)}
+              >
+                Advanced Options
+              </Button>
+            </Box>
+          )}
+        </Box>
+
+        {showDelete && (
+          <>
+            <Box paddingTop={4} paddingBottom={4}>
+              <Divider />
+            </Box>
+            <DialogContentText>
+              You also have the option to delete the entity itself from the
+              catalog. Note that this should only be done if you know that the
+              catalog file has been deleted at, or moved from, its origin
+              location. If that is not the case, the entity will reappear
+              shortly as the next refresh round is performed by the catalog.
+            </DialogContentText>
+            <Button
+              variant="contained"
+              color="secondary"
+              disabled={busy}
+              onClick={onDelete}
+            >
+              Delete Entity
+            </Button>
+          </>
+        )}
+      </>
+    );
+  }
+
+  return <Alert severity="error">Internal error: Unknown state</Alert>;
+};
 
 export const UnregisterEntityDialog = ({
   open,
   onConfirm,
   onClose,
   entity,
-}: Props) => {
-  const { value: entities, loading, error } = useColocatedEntities(entity);
-  const catalogApi = useApi(catalogApiRef);
-  const alertApi = useApi(alertApiRef);
-  const configApi = useApi(configApiRef);
-
-  const removeEntity = async () => {
-    const uid = entity.metadata.uid;
-    try {
-      await catalogApi.removeEntityByUid(uid!);
-    } catch (err) {
-      alertApi.post({ message: err.message });
-    }
-
-    onConfirm();
-  };
-
-  return (
-    <Dialog open={open} onClose={onClose}>
-      <DialogTitle id="responsive-dialog-title">
-        Are you sure you want to unregister this entity?
-      </DialogTitle>
-
-      <DialogContent>
-        {loading ? <Progress /> : null}
-
-        {error ? (
-          <Alert severity="error" style={{ wordBreak: 'break-word' }}>
-            {error.name === 'DeniedLocationException' ? (
-              <>
-                You cannot unregister this entity, since it originates from a
-                protected Backstage configuration (location{' '}
-                {`"${(error as DeniedLocationException).locationName}"`}). If
-                you believe this is in error, please contact the{' '}
-                {configApi.getOptionalString('app.title') ?? 'Backstage'}{' '}
-                integrator.
-              </>
-            ) : (
-              error.toString()
-            )}
-          </Alert>
-        ) : null}
-
-        {entities?.length ? (
-          <>
-            <DialogContentText>
-              This action will unregister the following entities:
-            </DialogContentText>
-            <Typography component="div">
-              <ul>
-                {entities.map(e => {
-                  const fullName = formatEntityRefTitle(e);
-                  return <li key={fullName}>{fullName}</li>;
-                })}
-              </ul>
-            </Typography>
-            <DialogContentText>
-              That are located at the following location:
-            </DialogContentText>
-            <Typography component="div">
-              <ul style={{ wordBreak: 'break-word' }}>
-                <li>
-                  {
-                    entities[0]?.metadata.annotations?.[
-                      ORIGIN_LOCATION_ANNOTATION
-                    ]
-                  }
-                </li>
-              </ul>
-            </Typography>
-            <DialogContentText>
-              To undo, just re-register the entity in Backstage.
-            </DialogContentText>
-          </>
-        ) : null}
-      </DialogContent>
-
-      <DialogActions>
-        <Button onClick={onClose} color="primary">
-          Cancel
-        </Button>
-        <Button
-          disabled={!!(loading || error)}
-          onClick={removeEntity}
-          color="secondary"
-        >
-          Unregister
-        </Button>
-      </DialogActions>
-    </Dialog>
-  );
-};
+}: Props) => (
+  <Dialog open={open} onClose={onClose}>
+    <DialogTitle id="responsive-dialog-title">
+      Are you sure you want to unregister this entity?
+    </DialogTitle>
+    <DialogContent>
+      <Contents entity={entity} onConfirm={onConfirm} />
+    </DialogContent>
+    <DialogActions>
+      <Button onClick={onClose} color="primary">
+        Cancel
+      </Button>
+    </DialogActions>
+  </Dialog>
+);

plugins/catalog/src/components/UnregisterEntityDialog/useUnregisterEntityDialogState.test.tsx

@@ -0,0 +1,183 @@
+/*
+ * 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 {
+  Entity,
+  Location,
+  ORIGIN_LOCATION_ANNOTATION,
+} from '@backstage/catalog-model';
+import { ApiProvider, ApiRegistry } from '@backstage/core';
+import { CatalogApi, catalogApiRef } from '@backstage/plugin-catalog-react';
+import {
+  act,
+  renderHook,
+  RenderHookResult,
+} from '@testing-library/react-hooks';
+import React from 'react';
+import {
+  UseUnregisterEntityDialogState,
+  useUnregisterEntityDialogState,
+} from './useUnregisterEntityDialogState';
+
+function defer<T>(): { promise: Promise<T>; resolve: (value: T) => void } {
+  let resolve: (value: T) => void = () => {};
+  const promise = new Promise<T>(_resolve => {
+    resolve = _resolve;
+  });
+  return { promise, resolve };
+}
+
+describe('useUnregisterEntityDialogState', () => {
+  const catalogApiMock = {
+    getOriginLocationByEntity: jest.fn(),
+    getEntities: jest.fn(),
+    removeLocationById: jest.fn(),
+    removeEntityByUid: jest.fn(),
+  };
+  const catalogApi = (catalogApiMock as Partial<CatalogApi>) as CatalogApi;
+
+  const Wrapper = ({ children }: { children?: React.ReactNode }) => (
+    <ApiProvider apis={ApiRegistry.with(catalogApiRef, catalogApi)}>
+      {children}
+    </ApiProvider>
+  );
+
+  let entity: Entity;
+  let resolveLocation: (location: Location | undefined) => void;
+  let resolveColocatedEntities: (entities: Entity[]) => void;
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+
+    const deferredLocation = defer<Location | undefined>();
+    const deferredColocatedEntities = defer<Entity[]>();
+
+    resolveLocation = deferredLocation.resolve;
+    resolveColocatedEntities = deferredColocatedEntities.resolve;
+
+    catalogApiMock.getOriginLocationByEntity.mockReturnValue(
+      deferredLocation.promise,
+    );
+    catalogApiMock.getEntities.mockReturnValue(
+      deferredColocatedEntities.promise.then(items => ({ items })),
+    );
+
+    entity = {
+      apiVersion: 'backstage.io/v1alpha1',
+      kind: 'Component',
+      metadata: {
+        name: 'n',
+        namespace: 'ns',
+        annotations: {
+          [ORIGIN_LOCATION_ANNOTATION]: 'url:https://example.com',
+        },
+      },
+      spec: {},
+    };
+  });
+
+  it('goes through the happy unregister path', async () => {
+    let rendered: RenderHookResult<unknown, UseUnregisterEntityDialogState>;
+    act(() => {
+      rendered = renderHook(() => useUnregisterEntityDialogState(entity), {
+        wrapper: Wrapper,
+      });
+    });
+
+    expect(rendered!.result.current).toEqual({ type: 'loading' });
+
+    resolveLocation({ type: 'url', target: 'https://example.com', id: 'x' });
+    resolveColocatedEntities([entity]);
+
+    await act(async () => {
+      await rendered!.waitForNextUpdate();
+    });
+
+    expect(rendered!.result.current).toEqual({
+      type: 'unregister',
+      location: 'url:https://example.com',
+      colocatedEntities: [{ kind: 'Component', namespace: 'ns', name: 'n' }],
+      unregisterLocation: expect.any(Function),
+      deleteEntity: expect.any(Function),
+    });
+  });
+
+  it('chooses the bootstrap path when necessary', async () => {
+    entity.metadata.annotations![ORIGIN_LOCATION_ANNOTATION] =
+      'bootstrap:bootstrap';
+
+    let rendered: RenderHookResult<unknown, UseUnregisterEntityDialogState>;
+    act(() => {
+      rendered = renderHook(() => useUnregisterEntityDialogState(entity), {
+        wrapper: Wrapper,
+      });
+    });
+
+    resolveLocation({ type: 'bootstrap', target: 'bootstrap', id: 'x' });
+    resolveColocatedEntities([]);
+    await act(async () => {
+      await rendered!.waitForNextUpdate();
+    });
+
+    expect(rendered!.result.current).toEqual({
+      type: 'bootstrap',
+      location: 'bootstrap:bootstrap',
+      deleteEntity: expect.any(Function),
+    });
+  });
+
+  it('chooses only-delete when there was no location annotation', async () => {
+    delete entity.metadata.annotations![ORIGIN_LOCATION_ANNOTATION];
+
+    let rendered: RenderHookResult<unknown, UseUnregisterEntityDialogState>;
+    act(() => {
+      rendered = renderHook(() => useUnregisterEntityDialogState(entity), {
+        wrapper: Wrapper,
+      });
+    });
+
+    resolveLocation(undefined);
+    resolveColocatedEntities([]);
+    await act(async () => {
+      await rendered!.waitForNextUpdate();
+    });
+
+    expect(rendered!.result.current).toEqual({
+      type: 'only-delete',
+      deleteEntity: expect.any(Function),
+    });
+  });
+
+  it('chooses only-delete when the location could not be found', async () => {
+    let rendered: RenderHookResult<unknown, UseUnregisterEntityDialogState>;
+    act(() => {
+      rendered = renderHook(() => useUnregisterEntityDialogState(entity), {
+        wrapper: Wrapper,
+      });
+    });
+
+    resolveLocation(undefined);
+    resolveColocatedEntities([]);
+    await act(async () => {
+      await rendered!.waitForNextUpdate();
+    });
+
+    expect(rendered!.result.current).toEqual({
+      type: 'only-delete',
+      deleteEntity: expect.any(Function),
+    });
+  });
+});

plugins/catalog/src/components/UnregisterEntityDialog/useUnregisterEntityDialogState.ts

@@ -0,0 +1,148 @@
+/*
+ * 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 {
+  Entity,
+  EntityName,
+  getEntityName,
+  ORIGIN_LOCATION_ANNOTATION,
+} from '@backstage/catalog-model';
+import { useApi } from '@backstage/core';
+import { catalogApiRef } from '@backstage/plugin-catalog-react';
+import { useCallback } from 'react';
+import { useAsync } from 'react-use';
+
+/**
+ * Each distinct state that the dialog can be in at any given time.
+ */
+export type UseUnregisterEntityDialogState =
+  | {
+      type: 'loading';
+    }
+  | {
+      type: 'error';
+      error: Error;
+    }
+  | {
+      type: 'bootstrap';
+      location: string;
+      deleteEntity: () => Promise<void>;
+    }
+  | {
+      type: 'unregister';
+      location: string;
+      colocatedEntities: EntityName[];
+      unregisterLocation: () => Promise<void>;
+      deleteEntity: () => Promise<void>;
+    }
+  | {
+      type: 'only-delete';
+      deleteEntity: () => Promise<void>;
+    };
+
+/**
+ * Houses the main logic for unregistering entities and their locations.
+ */
+export function useUnregisterEntityDialogState(
+  entity: Entity,
+): UseUnregisterEntityDialogState {
+  const catalogApi = useApi(catalogApiRef);
+  const locationRef = entity.metadata.annotations?.[ORIGIN_LOCATION_ANNOTATION];
+  const uid = entity.metadata.uid;
+  const isBootstrap = locationRef === 'bootstrap:bootstrap';
+
+  // Load the prerequisite data: what entities that are colocated with us, and
+  // what location that spawned us
+  const prerequisites = useAsync(async () => {
+    const locationPromise = catalogApi.getOriginLocationByEntity(entity);
+
+    let colocatedEntitiesPromise: Promise<Entity[]>;
+    if (!locationRef) {
+      colocatedEntitiesPromise = Promise.resolve([]);
+    } else {
+      const locationAnnotationFilter = `metadata.annotations.${ORIGIN_LOCATION_ANNOTATION}`;
+      colocatedEntitiesPromise = catalogApi
+        .getEntities({
+          filter: { [locationAnnotationFilter]: locationRef },
+          fields: [
+            'kind',
+            'metadata.uid',
+            'metadata.name',
+            'metadata.namespace',
+          ],
+        })
+        .then(response => response.items);
+    }
+
+    return Promise.all([locationPromise, colocatedEntitiesPromise]).then(
+      ([location, colocatedEntities]) => ({
+        location,
+        colocatedEntities,
+      }),
+    );
+  }, [catalogApi, entity]);
+
+  // Unregisters the underlying location and removes all of the entities that
+  // are spawned from it. Can only ever be called when the prerequisites have
+  // finished loading successfully, and if there was a matching location.
+  const unregisterLocation = useCallback(
+    async function unregisterLocationFn() {
+      const { location, colocatedEntities } = prerequisites.value!;
+      await catalogApi.removeLocationById(location!.id);
+      await Promise.allSettled(
+        colocatedEntities.map(e =>
+          catalogApi.removeEntityByUid(e.metadata.uid!),
+        ),
+      );
+    },
+    [catalogApi, prerequisites],
+  );
+
+  // Just removes the entity, without affecting locations in any way.
+  const deleteEntity = useCallback(
+    async function deleteEntityFn() {
+      await catalogApi.removeEntityByUid(uid!);
+    },
+    [catalogApi, uid],
+  );
+
+  // If this is a bootstrap location entity, don't even block on loading
+  // prerequisites. We know that all that we will do is to offer to remove the
+  // entity, and that doesn't require anything from the prerequisites.
+  if (isBootstrap) {
+    return { type: 'bootstrap', location: locationRef!, deleteEntity };
+  }
+
+  // Return early if prerequisites still loading or failing
+  const { loading, error, value } = prerequisites;
+  if (loading) {
+    return { type: 'loading' };
+  } else if (error) {
+    return { type: 'error', error };
+  }
+
+  const { location, colocatedEntities } = value!;
+  if (!location) {
+    return { type: 'only-delete', deleteEntity };
+  }
+  return {
+    type: 'unregister',
+    location: locationRef!,
+    colocatedEntities: colocatedEntities.map(getEntityName),
+    unregisterLocation,
+    deleteEntity,
+  };
+}

plugins/catalog/src/filter/useEntityFilterGroup.test.tsx

@@ -32,8 +32,10 @@ describe('useEntityFilterGroup', () => {
       /* eslint-disable-next-line @typescript-eslint/no-unused-vars */
       addLocation: jest.fn(_a => new Promise(() => {})),
       getEntities: jest.fn(),
+      getOriginLocationByEntity: jest.fn(),
       getLocationByEntity: jest.fn(),
       getLocationById: jest.fn(),
+      removeLocationById: jest.fn(),
       removeEntityByUid: jest.fn(),
       getEntityByName: jest.fn(),
     };

plugins/explore/src/components/DomainExplorerContent/DomainExplorerContent.test.tsx

@@ -27,8 +27,10 @@ describe('<DomainExplorerContent />', () => {
   const catalogApi: jest.Mocked<typeof catalogApiRef.T> = {
     addLocation: jest.fn(_a => new Promise(() => {})),
     getEntities: jest.fn(),
+    getOriginLocationByEntity: jest.fn(),
     getLocationByEntity: jest.fn(),
     getLocationById: jest.fn(),
+    removeLocationById: jest.fn(),
     removeEntityByUid: jest.fn(),
     getEntityByName: jest.fn(),
   };

plugins/register-component/src/components/RegisterComponentPage/RegisterComponentPage.test.tsx

@@ -36,8 +36,10 @@ const catalogApi: jest.Mocked<typeof catalogApiRef.T> = {
   /* eslint-disable-next-line @typescript-eslint/no-unused-vars */
   addLocation: jest.fn(_a => new Promise(() => {})),
   getEntities: jest.fn(),
+  getOriginLocationByEntity: jest.fn(),
   getLocationByEntity: jest.fn(),
   getLocationById: jest.fn(),
+  removeLocationById: jest.fn(),
   removeEntityByUid: jest.fn(),
   getEntityByName: jest.fn(),
 };

plugins/scaffolder/src/filter/useEntityFilterGroup.test.tsx

@@ -32,8 +32,10 @@ describe('useEntityFilterGroup', () => {
       /* eslint-disable-next-line @typescript-eslint/no-unused-vars */
       addLocation: jest.fn(_a => new Promise(() => {})),
       getEntities: jest.fn(),
+      getOriginLocationByEntity: jest.fn(),
       getLocationByEntity: jest.fn(),
       getLocationById: jest.fn(),
+      removeLocationById: jest.fn(),
       removeEntityByUid: jest.fn(),
       getEntityByName: jest.fn(),
     };

plugins/todo-backend/src/service/TodoReaderService.test.ts

@@ -45,8 +45,10 @@ function mockCatalogClient(entity?: Entity): jest.Mocked<CatalogApi> {
     addLocation: jest.fn(),
     getEntities: jest.fn(),
     getEntityByName: jest.fn(),
+    getOriginLocationByEntity: jest.fn(),
     getLocationByEntity: jest.fn(),
     getLocationById: jest.fn(),
+    removeLocationById: jest.fn(),
     removeEntityByUid: jest.fn(),
   };
   if (entity) {