From b7711bf21ae25f7faf97aab945e5fa6c35c1de16 Mon Sep 17 00:00:00 2001
From: Alberto Miorin <32617+amiorin@users.noreply.github.com>
Date: Fri, 3 Dec 2021 03:54:14 +0100
Subject: [PATCH] Fix #1527: List locations is not including the service and
 the serviceType in the result (#1530)

---
 .../catalog/jdbi3/LocationRepository.java     | 23 +++++++++++--------
 .../locations/LocationResourceTest.java       |  3 ++-
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/LocationRepository.java b/catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/LocationRepository.java
index b13aa2bbb82..a6f3d8841f2 100644
--- a/catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/LocationRepository.java
+++ b/catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/LocationRepository.java
@@ -17,6 +17,7 @@ import com.fasterxml.jackson.core.JsonProcessingException;
 import org.jdbi.v3.sqlobject.transaction.Transaction;
 import org.openmetadata.catalog.Entity;
 import org.openmetadata.catalog.entity.data.Location;
+import org.openmetadata.catalog.entity.services.StorageService;
 import org.openmetadata.catalog.exception.CatalogExceptionMessage;
 import org.openmetadata.catalog.resources.locations.LocationResource;
 import org.openmetadata.catalog.type.ChangeDescription;
@@ -57,8 +58,8 @@ public class LocationRepository extends EntityRepository<Location> {
 
   @Override
   public Location setFields(Location location, Fields fields) throws IOException {
+    location.setService(getService(location));
     location.setOwner(fields.contains("owner") ? getOwner(location) : null);
-    location.setService(fields.contains("service") ? getService(location) : null);
     location.setFollowers(fields.contains("followers") ? getFollowers(location) : null);
     location.setTags(fields.contains("tags") ? getTags(location.getFullyQualifiedName()) : null);
     return location;
@@ -141,18 +142,22 @@ public class LocationRepository extends EntityRepository<Location> {
   public EntityReference getOwnerReference(Location location) throws IOException {
     return EntityUtil.populateOwner(dao.userDAO(), dao.teamDAO(), location.getOwner());
   }
+  
+  private StorageService getService(UUID serviceId, String entityType) throws IOException {
+    if (entityType.equalsIgnoreCase(Entity.STORAGE_SERVICE)) {
+      return dao.storageServiceDAO().findEntityById(serviceId);
+    }
+    throw new IllegalArgumentException(CatalogExceptionMessage.invalidServiceEntity(entityType, Entity.LOCATION));
+  }
 
   @Override
   public void prepare(Location location) throws IOException {
-    EntityReference storageService = getService(location.getService());
-    location.setService(storageService);
-    // Set data in location entity based on storage relationship
+    StorageService storageService = getService(location.getService().getId(), location.getService().getType());
+    location.setService(new StorageServiceRepository.StorageServiceEntityInterface(storageService)
+            .getEntityReference());
+    location.setServiceType(storageService.getServiceType());
     location.setFullyQualifiedName(getFQN(location));
-
-    // Check if owner is valid and set the relationship
-    location.setOwner(EntityUtil.populateOwner(dao.userDAO(), dao.teamDAO(), location.getOwner()));
-
-    // Validate location tags and add derived tags to the list
+    EntityUtil.populateOwner(dao.userDAO(), dao.teamDAO(), location.getOwner()); // Validate owner
     location.setTags(EntityUtil.addDerivedTags(dao.tagDAO(), location.getTags()));
   }
 
diff --git a/catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/locations/LocationResourceTest.java b/catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/locations/LocationResourceTest.java
index 5c67c4650b3..12675f1ba48 100644
--- a/catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/locations/LocationResourceTest.java
+++ b/catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/locations/LocationResourceTest.java
@@ -674,7 +674,8 @@ public class LocationResourceTest extends CatalogApplicationTest {
     location = byName ? getLocationByName(location.getFullyQualifiedName(), fields, adminAuthHeaders()) :
             getLocation(location.getId(), fields, adminAuthHeaders());
     assertNotNull(location.getOwner());
-    assertNull(location.getService());
+    assertNotNull(location.getService()); // We always return the service
+    assertNotNull(location.getServiceType()); // We always return the service
 
     // .../locations?fields=owner,service
     fields = "owner,service";