feat(openapi-v3): add additional delete options (#11347)

datahub-project · Sep 12, 2024 · 2890b6b · 2890b6b
1 parent d385076
commit 2890b6b
Show file tree

Hide file tree

Showing 5 changed files with 100 additions and 10 deletions.
diff --git a/...servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java b/...servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java
@@ -406,10 +406,11 @@ public ResponseEntity<Object> headAspect(
   public void deleteEntity(
       HttpServletRequest request,
       @PathVariable("entityName") String entityName,
-      @PathVariable("entityUrn") String entityUrn)
+      @PathVariable("entityUrn") String entityUrn,
+      @RequestParam(value = "aspects", required = false, defaultValue = "") Set<String> aspects,
+      @RequestParam(value = "clear", required = false, defaultValue = "false") boolean clear)
       throws InvalidUrnException {
 
-    EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName);
     Urn urn = validatedUrn(entityUrn);
     Authentication authentication = AuthenticationContext.getAuthentication();
     OperationContext opContext =
@@ -427,7 +428,26 @@ public void deleteEntity(
           authentication.getActor().toUrnStr() + " is unauthorized to " + DELETE + " entities.");
     }
 
-    entityService.deleteUrn(opContext, urn);
+    EntitySpec entitySpec = entityRegistry.getEntitySpec(urn.getEntityType());
+
+    if (clear) {
+      // remove all aspects, preserve entity by retaining key aspect
+      aspects =
+          entitySpec.getAspectSpecs().stream()
+              .map(AspectSpec::getName)
+              .filter(name -> !name.equals(entitySpec.getKeyAspectName()))
+              .collect(Collectors.toSet());
+    }
+
+    if (aspects == null || aspects.isEmpty() || aspects.contains(entitySpec.getKeyAspectName())) {
+      entityService.deleteUrn(opContext, urn);
+    } else {
+      aspects.stream()
+          .map(aspectName -> lookupAspectSpec(urn, aspectName).getName())
+          .forEach(
+              aspectName ->
+                  entityService.deleteAspect(opContext, entityUrn, aspectName, Map.of(), true));
+    }
   }
 
   @Tag(name = "Generic Entities")

diff --git a/...ervice/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/OpenAPIV3Generator.java b/...ervice/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/OpenAPIV3Generator.java
@@ -257,7 +257,17 @@ private static PathItem buildSingleEntityPath(final EntitySpec entity) {
                         .in(NAME_PATH)
                         .name("urn")
                         .description("The entity's unique URN id.")
-                        .schema(new Schema().type(TYPE_STRING))))
+                        .schema(new Schema().type(TYPE_STRING)),
+                    new Parameter()
+                        .in(NAME_QUERY)
+                        .name("clear")
+                        .description("Delete all aspects, preserving the entity's key aspect.")
+                        .schema(new Schema().type(TYPE_BOOLEAN)._default(false)),
+                    new Parameter()
+                        .$ref(
+                            String.format(
+                                "#/components/parameters/%s",
+                                aspectParameterName + MODEL_VERSION))))
             .tags(List.of(entity.getName() + " Entity"))
             .responses(new ApiResponses().addApiResponse("200", successDeleteResponse));
 
@@ -507,13 +517,13 @@ private static Parameter buildParameterSchema(
             .items(
                 new Schema()
                     .type(TYPE_STRING)
-                    ._enum(aspectNames)
+                    ._enum(aspectNames.stream().sorted().toList())
                     ._default(aspectNames.stream().findFirst().orElse(null)));
     return new Parameter()
         .in(NAME_QUERY)
         .name("aspects")
         .explode(true)
-        .description("Aspects to include in response.")
+        .description("Aspects to include.")
         .example(aspectNames)
         .schema(schema);
   }

diff --git a/...i-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java b/...i-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java
@@ -1,12 +1,16 @@
 package io.datahubproject.openapi.v3.controller;
 
+import static com.linkedin.metadata.Constants.DATASET_ENTITY_NAME;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyMap;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
 import static org.testng.Assert.assertNotNull;
@@ -26,6 +30,7 @@
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.entity.EntityServiceImpl;
 import com.linkedin.metadata.graph.elastic.ElasticSearchGraphService;
