Fixes #3021 - Restore a soft-deleted container entity along with the …

…children

catalog-rest-service/src/main/java/org/openmetadata/catalog/Entity.java

@@ -139,6 +139,15 @@ public static EntityReference getEntityReferenceByName(@NonNull String entity, @
     return dao.findEntityReferenceByName(fqn);
   }
 
+  public static EntityReference getEntityReferenceByName(@NonNull String entity, @NonNull String fqn, Include include)
+      throws IOException {
+    EntityDAO<?> dao = DAO_MAP.get(entity);
+    if (dao == null) {
+      throw EntityNotFoundException.byMessage(CatalogExceptionMessage.entityTypeNotFound(entity));
+    }
+    return dao.findEntityReferenceByName(fqn, include);
+  }
+
   public static <T> EntityReference getEntityReference(T entity) {
     String entityType = getEntityTypeFromObject(entity);
 
@@ -224,15 +233,24 @@ public static <T> EntityRepository<T> getEntityRepository(@NonNull String entity
     return entityRepository;
   }
 
-  public static void deleteEntity(String updatedBy, String entity, UUID entityId, boolean recursive)
+  public static void deleteEntity(String updatedBy, String entityType, UUID entityId, boolean recursive)
       throws IOException, ParseException {
-    EntityRepository<?> dao = ENTITY_REPOSITORY_MAP.get(entity);
+    EntityRepository<?> dao = ENTITY_REPOSITORY_MAP.get(entityType);
     if (dao == null) {
-      throw EntityNotFoundException.byMessage(CatalogExceptionMessage.entityTypeNotFound(entity));
+      throw EntityNotFoundException.byMessage(CatalogExceptionMessage.entityTypeNotFound(entityType));
     }
     dao.delete(updatedBy, entityId.toString(), recursive);
   }
 
+  public static void restoreEntity(String updatedBy, String entityType, UUID entityId)
+      throws IOException, ParseException {
+    EntityRepository<?> dao = ENTITY_REPOSITORY_MAP.get(entityType);
+    if (dao == null) {
+      throw EntityNotFoundException.byMessage(CatalogExceptionMessage.entityTypeNotFound(entityType));
+    }
+    dao.restoreEntity(updatedBy, entityType, entityId);
+  }
+
   public static <T> EntityRepository<T> getEntityRepositoryForClass(@NonNull Class<T> clazz) {
     @SuppressWarnings("unchecked")
     EntityRepository<T> entityRepository = (EntityRepository<T>) CLASS_ENTITY_REPOSITORY_MAP.get(clazz);

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/CollectionDAO.java

@@ -502,7 +502,7 @@ int deleteTo(
     void softDeleteAll(@Bind("id") String id, @Bind("entity") String entity);
 
     @SqlUpdate("UPDATE entity_relationship SET deleted = false WHERE toId = :id OR fromId = :id")
-    void recoverSoftDeleteAll(@Bind("id") String id);
+    int recoverSoftDeleteAll(@Bind("id") String id);
   }
 
   interface FeedDAO {

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/EntityDAO.java

@@ -195,6 +195,14 @@ default EntityReference findEntityReferenceByName(String fqn) throws IOException
     return getEntityReference(findEntityByName(fqn));
   }
 
+  default EntityReference findEntityReferenceById(UUID id, Include include) throws IOException {
+    return getEntityReference(findEntityById(id, include));
+  }
+
+  default EntityReference findEntityReferenceByName(String fqn, Include include) throws IOException {
+    return getEntityReference(findEntityByName(fqn, include));
+  }
+
   default String findJsonById(String id, Include include) {
     return findById(getTableName(), id, toBoolean(include));
   }

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/EntityRepository.java

@@ -374,16 +374,21 @@ public final PutResponse<T> createOrUpdate(UriInfo uriInfo, T updated) throws IO
   public final PutResponse<T> createOrUpdate(UriInfo uriInfo, T updated, boolean allowEdits)
       throws IOException, ParseException {
     prepare(updated);
+    EntityInterface<T> updatedInterface = getEntityInterface(updated);
+
     // Check if there is any original, deleted or not
     T original = JsonUtils.readValue(dao.findJsonByFqn(getFullyQualifiedName(updated), Include.ALL), entityClass);
     if (original == null) {
       return new PutResponse<>(Status.CREATED, withHref(uriInfo, createNewEntity(updated)), RestUtil.ENTITY_CREATED);
     }
-
     // Get all the fields in the original entity that can be updated during PUT operation
     setFields(original, putFields);
-    // Recover relationships if original was deleted before setFields
-    recoverDeletedRelationships(original);
+
+    // If the entity state is soft-deleted, recursively undelete the entity and it's children
+    EntityInterface<T> origInterface = getEntityInterface(original);
+    if (origInterface.isDeleted()) {
+      restoreEntity(updatedInterface.getUpdatedBy(), entityType, origInterface.getId());
+    }
 
     // Update the attributes and relationships of an entity
     EntityUpdater entityUpdater = getUpdater(original, updated, Operation.PUT);
@@ -622,13 +627,29 @@ public URI getHref(UriInfo uriInfo, UUID id) {
     return RestUtil.getHref(uriInfo, collectionPath, id);
   }
 
-  private void recoverDeletedRelationships(T original) {
-    // If original is deleted, we need to recover the relationships before setting the fields
-    // or we won't find the related services
-    EntityInterface<T> originalRef = getEntityInterface(original);
-    if (Boolean.TRUE.equals(originalRef.isDeleted())) {
-      daoCollection.relationshipDAO().recoverSoftDeleteAll(originalRef.getId().toString());
+  public void restoreEntity(String updatedBy, String entityType, UUID id) throws IOException, ParseException {
+    // If an entity being restored contains other **deleted** children entities, restore them
+    List<EntityReference> contains =
+        daoCollection
+            .relationshipDAO()
+            .findTo(id.toString(), entityType, Relationship.CONTAINS.ordinal(), toBoolean(Include.DELETED));
+
+    if (!contains.isEmpty()) {
+      // Restore all the contained entities
+      for (EntityReference entityReference : contains) {
+        LOG.info("Recursively restoring {} {}", entityReference.getType(), entityReference.getId());
+        Entity.restoreEntity(updatedBy, entityReference.getType(), entityReference.getId());
+      }
     }
+
+    // Restore all the relationships from and to the entity as not deleted
+    daoCollection.relationshipDAO().recoverSoftDeleteAll(id.toString());
+
+    // Finally set entity deleted flag to false
+    T entity = dao.findEntityById(id, DELETED);
+    EntityInterface<T> entityInterface = getEntityInterface(entity);
+    entityInterface.setDeleted(false);
+    dao.update(entityInterface.getId(), JsonUtils.pojoToJson(entity));
   }
 
   /** Builder method for EntityHandler */
@@ -758,7 +779,7 @@ public EntityReference getContainer(String containerEntityType) throws IOExcepti
                   Relationship.CONTAINS.ordinal(),
                   // FIXME: containerEntityName should be a property of the entity decorated.
                   containerEntityType,
-                  toBoolean(isDeleted));
+                  null);
       if (refs.isEmpty()) {
         throw new UnhandledServerException(CatalogExceptionMessage.entityTypeNotFound(containerEntityType));
       } else if (refs.size() > 1) {
@@ -900,6 +921,7 @@ public class EntityUpdater {
     protected final Operation operation;
     protected final ChangeDescription changeDescription = new ChangeDescription();
     protected boolean majorVersionChange = false;
+    private boolean entityRestored = false;
 
     public EntityUpdater(T original, T updated, Operation operation) {
       this.original = getEntityInterface(original);
@@ -955,6 +977,7 @@ private void updateDeleted() throws JsonProcessingException {
         if (Boolean.TRUE.equals(original.isDeleted())) {
           updated.setDeleted(false);
           recordChange("deleted", true, false);
+          entityRestored = true;
         }
       } else {
         recordChange("deleted", original.isDeleted(), updated.isDeleted());
@@ -1035,6 +1058,10 @@ public boolean fieldsChanged() {
           || !changeDescription.getFieldsDeleted().isEmpty();
     }
 
+    public boolean isEntityRestored() {
+      return entityRestored;
+    }
+
     public final <K> boolean recordChange(String field, K orig, K updated) throws JsonProcessingException {
       return recordChange(field, orig, updated, false, objectMatch);
     }

catalog-rest-service/src/main/java/org/openmetadata/catalog/util/EntityUtil.java

@@ -256,7 +256,7 @@ public static EntityReference validateEntityLink(EntityLink entityLink) throws I
 
     // TODO: add more validation for field name and array fields
 
-    return Entity.getEntityReferenceByName(entityType, fqn);
+    return Entity.getEntityReferenceByName(entityType, fqn, ALL);
   }
 
   public static UsageDetails getLatestUsage(UsageDAO usageDAO, UUID entityId) {

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/EntityResourceTest.java

@@ -424,7 +424,7 @@ public T beforeDeletion(TestInfo test, T entity) throws HttpResponseException {
   // Get container entity based on create request that has CONTAINS relationship to the entity created with this
   // request has . For table, it is database. For database, it is databaseService. See Relationship.CONTAINS for
   // details.
-  public EntityReference getContainer(K createRequest) {
+  public EntityReference getContainer() {
     return null;
   }
 
@@ -611,11 +611,16 @@ void get_entityListWithPagination_200(TestInfo test) throws IOException {
   /** At the end of test for an entity, delete the parent container to test recursive delete functionality */
   private void delete_recursiveTest() throws IOException {
     // Finally, delete the container that contains the entities created for this test
-    EntityReference container = getContainer(createRequest("deleteRecursive", "", "", null));
+    EntityReference container = getContainer();
     if (container != null) {
-      ResultList<T> listBeforeDeletion = listEntities(null, 1000, null, null, ADMIN_AUTH_HEADERS);
+      // List both deleted and non deleted entities
+      Map<String, String> queryParams = new HashMap<>();
+      queryParams.put("include", "all");
+      ResultList<T> listBeforeDeletion = listEntities(queryParams, 1000, null, null, ADMIN_AUTH_HEADERS);
+
       // Delete non-empty container entity and ensure deletion is not allowed
-      EntityResourceTest<?, ?> containerTest = ENTITY_RESOURCE_TEST_MAP.get(container.getType());
+      EntityResourceTest<Object, Object> containerTest =
+          (EntityResourceTest<Object, Object>) ENTITY_RESOURCE_TEST_MAP.get(container.getType());
       assertResponse(
           () -> containerTest.deleteEntity(container.getId(), ADMIN_AUTH_HEADERS),
           BAD_REQUEST,
@@ -624,12 +629,26 @@ private void delete_recursiveTest() throws IOException {
       // Now delete the container with recursive flag on
       containerTest.deleteEntity(container.getId(), true, ADMIN_AUTH_HEADERS);
 
-      // Make sure entities contained are deleted and the new list operation returns 0 entities
+      // Make sure entities that belonged to the container are deleted and the new list operation returns less entities
       ResultList<T> listAfterDeletion = listEntities(null, 1000, null, null, ADMIN_AUTH_HEADERS);
       listAfterDeletion
           .getData()
           .forEach(e -> assertNotEquals(getEntityInterface(e).getContainer().getId(), container.getId()));
       assertTrue(listAfterDeletion.getData().size() < listBeforeDeletion.getData().size());
+
+      // Restore the soft-deleted container by PUT operation and make sure it is restored
+      String containerName = container.getName();
+      if (containerTest.getContainer() != null) {
+        // Find container name by removing parentContainer fqn from container fqn
+        // Example: remove "service" from "service.database" to get "database" container name for table
+        String parentOfContainer = containerTest.getContainer().getName();
+        containerName = container.getName().replace(parentOfContainer + ".", "");
+      }
+      Object request = containerTest.createRequest(containerName, "", "", null);
+      containerTest.updateEntity(request, Status.OK, ADMIN_AUTH_HEADERS);
+
+      ResultList<T> listAfterRestore = listEntities(null, 1000, null, null, ADMIN_AUTH_HEADERS);
+      assertEquals(listBeforeDeletion.getData().size(), listAfterRestore.getData().size());
     }
   }
 

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/charts/ChartResourceTest.java

@@ -130,13 +130,13 @@ public CreateChart createRequest(String name, String description, String display
         .withDescription(description)
         .withDisplayName(displayName)
         .withOwner(owner)
-        .withService(SUPERSET_REFERENCE)
+        .withService(getContainer())
         .withChartType(ChartType.Area);
   }
 
   @Override
-  public EntityReference getContainer(CreateChart createRequest) {
-    return createRequest.getService();
+  public EntityReference getContainer() {
+    return SUPERSET_REFERENCE;
   }
 
   @Override

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/dashboards/DashboardResourceTest.java

@@ -234,16 +234,16 @@ private static void validateDashboardCharts(Dashboard dashboard, List<EntityRefe
   public CreateDashboard createRequest(String name, String description, String displayName, EntityReference owner) {
     return new CreateDashboard()
         .withName(name)
-        .withService(SUPERSET_REFERENCE)
+        .withService(getContainer())
         .withCharts(CHART_REFERENCES)
         .withDescription(description)
         .withDisplayName(displayName)
         .withOwner(owner);
   }
 
   @Override
-  public EntityReference getContainer(CreateDashboard createRequest) {
-    return createRequest.getService();
+  public EntityReference getContainer() {
+    return SUPERSET_REFERENCE;
   }
 
   @Override

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/databases/DatabaseResourceTest.java

@@ -161,12 +161,12 @@ public CreateDatabase createRequest(String name, String description, String disp
         .withName(name)
         .withDescription(description)
         .withOwner(owner)
-        .withService(SNOWFLAKE_REFERENCE);
+        .withService(getContainer());
   }
 
   @Override
-  public EntityReference getContainer(CreateDatabase createRequest) {
-    return createRequest.getService();
+  public EntityReference getContainer() {
+    return SNOWFLAKE_REFERENCE;
   }
 
   @Override

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/databases/TableResourceTest.java

@@ -1791,7 +1791,7 @@ public CreateTable createRequest(String name, String description, String display
         new TableConstraint().withConstraintType(ConstraintType.UNIQUE).withColumns(List.of(COLUMNS.get(0).getName()));
     return new CreateTable()
         .withName(name)
-        .withDatabase(DATABASE_REFERENCE)
+        .withDatabase(getContainer())
         .withColumns(COLUMNS)
         .withTableConstraints(List.of(constraint))
         .withDescription(description)
@@ -1809,8 +1809,8 @@ public Table beforeDeletion(TestInfo test, Table table) throws HttpResponseExcep
   }
 
   @Override
-  public EntityReference getContainer(CreateTable createRequest) {
-    return Entity.getEntityReference(DATABASE); // TODO clean this up
+  public EntityReference getContainer() {
+    return DATABASE_REFERENCE;
   }
 
   @Override

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/glossary/GlossaryResourceTest.java

@@ -66,11 +66,6 @@ public CreateGlossary createRequest(String name, String description, String disp
         .withOwner(owner);
   }
 
-  @Override
-  public EntityReference getContainer(CreateGlossary createRequest) {
-    return null;
-  }
-
   @Override
   public void validateCreatedEntity(
       Glossary createdEntity, CreateGlossary createRequest, Map<String, String> authHeaders)

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/glossary/GlossaryTermResourceTest.java

@@ -214,11 +214,6 @@ public CreateGlossaryTerm createRequest(String name, String description, String
         .withReviewers(List.of(USER_OWNER1));
   }
 
-  @Override
-  public EntityReference getContainer(CreateGlossaryTerm createRequest) {
-    return null;
-  }
-
   @Override
   public void validateCreatedEntity(GlossaryTerm entity, CreateGlossaryTerm request, Map<String, String> authHeaders)
       throws HttpResponseException {

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/locations/LocationResourceTest.java

@@ -72,14 +72,14 @@ public void setup(TestInfo test) throws IOException, URISyntaxException {
   public CreateLocation createRequest(String name, String description, String displayName, EntityReference owner) {
     return new CreateLocation()
         .withName(name)
-        .withService(AWS_STORAGE_SERVICE_REFERENCE)
+        .withService(getContainer())
         .withDescription(description)
         .withOwner(owner);
   }
 
   @Override
-  public EntityReference getContainer(CreateLocation createRequest) {
-    return createRequest.getService();
+  public EntityReference getContainer() {
+    return AWS_STORAGE_SERVICE_REFERENCE;
   }
 
   @Override

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/operations/AirflowPipelineResourceTest.java

@@ -129,7 +129,7 @@ public CreateAirflowPipeline createRequest(
     return new CreateAirflowPipeline()
         .withName(name)
         .withPipelineType(PipelineType.METADATA)
-        .withService(BIGQUERY_REFERENCE)
+        .withService(getContainer())
         .withPipelineConfig(INGESTION_CONFIG)
         .withStartDate("2021-11-21")
         .withDescription(description)
@@ -138,8 +138,8 @@ public CreateAirflowPipeline createRequest(
   }
 
   @Override
-  public EntityReference getContainer(CreateAirflowPipeline createRequest) {
-    return createRequest.getService();
+  public EntityReference getContainer() {
+    return BIGQUERY_REFERENCE;
   }
 
   @Override

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/pipelines/PipelineResourceTest.java

@@ -101,16 +101,16 @@ public void setup(TestInfo test) throws IOException, URISyntaxException {
   public CreatePipeline createRequest(String name, String description, String displayName, EntityReference owner) {
     return new CreatePipeline()
         .withName(name)
-        .withService(AIRFLOW_REFERENCE)
+        .withService(getContainer())
         .withDescription(description)
         .withDisplayName(displayName)
         .withOwner(owner)
         .withTasks(TASKS);
   }
 
   @Override
-  public EntityReference getContainer(CreatePipeline createRequest) {
-    return createRequest.getService();
+  public EntityReference getContainer() {
+    return AIRFLOW_REFERENCE;
   }
 
   @Override