Fixes #5332 - CustomProperties of an entity are deleted during restart (#5333)

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

@@ -21,6 +21,8 @@ import java.util.concurrent.ConcurrentHashMap;
 import lombok.extern.slf4j.Slf4j;
 import org.openmetadata.catalog.entity.Type;
 import org.openmetadata.catalog.entity.type.CustomProperty;
+import org.openmetadata.catalog.exception.CatalogExceptionMessage;
+import org.openmetadata.catalog.exception.EntityNotFoundException;
 import org.openmetadata.catalog.util.FullyQualifiedName;
 import org.openmetadata.catalog.util.JsonUtils;
 
@@ -48,7 +50,6 @@ public class TypeRegistry {
 
   public void addType(Type type) {
     TYPES.put(type.getName(), type);
-    LOG.info("Updated type {}\n", type.getName());
 
     // Store custom properties added to a type
     for (CustomProperty property : type.getCustomProperties()) {
@@ -58,18 +59,17 @@ public class TypeRegistry {
 
   public void removeType(String typeName) {
     TYPES.remove(typeName);
-    LOG.info("Deleted type {}\n", typeName);
+    LOG.info("Deleted type {}", typeName);
     // TODO cleanup custom properties
   }
 
   private void addCustomProperty(String entityType, String propertyName, CustomProperty customProperty) {
     String customPropertyFQN = getCustomPropertyFQN(entityType, propertyName);
     CUSTOM_PROPERTIES.put(customPropertyFQN, customProperty);
-    LOG.info("Adding custom property {}\n", customPropertyFQN);
 
     JsonSchema jsonSchema = JsonUtils.getJsonSchema(TYPES.get(customProperty.getPropertyType().getName()).getSchema());
     CUSTOM_PROPERTY_SCHEMAS.put(customPropertyFQN, jsonSchema);
-    LOG.info("Adding json schema for {} {}\n", customProperty, jsonSchema);
+    LOG.info("Adding custom property {} with JSON schema {}", customPropertyFQN, jsonSchema);
   }
 
   public JsonSchema getSchema(String entityType, String propertyName) {
@@ -81,9 +81,8 @@ public class TypeRegistry {
     // Validate custom properties added to a type
     for (CustomProperty property : listOrEmpty(type.getCustomProperties())) {
       if (TYPES.get(property.getPropertyType().getName()) == null) {
-        throw new IllegalArgumentException(
-            String.format(
-                "Type %s not found for custom property %s", property.getPropertyType().getName(), property.getName()));
+        throw EntityNotFoundException.byMessage(
+            CatalogExceptionMessage.entityNotFound(Entity.TYPE, property.getPropertyType().getId()));
       }
     }
   }

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

@@ -87,11 +87,6 @@ public class TypeRepository extends EntityRepository<Type> {
     TypeRegistry.instance().removeType(entity.getName());
   }
 
-  @Override
-  protected void postUpdate(Type entity) {
-    TypeRegistry.instance().addType(entity);
-  }
-
   @Override
   public EntityUpdater getUpdater(Type original, Type updated, Operation operation) {
     return new TypeUpdater(original, updated, operation);
@@ -102,7 +97,7 @@ public class TypeRepository extends EntityRepository<Type> {
     Type type = dao.findEntityById(UUID.fromString(id), Include.NON_DELETED);
     property.setPropertyType(dao.findEntityReferenceById(property.getPropertyType().getId(), Include.NON_DELETED));
     if (type.getCategory().equals(Category.Field)) {
-      throw new IllegalArgumentException("Property types can't be extended");
+      throw new IllegalArgumentException("Only entity types can be extended and field types can't be extended");
     }
     setFields(type, putFields);
 
@@ -148,7 +143,7 @@ public class TypeRepository extends EntityRepository<Type> {
       List<CustomProperty> origFields = listOrEmpty(original.getCustomProperties());
       List<CustomProperty> added = new ArrayList<>();
       List<CustomProperty> deleted = new ArrayList<>();
-      recordListChange("customProperty", origFields, updatedFields, added, deleted, EntityUtil.customFieldMatch);
+      recordListChange("customProperties", origFields, updatedFields, added, deleted, EntityUtil.customFieldMatch);
       for (CustomProperty property : added) {
         String customPropertyFQN = getCustomPropertyFQN(updated.getName(), property.getName());
         String customPropertyJson = JsonUtils.pojoToJson(property);

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/types/TypeResource.java

@@ -56,6 +56,7 @@ import org.openmetadata.catalog.CatalogApplicationConfig;
 import org.openmetadata.catalog.Entity;
 import org.openmetadata.catalog.api.CreateType;
 import org.openmetadata.catalog.entity.Type;
+import org.openmetadata.catalog.entity.type.Category;
 import org.openmetadata.catalog.entity.type.CustomProperty;
 import org.openmetadata.catalog.jdbi3.CollectionDAO;
 import org.openmetadata.catalog.jdbi3.ListFilter;
@@ -66,6 +67,7 @@ import org.openmetadata.catalog.security.Authorizer;
 import org.openmetadata.catalog.security.SecurityUtil;
 import org.openmetadata.catalog.type.EntityHistory;
 import org.openmetadata.catalog.type.Include;
+import org.openmetadata.catalog.util.EntityUtil.Fields;
 import org.openmetadata.catalog.util.JsonUtils;
 import org.openmetadata.catalog.util.RestUtil.PutResponse;
 import org.openmetadata.catalog.util.ResultList;
@@ -92,7 +94,7 @@ public class TypeResource extends EntityResource<Type, TypeRepository> {
 
   @SuppressWarnings("unused") // Method used for reflection
   public void initialize(CatalogApplicationConfig config) throws IOException {
-    // Find tag definitions and load tag categories from the json file, if necessary
+    // Load types defined in OpenMetadata schemas
     long now = System.currentTimeMillis();
     List<Type> types = JsonUtils.getTypes();
     types.forEach(
@@ -100,6 +102,17 @@ public class TypeResource extends EntityResource<Type, TypeRepository> {
           type.withId(UUID.randomUUID()).withUpdatedBy("admin").withUpdatedAt(now);
           LOG.info("Loading type {}", type.getName());
           try {
+            Fields fields = getFields("customProperties");
+            try {
+              Type storedType = dao.getByName(null, type.getName(), fields);
+              type.setId(storedType.getId());
+              // If entity type already exists, then carry forward custom properties
+              if (storedType.getCategory().equals(Category.Entity)) {
+                type.setCustomProperties(storedType.getCustomProperties());
+              }
+            } catch (Exception e) {
+              LOG.debug("Creating entity that does not exist ", e);
+            }
             this.dao.createOrUpdate(null, type);
           } catch (IOException e) {
             LOG.error("Error loading type {}", type.getName(), e);

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

@@ -159,7 +159,7 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
   protected boolean supportsFieldsQueryParam = true;
   protected boolean supportsEmptyDescription = true;
   protected boolean supportsNameWithDot = true;
-  protected final boolean supportsCustomAttributes;
+  protected final boolean supportsCustomExtension;
 
   public static final String DATA_STEWARD_ROLE_NAME = "DataSteward";
   public static final String DATA_CONSUMER_ROLE_NAME = "DataConsumer";
@@ -248,7 +248,7 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
     this.supportsOwner = allowedFields.contains(FIELD_OWNER);
     this.supportsTags = allowedFields.contains(FIELD_TAGS);
     this.supportsSoftDelete = allowedFields.contains(FIELD_DELETED);
-    this.supportsCustomAttributes = allowedFields.contains(FIELD_EXTENSION);
+    this.supportsCustomExtension = allowedFields.contains(FIELD_EXTENSION);
   }
 
   @BeforeAll
@@ -1130,38 +1130,55 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
   }
 
   @Test
-  void put_entityExtension(TestInfo test) throws IOException {
-    if (!supportsCustomAttributes) {
+  void put_addEntityCustomAttributes(TestInfo test) throws IOException {
+    if (!supportsCustomExtension) {
       return;
     }
 
-    // Add valid custom fields to the entity
+    // 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);
     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();
+
+    // PATCH valid custom field stringB
     CustomProperty fieldB =
         new CustomProperty()
             .withName("stringB")
             .withDescription("stringB")
             .withPropertyType(STRING_TYPE.getEntityReference());
-    typeResourceTest.addCustomProperty(entity.getId(), fieldA, OK, ADMIN_AUTH_HEADERS);
-    typeResourceTest.addCustomProperty(entity.getId(), fieldB, OK, ADMIN_AUTH_HEADERS);
+    String json = JsonUtils.pojoToJson(entity);
 
-    // Add invalid custom fields to the entity - custom field has invalid type
-    Type INVALID_TYPE =
+    ChangeDescription change = getChangeDescription(entity.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);
+
+    // PUT invalid custom fields to the entity - custom field has invalid type
+    Type invalidType =
         new Type().withId(UUID.randomUUID()).withName("invalid").withCategory(Category.Field).withSchema("{}");
     CustomProperty fieldInvalid =
         new CustomProperty()
             .withName("invalid")
             .withDescription("invalid")
-            .withPropertyType(INVALID_TYPE.getEntityReference());
+            .withPropertyType(invalidType.getEntityReference());
     assertResponse(
-        () -> typeResourceTest.addCustomProperty(entity.getId(), fieldInvalid, OK, ADMIN_AUTH_HEADERS),
+        () -> typeResourceTest.addCustomProperty(id, fieldInvalid, OK, ADMIN_AUTH_HEADERS),
         NOT_FOUND,
-        CatalogExceptionMessage.entityNotFound(Entity.TYPE, INVALID_TYPE.getId()));
+        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;
+    assertResponse(
+        () -> typeResourceTest.patchEntity(id, json1, finalEntity, ADMIN_AUTH_HEADERS),
+        NOT_FOUND,
+        CatalogExceptionMessage.entityNotFound(Entity.TYPE, invalidType.getId()));
 
     // Now create an entity with custom field
     ObjectMapper mapper = new ObjectMapper();

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

@@ -22,6 +22,7 @@ import static org.openmetadata.common.utils.CommonUtil.listOrEmpty;
 
 import java.io.IOException;
 import java.net.URISyntaxException;
+import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 import javax.ws.rs.client.WebTarget;
@@ -42,6 +43,8 @@ import org.openmetadata.catalog.resources.EntityResourceTest;
 import org.openmetadata.catalog.resources.types.TypeResource;
 import org.openmetadata.catalog.resources.types.TypeResource.TypeList;
 import org.openmetadata.catalog.type.EntityReference;
+import org.openmetadata.catalog.util.EntityUtil;
+import org.openmetadata.catalog.util.JsonUtils;
 import org.openmetadata.catalog.util.TestUtils;
 
 @Slf4j
@@ -112,7 +115,7 @@ public class TypeResourceTest extends EntityResourceTest<Type, CreateType> {
     assertResponse(
         () -> addCustomProperty(INT_TYPE.getId(), field, Status.CREATED, ADMIN_AUTH_HEADERS),
         Status.BAD_REQUEST,
-        "Property types can't be extended");
+        "Only entity types can be extended and field types can't be extended");
   }
 
   @Override
@@ -125,6 +128,17 @@ public class TypeResourceTest extends EntityResourceTest<Type, CreateType> {
     return type;
   }
 
+  public Type addAndCheckCustomProperty(
+      UUID entityTypeId, CustomProperty customProperty, Status status, Map<String, String> authHeaders)
+      throws HttpResponseException {
+    Type updated = addCustomProperty(entityTypeId, customProperty, status, authHeaders);
+    assertTrue(updated.getCustomProperties().contains(customProperty));
+
+    Type get = getEntity(entityTypeId, "customProperties", authHeaders);
+    assertTrue(get.getCustomProperties().contains(customProperty));
+    return updated;
+  }
+
   public Type addCustomProperty(
       UUID entityTypeId, CustomProperty customProperty, Status status, Map<String, String> authHeaders)
       throws HttpResponseException {
@@ -159,6 +173,15 @@ public class TypeResourceTest extends EntityResourceTest<Type, CreateType> {
     if (expected == actual) {
       return;
     }
-    assertCommonFieldChange(fieldName, expected, actual);
+    if (fieldName.equals("customProperties")) {
+      List<CustomProperty> expectedProperties = (List<CustomProperty>) expected;
+      List<CustomProperty> actualProperties = JsonUtils.readObjects(actual.toString(), CustomProperty.class);
+      assertEquals(expectedProperties.size(), actualProperties.size());
+      expectedProperties.sort(EntityUtil.compareCustomProperty);
+      actualProperties.sort(EntityUtil.compareCustomProperty);
+      assertEquals(expectedProperties, actualProperties);
+    } else {
+      assertCommonFieldChange(fieldName, expected, actual);
+    }
   }
 }