From 2214202d90cf88fc8b50968a662f7906a6568d42 Mon Sep 17 00:00:00 2001
From: Rickard Dybeck <dybeck@spotify.com>
Date: Fri, 29 Dec 2023 10:25:53 -0500
Subject: [PATCH] review fixups

Signed-off-by: Rickard Dybeck <dybeck@spotify.com>
---
 packages/catalog-client/src/CatalogClient.ts               | 7 +++++--
 packages/catalog-client/src/types/api.ts                   | 4 ++--
 .../src/modules/core/DefaultLocationStore.ts               | 4 ++--
 .../src/service/AuthorizedLocationService.ts               | 2 +-
 .../catalog-backend/src/service/DefaultLocationService.ts  | 7 +++++--
 plugins/catalog-backend/src/service/types.ts               | 4 ++--
 6 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/packages/catalog-client/src/CatalogClient.ts b/packages/catalog-client/src/CatalogClient.ts
index eb99fa7830..07fedd0e9c 100644
--- a/packages/catalog-client/src/CatalogClient.ts
+++ b/packages/catalog-client/src/CatalogClient.ts
@@ -92,11 +92,14 @@ export class CatalogClient implements CatalogApi {
    * {@inheritdoc CatalogApi.getLocationByEntity}
    */
   async getLocationByEntity(
-    entityRef: CompoundEntityRef,
+    entityRef: CompoundEntityRef | string,
     options?: CatalogRequestOptions,
   ): Promise<Location | undefined> {
     return await this.requestOptional(
-      await this.apiClient.getLocationByEntity({ path: entityRef }, options),
+      await this.apiClient.getLocationByEntity(
+        { path: parseEntityRef(entityRef) },
+        options,
+      ),
     );
   }
 
diff --git a/packages/catalog-client/src/types/api.ts b/packages/catalog-client/src/types/api.ts
index 6d9f95b940..ecf2be7f44 100644
--- a/packages/catalog-client/src/types/api.ts
+++ b/packages/catalog-client/src/types/api.ts
@@ -632,11 +632,11 @@ export interface CatalogApi {
   /**
    * Gets a location associated with an entity.
    *
-   * @param entityRef - A reference to an entity
+   * @param entityRef - A complete entity ref, either on string or compound form
    * @param options - Additional options
    */
   getLocationByEntity(
-    entityRef: CompoundEntityRef,
+    entityRef: string | CompoundEntityRef,
     options?: CatalogRequestOptions,
   ): Promise<Location | undefined>;
 
diff --git a/plugins/catalog-backend/src/modules/core/DefaultLocationStore.ts b/plugins/catalog-backend/src/modules/core/DefaultLocationStore.ts
index 267abce046..ac506051d2 100644
--- a/plugins/catalog-backend/src/modules/core/DefaultLocationStore.ts
+++ b/plugins/catalog-backend/src/modules/core/DefaultLocationStore.ts
@@ -33,6 +33,7 @@ import { LocationInput, LocationStore } from '../../service/types';
 import {
   ANNOTATION_ORIGIN_LOCATION,
   CompoundEntityRef,
+  parseLocationRef,
   stringifyEntityRef,
 } from '@backstage/catalog-model';
 
@@ -143,9 +144,8 @@ export class DefaultLocationStore implements LocationStore, EntityProvider {
         `found no origin annotation for ref ${entityRefString}`,
       );
     }
-    const [type, ...rest] = locationKeyValue.value?.split(':') ?? [];
-    const target = rest.join(':');
 
+    const { type, target } = parseLocationRef(entityRefString);
     // const kind, target = split[0], split[1];
     const [location] = await this.db<DbLocationsRow>('locations')
       .where({ type, target })
diff --git a/plugins/catalog-backend/src/service/AuthorizedLocationService.ts b/plugins/catalog-backend/src/service/AuthorizedLocationService.ts
index fe4547a8fc..0b9d50b20e 100644
--- a/plugins/catalog-backend/src/service/AuthorizedLocationService.ts
+++ b/plugins/catalog-backend/src/service/AuthorizedLocationService.ts
@@ -113,7 +113,7 @@ export class AuthorizedLocationService implements LocationService {
   }
 
   async getLocationByEntity(
-    entityRef: CompoundEntityRef,
+    entityRef: CompoundEntityRef | string,
     options?: { authorizationToken?: string | undefined } | undefined,
   ): Promise<Location> {
     const authorizationResponse = (
diff --git a/plugins/catalog-backend/src/service/DefaultLocationService.ts b/plugins/catalog-backend/src/service/DefaultLocationService.ts
index 5ce9456212..d0b603b90b 100644
--- a/plugins/catalog-backend/src/service/DefaultLocationService.ts
+++ b/plugins/catalog-backend/src/service/DefaultLocationService.ts
@@ -20,6 +20,7 @@ import {
   ANNOTATION_ORIGIN_LOCATION,
   stringifyEntityRef,
   CompoundEntityRef,
+  parseEntityRef,
 } from '@backstage/catalog-model';
 import { Location } from '@backstage/catalog-client';
 import { CatalogProcessingOrchestrator } from '../processing/types';
@@ -69,8 +70,10 @@ export class DefaultLocationService implements LocationService {
     return this.store.deleteLocation(id);
   }
 
-  getLocationByEntity(entityRef: CompoundEntityRef): Promise<Location> {
-    return this.store.getLocationByEntity(entityRef);
+  getLocationByEntity(
+    entityRef: CompoundEntityRef | string,
+  ): Promise<Location> {
+    return this.store.getLocationByEntity(parseEntityRef(entityRef));
   }
 
   private async processEntities(
diff --git a/plugins/catalog-backend/src/service/types.ts b/plugins/catalog-backend/src/service/types.ts
index db4dfc7f34..878afe5662 100644
--- a/plugins/catalog-backend/src/service/types.ts
+++ b/plugins/catalog-backend/src/service/types.ts
@@ -49,7 +49,7 @@ export interface LocationService {
     options?: { authorizationToken?: string },
   ): Promise<void>;
   getLocationByEntity(
-    entityRef: CompoundEntityRef,
+    entityRef: CompoundEntityRef | string,
     options?: { authorizationToken?: string },
   ): Promise<Location>;
 }
@@ -86,5 +86,5 @@ export interface LocationStore {
   listLocations(): Promise<Location[]>;
   getLocation(id: string): Promise<Location>;
   deleteLocation(id: string): Promise<void>;
-  getLocationByEntity(entityRef: CompoundEntityRef): Promise<Location>;
+  getLocationByEntity(entityRef: CompoundEntityRef | string): Promise<Location>;
 }