+import com.linkedin.metadata.models.AspectSpec;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.query.filter.SortOrder;
@@ -68,6 +73,7 @@ public class EntityControllerTest extends AbstractTestNGSpringContextTests {
   @Autowired private MockMvc mockMvc;
   @Autowired private SearchService mockSearchService;
   @Autowired private EntityService<?> mockEntityService;
+  @Autowired private EntityRegistry entityRegistry;
 
   @Test
   public void initTest() {
@@ -171,6 +177,57 @@ public void testSearchOrderPreserved() throws Exception {
             MockMvcResultMatchers.jsonPath("$.entities[2].urn").value(TEST_URNS.get(0).toString()));
   }
 
+  @Test
+  public void testDeleteEntity() throws Exception {
+    Urn TEST_URN = UrnUtils.getUrn("urn:li:dataset:(urn:li:dataPlatform:testPlatform,4,PROD)");
+
+    // test delete entity
+    mockMvc
+        .perform(
+            MockMvcRequestBuilders.delete(String.format("/v3/entity/dataset/%s", TEST_URN))
+                .accept(MediaType.APPLICATION_JSON))
+        .andExpect(status().is2xxSuccessful());
+
+    // test delete entity by aspect key
+    mockMvc
+        .perform(
+            MockMvcRequestBuilders.delete(String.format("/v3/entity/dataset/%s", TEST_URN))
+                .param("aspects", "datasetKey")
+                .accept(MediaType.APPLICATION_JSON))
+        .andExpect(status().is2xxSuccessful());
+
+    verify(mockEntityService, times(2)).deleteUrn(any(), eq(TEST_URN));
+
+    // test delete entity by non-key aspect
+    reset(mockEntityService);
+    mockMvc
+        .perform(
+            MockMvcRequestBuilders.delete(String.format("/v3/entity/dataset/%s", TEST_URN))
+                .param("aspects", "status")
+                .accept(MediaType.APPLICATION_JSON))
+        .andExpect(status().is2xxSuccessful());
+    verify(mockEntityService, times(1))
+        .deleteAspect(any(), eq(TEST_URN.toString()), eq("status"), anyMap(), eq(true));
+
+    // test delete entity clear
+    reset(mockEntityService);
+    mockMvc
+        .perform(
+            MockMvcRequestBuilders.delete(String.format("/v3/entity/dataset/%s", TEST_URN))
+                .param("clear", "true")
+                .accept(MediaType.APPLICATION_JSON))
+        .andExpect(status().is2xxSuccessful());
+
+    entityRegistry.getEntitySpec(DATASET_ENTITY_NAME).getAspectSpecs().stream()
+        .map(AspectSpec::getName)
+        .filter(aspectName -> !"datasetKey".equals(aspectName))
+        .forEach(
+            aspectName ->
+                verify(mockEntityService)
+                    .deleteAspect(
+                        any(), eq(TEST_URN.toString()), eq(aspectName), anyMap(), eq(true)));
+  }
+
   @TestConfiguration
   public static class EntityControllerTestConfig {
     @MockBean public EntityServiceImpl entityService;

diff --git a/...data-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java b/...data-service/services/src/main/java/com/linkedin/metadata/entity/DeleteEntityService.java
@@ -49,6 +49,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
@@ -342,9 +343,9 @@ private void deleteReference(
    */
   private void deleteAspect(
       @Nonnull OperationContext opContext, Urn urn, String aspectName, RecordTemplate prevAspect) {
-    final RollbackResult rollbackResult =
+    final Optional<RollbackResult> rollbackResult =
         _entityService.deleteAspect(opContext, urn.toString(), aspectName, new HashMap<>(), true);
-    if (rollbackResult == null || rollbackResult.getNewValue() != null) {
+    if (rollbackResult.isEmpty() || rollbackResult.get().getNewValue() != null) {
       log.error(
           "Failed to delete aspect with references. Before {}, after: null, please check GMS logs"
               + " logs for more information",

diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/entity/EntityService.java
@@ -26,6 +26,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Future;
 import java.util.function.Consumer;
@@ -458,7 +459,7 @@ void ingestEntity(
 
   void setRetentionService(RetentionService<U> retentionService);
 
-  default RollbackResult deleteAspect(
+  default Optional<RollbackResult> deleteAspect(
       @Nonnull OperationContext opContext,
       String urn,
       String aspectName,
@@ -468,7 +469,8 @@ default RollbackResult deleteAspect(
         new AspectRowSummary().setUrn(urn).setAspectName(aspectName);
     return rollbackWithConditions(opContext, List.of(aspectRowSummary), conditions, hardDelete)
         .getRollbackResults()
-        .get(0);
+        .stream()
+        .findFirst();
   }
 
   RollbackRunResult deleteUrn(@Nonnull OperationContext opContext, Urn urn);