Fixes #5345 - PUT and PATCH API to support storing custom properties added through entity extension (#5346)

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

@@ -306,9 +306,13 @@ public interface CollectionDAO {
     @RegisterRowMapper(ExtensionMapper.class)
     @SqlQuery(
         "SELECT extension, json FROM entity_extension WHERE id = :id AND extension "
-            + "LIKE CONCAT (:extensionPrefix, '.%')")
+            + "LIKE CONCAT (:extensionPrefix, '.%') "
+            + "ORDER BY extension")
     List<ExtensionRecord> getExtensions(@Bind("id") String id, @Bind("extensionPrefix") String extensionPrefix);
 
+    @SqlUpdate("DELETE FROM entity_extension WHERE id = :id AND extension = :extension")
+    void delete(@Bind("id") String id, @Bind("extension") String extension);
+
     @SqlUpdate("DELETE FROM entity_extension WHERE id = :id")
     void deleteAll(@Bind("id") String id);
   }

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

@@ -572,7 +572,7 @@ public abstract class EntityRepository<T extends EntityInterface> {
     }
   }
 
-  protected void cleanup(EntityInterface entityInterface) {
+  protected void cleanup(EntityInterface entityInterface) throws JsonProcessingException {
     String id = entityInterface.getId().toString();
 
     // Delete all the relationships to other entities
@@ -590,6 +590,9 @@ public abstract class EntityRepository<T extends EntityInterface> {
     // Delete all the usage data
     daoCollection.usageDAO().delete(id);
 
+    // Delete the extension data storing custom properties
+    removeExtension(entityInterface);
+
     // Finally, delete the entity
     dao.delete(id);
   }
@@ -671,24 +674,39 @@ public abstract class EntityRepository<T extends EntityInterface> {
     }
   }
 
-  public void storeExtension(T entity) throws JsonProcessingException {
+  public void storeExtension(EntityInterface entity) throws JsonProcessingException {
     JsonNode jsonNode = JsonUtils.valueToTree(entity.getExtension());
     Iterator<Entry<String, JsonNode>> customFields = jsonNode.fields();
     while (customFields.hasNext()) {
       Entry<String, JsonNode> entry = customFields.next();
       String fieldName = entry.getKey();
       JsonNode value = entry.getValue();
-      storeCustomField(entity, fieldName, value);
+      storeCustomProperty(entity, fieldName, value);
     }
   }
 
-  private void storeCustomField(T entity, String fieldName, JsonNode value) throws JsonProcessingException {
+  public void removeExtension(EntityInterface entity) throws JsonProcessingException {
+    JsonNode jsonNode = JsonUtils.valueToTree(entity.getExtension());
+    Iterator<Entry<String, JsonNode>> customFields = jsonNode.fields();
+    while (customFields.hasNext()) {
+      Entry<String, JsonNode> entry = customFields.next();
+      removeCustomProperty(entity, entry.getKey());
+    }
+  }
+
+  private void storeCustomProperty(EntityInterface entity, String fieldName, JsonNode value)
+      throws JsonProcessingException {
     String fieldFQN = TypeRegistry.getCustomPropertyFQN(entityType, fieldName);
     daoCollection
         .entityExtensionDAO()
         .insert(entity.getId().toString(), fieldFQN, "customFieldSchema", JsonUtils.pojoToJson(value));
   }
 
+  private void removeCustomProperty(EntityInterface entity, String fieldName) {
+    String fieldFQN = TypeRegistry.getCustomPropertyFQN(entityType, fieldName);
+    daoCollection.entityExtensionDAO().delete(entity.getId().toString(), fieldFQN);
+  }
+
   public ObjectNode getExtension(T entity) throws JsonProcessingException {
     String fieldFQNPrefix = TypeRegistry.getCustomPropertyFQNPrefix(entityType);
     List<ExtensionRecord> records =
@@ -1046,6 +1064,7 @@ public abstract class EntityRepository<T extends EntityInterface> {
         updateDescription();
         updateDisplayName();
         updateOwner();
+        updateExtension();
         updateTags(updated.getFullyQualifiedName(), FIELD_TAGS, original.getTags(), updated.getTags());
         entitySpecificUpdate();
       }
@@ -1135,6 +1154,12 @@ public abstract class EntityRepository<T extends EntityInterface> {
       applyTags(updatedTags, fqn);
     }
 
+    private void updateExtension() throws JsonProcessingException {
+      removeExtension(original);
+      storeExtension(updated);
+      // TODO change descriptions for custom attributes
+    }
+
     public final boolean updateVersion(Double oldVersion) {
       Double newVersion = oldVersion;
       if (majorVersionChange) {

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

@@ -90,9 +90,10 @@ import org.openmetadata.common.utils.CommonUtil;
 public class TableRepository extends EntityRepository<Table> {
 
   // Table fields that can be patched in a PATCH request
-  static final String TABLE_PATCH_FIELDS = "owner,tags,tableConstraints,tablePartition";
+  static final String TABLE_PATCH_FIELDS = "owner,tags,tableConstraints,tablePartition,extension";
   // Table fields that can be updated in a PUT request
-  static final String TABLE_UPDATE_FIELDS = "owner,tags,tableConstraints,tablePartition,dataModel,profileSample";
+  static final String TABLE_UPDATE_FIELDS =
+      "owner,tags,tableConstraints,tablePartition,dataModel,profileSample," + "extension";
 
   public static final String FIELD_RELATION_COLUMN_TYPE = "table.columns.column";
   public static final String FIELD_RELATION_TABLE_TYPE = "table";

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

@@ -1138,25 +1138,25 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
     // PUT valid custom field intA to the entity type
     TypeResourceTest typeResourceTest = new TypeResourceTest();
     INT_TYPE = typeResourceTest.getEntityByName("integer", "", ADMIN_AUTH_HEADERS);
-    STRING_TYPE = typeResourceTest.getEntityByName("string", "", ADMIN_AUTH_HEADERS);
-    Type entity = typeResourceTest.getEntityByName(this.entityType, "customProperties", ADMIN_AUTH_HEADERS);
+    Type entityType = typeResourceTest.getEntityByName(this.entityType, "customProperties", ADMIN_AUTH_HEADERS);
     CustomProperty fieldA =
         new CustomProperty().withName("intA").withDescription("intA").withPropertyType(INT_TYPE.getEntityReference());
-    entity = typeResourceTest.addAndCheckCustomProperty(entity.getId(), fieldA, OK, ADMIN_AUTH_HEADERS);
-    final UUID id = entity.getId();
+    entityType = typeResourceTest.addAndCheckCustomProperty(entityType.getId(), fieldA, OK, ADMIN_AUTH_HEADERS);
+    final UUID id = entityType.getId();
 
     // PATCH valid custom field stringB
+    STRING_TYPE = typeResourceTest.getEntityByName("string", "", ADMIN_AUTH_HEADERS);
     CustomProperty fieldB =
         new CustomProperty()
             .withName("stringB")
             .withDescription("stringB")
             .withPropertyType(STRING_TYPE.getEntityReference());
-    String json = JsonUtils.pojoToJson(entity);
 
-    ChangeDescription change = getChangeDescription(entity.getVersion());
+    String json = JsonUtils.pojoToJson(entityType);
+    ChangeDescription change = getChangeDescription(entityType.getVersion());
     change.getFieldsAdded().add(new FieldChange().withName("customProperties").withNewValue(Arrays.asList(fieldB)));
-    entity.getCustomProperties().add(fieldB);
-    entity = typeResourceTest.patchEntityAndCheck(entity, json, ADMIN_AUTH_HEADERS, MINOR_UPDATE, change);
+    entityType.getCustomProperties().add(fieldB);
+    entityType = typeResourceTest.patchEntityAndCheck(entityType, json, ADMIN_AUTH_HEADERS, MINOR_UPDATE, change);
 
     // PUT invalid custom fields to the entity - custom field has invalid type
     Type invalidType =
@@ -1172,35 +1172,65 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
         CatalogExceptionMessage.entityNotFound(Entity.TYPE, invalidType.getId()));
 
     // PATCH invalid custom fields to the entity - custom field has invalid type
-    String json1 = JsonUtils.pojoToJson(entity);
-    entity.getCustomProperties().add(fieldInvalid);
-    Type finalEntity = entity;
+    String json1 = JsonUtils.pojoToJson(entityType);
+    entityType.getCustomProperties().add(fieldInvalid);
+    Type finalEntity = entityType;
     assertResponse(
         () -> typeResourceTest.patchEntity(id, json1, finalEntity, ADMIN_AUTH_HEADERS),
         NOT_FOUND,
         CatalogExceptionMessage.entityNotFound(Entity.TYPE, invalidType.getId()));
 
-    // Now create an entity with custom field
+    // Now POST an entity with extension that includes custom field intA
     ObjectMapper mapper = new ObjectMapper();
     ObjectNode jsonNode = mapper.createObjectNode();
     jsonNode.set("intA", mapper.convertValue(1, JsonNode.class));
-    jsonNode.set("stringB", mapper.convertValue("string", JsonNode.class));
     K create = createRequest(test).withExtension(jsonNode);
-    createAndCheckEntity(create, ADMIN_AUTH_HEADERS);
+    T entity = createAndCheckEntity(create, ADMIN_AUTH_HEADERS);
 
-    // Now set the entity custom field with an unknown field name
-    jsonNode.set("stringC", mapper.convertValue("string", JsonNode.class)); // Unknown field
-    assertResponse(
-        () -> createEntity(createRequest(test, 1).withExtension(jsonNode), ADMIN_AUTH_HEADERS),
-        BAD_REQUEST,
-        CatalogExceptionMessage.unknownCustomField("stringC"));
+    // PUT and update the entity with extension field intA to a new value
+    // TODO to do change description for stored customProperties
+    jsonNode.set("intA", mapper.convertValue(2, JsonNode.class));
+    create = createRequest(test).withExtension(jsonNode);
+    entity = updateEntity(create, Status.OK, ADMIN_AUTH_HEADERS);
+    assertEquals(JsonUtils.valueToTree(create.getExtension()), JsonUtils.valueToTree(entity.getExtension()));
 
-    // Now set the entity custom field to an invalid value
+    // PATCH and update the entity with extension field stringB
+    // TODO to do change description for stored customProperties
+    json = JsonUtils.pojoToJson(entity);
+    jsonNode.set("stringB", mapper.convertValue("stringB", JsonNode.class));
+    entity.setExtension(jsonNode);
+    entity = patchEntity(entity.getId(), json, entity, ADMIN_AUTH_HEADERS);
+    assertEquals(JsonUtils.valueToTree(jsonNode), JsonUtils.valueToTree(entity.getExtension()));
+
+    // PUT and remove field intA from the the entity extension
+    // TODO to do change description for stored customProperties
+    jsonNode.remove("intA");
+    create = createRequest(test).withExtension(jsonNode);
+    entity = updateEntity(create, Status.OK, ADMIN_AUTH_HEADERS);
+    assertEquals(JsonUtils.valueToTree(create.getExtension()), JsonUtils.valueToTree(entity.getExtension()));
+
+    // PATCH and remove field stringB from the the entity extension
+    // TODO to do change description for stored customProperties
+    json = JsonUtils.pojoToJson(entity);
+    jsonNode.remove("stringB");
+    entity.setExtension(jsonNode);
+    entity = patchEntity(entity.getId(), json, entity, ADMIN_AUTH_HEADERS);
+    assertEquals(JsonUtils.valueToTree(jsonNode), JsonUtils.valueToTree(entity.getExtension()));
+
+    // Now set the entity custom property to an invalid value
     jsonNode.set("intA", mapper.convertValue("stringInsteadOfNumber", JsonNode.class)); // String in integer field
     assertResponseContains(
         () -> createEntity(createRequest(test, 1).withExtension(jsonNode), ADMIN_AUTH_HEADERS),
         BAD_REQUEST,
         CatalogExceptionMessage.jsonValidationError("intA", ""));
+
+    // Now set the entity custom property with an unknown field name
+    jsonNode.remove("intA");
+    jsonNode.set("stringC", mapper.convertValue("string", JsonNode.class)); // Unknown field
+    assertResponse(
+        () -> createEntity(createRequest(test, 1).withExtension(jsonNode), ADMIN_AUTH_HEADERS),
+        BAD_REQUEST,
+        CatalogExceptionMessage.unknownCustomField("stringC"));
   }
 
   ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////