From eca1cf7302d21f9d617df01999a72db4afa54ff5 Mon Sep 17 00:00:00 2001
From: Sachin Chaurasiya <sachinchaurasiyachotey87@gmail.com>
Date: Wed, 3 Jan 2024 19:22:21 +0530
Subject: [PATCH] FIX[#14565]: container and ml models custom property apis are
 sending wrong payload

---
 .../src/pages/ContainerPage/ContainerPage.tsx | 53 +++++++++++++------
 .../MlModelPage/MlModelPage.component.tsx     | 46 ++++++++++------
 2 files changed, 66 insertions(+), 33 deletions(-)

diff --git a/openmetadata-ui/src/main/resources/ui/src/pages/ContainerPage/ContainerPage.tsx b/openmetadata-ui/src/main/resources/ui/src/pages/ContainerPage/ContainerPage.tsx
index 81329581b3b7..16a213329f5b 100644
--- a/openmetadata-ui/src/main/resources/ui/src/pages/ContainerPage/ContainerPage.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/pages/ContainerPage/ContainerPage.tsx
@@ -245,11 +245,17 @@ const ContainerPage = () => {
     }
   };
 
-  const handleUpdateContainerData = (updatedData: Container) => {
-    const jsonPatch = compare(omitBy(containerData, isUndefined), updatedData);
+  const handleUpdateContainerData = useCallback(
+    (updatedData: Container) => {
+      const jsonPatch = compare(
+        omitBy(containerData, isUndefined),
+        updatedData
+      );
 
-    return patchContainerDetails(containerData?.id ?? '', jsonPatch);
-  };
+      return patchContainerDetails(containerData?.id ?? '', jsonPatch);
+    },
+    [containerData]
+  );
 
   const handleUpdateDescription = async (updatedDescription: string) => {
     try {
@@ -418,18 +424,33 @@ const ContainerPage = () => {
     }
   };
 
-  const handleExtensionUpdate = async (updatedContainer: Container) => {
-    try {
-      const response = await handleUpdateContainerData(updatedContainer);
-      setContainerData({
-        ...response,
-        tags: sortTagsCaseInsensitive(response.tags ?? []),
-      });
-      getEntityFeedCount();
-    } catch (error) {
-      showErrorToast(error as AxiosError);
-    }
-  };
+  const handleExtensionUpdate = useCallback(
+    async (updatedContainer: Container) => {
+      if (isUndefined(containerData)) {
+        return;
+      }
+
+      try {
+        const response = await handleUpdateContainerData({
+          ...containerData,
+          extension: updatedContainer.extension,
+        });
+        setContainerData({
+          ...response,
+          tags: sortTagsCaseInsensitive(response.tags ?? []),
+        });
+        getEntityFeedCount();
+      } catch (error) {
+        showErrorToast(error as AxiosError);
+      }
+    },
+    [
+      containerData,
+      handleUpdateContainerData,
+      getEntityFeedCount,
+      setContainerData,
+    ]
+  );
 
   const handleUpdateDataModel = async (
     updatedDataModel: Container['dataModel']
diff --git a/openmetadata-ui/src/main/resources/ui/src/pages/MlModelPage/MlModelPage.component.tsx b/openmetadata-ui/src/main/resources/ui/src/pages/MlModelPage/MlModelPage.component.tsx
index 5091f49e5b20..0ee41a951221 100644
--- a/openmetadata-ui/src/main/resources/ui/src/pages/MlModelPage/MlModelPage.component.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/pages/MlModelPage/MlModelPage.component.tsx
@@ -116,11 +116,17 @@ const MlModelPage = () => {
     }
   }, [mlModelPermissions, mlModelFqn]);
 
-  const saveUpdatedMlModelData = (updatedData: Mlmodel) => {
-    const jsonPatch = compare(omitBy(mlModelDetail, isUndefined), updatedData);
+  const saveUpdatedMlModelData = useCallback(
+    (updatedData: Mlmodel) => {
+      const jsonPatch = compare(
+        omitBy(mlModelDetail, isUndefined),
+        updatedData
+      );
 
-    return patchMlModelDetails(mlModelDetail.id, jsonPatch);
-  };
+      return patchMlModelDetails(mlModelDetail.id, jsonPatch);
+    },
+    [mlModelDetail]
+  );
 
   const descriptionUpdateHandler = async (updatedMlModel: Mlmodel) => {
     try {
@@ -230,19 +236,25 @@ const MlModelPage = () => {
     }
   };
 
-  const handleExtensionUpdate = async (updatedMlModel: Mlmodel) => {
-    try {
-      const data = await saveUpdatedMlModelData(updatedMlModel);
-      setMlModelDetail(data);
-    } catch (error) {
-      showErrorToast(
-        error as AxiosError,
-        t('server.entity-updating-error', {
-          entity: getEntityName(mlModelDetail),
-        })
-      );
-    }
-  };
+  const handleExtensionUpdate = useCallback(
+    async (updatedMlModel: Mlmodel) => {
+      try {
+        const data = await saveUpdatedMlModelData({
+          ...mlModelDetail,
+          extension: updatedMlModel.extension,
+        });
+        setMlModelDetail(data);
+      } catch (error) {
+        showErrorToast(
+          error as AxiosError,
+          t('server.entity-updating-error', {
+            entity: getEntityName(mlModelDetail),
+          })
+        );
+      }
+    },
+    [saveUpdatedMlModelData, setMlModelDetail, mlModelDetail]
+  );
 
   const createThread = async (data: CreateThread) => {
     try {