diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java
index 002ef94103b..9c6cada3a06 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java
@@ -2840,6 +2840,13 @@ public abstract class EntityRepository<T extends EntityInterface> {
         return;
       }
 
+      if (operation == Operation.PUT && updatedExtension == null) {
+        // Revert change to non-empty extension if it is being updated by a PUT request
+        // For PUT operations, existing extension can't be removed.
+        updated.setExtension(origExtension);
+        return;
+      }
+
       List<JsonNode> added = new ArrayList<>();
       List<JsonNode> deleted = new ArrayList<>();
       JsonNode origFields = JsonUtils.valueToTree(origExtension);
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/MlModelRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/MlModelRepository.java
index 2937030e0e6..e9ada22bd09 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/MlModelRepository.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/MlModelRepository.java
@@ -211,7 +211,8 @@ public class MlModelRepository extends EntityRepository<MlModel> {
                       .getFeatureSources()
                       .forEach(
                           mlFeatureSource -> {
-                            EntityReference targetEntity = getEntityReference(mlFeatureSource.getDataSource(), Include.ALL);
+                            EntityReference targetEntity =
+                                getEntityReference(mlFeatureSource.getDataSource(), Include.ALL);
                             if (targetEntity != null) {
                               addRelationship(
                                   targetEntity.getId(),