From caf4fea13a871a3d03687e3f3f021c6361b1404b Mon Sep 17 00:00:00 2001
From: John Joyce <john@acryl.io>
Date: Thu, 4 Aug 2022 09:43:44 -0700
Subject: [PATCH] Addressing comments

---
 .../mutate/BatchSetDomainResolver.java        |  2 +-
 .../resolvers/mutate/util/LabelUtils.java     |  1 +
 .../styled/search/action/DomainsDropdown.tsx  | 31 +++++++++--------
 .../profile/sidebar/Domain/SetDomainModal.tsx | 34 +++++++++----------
 4 files changed, 35 insertions(+), 33 deletions(-)

diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchSetDomainResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchSetDomainResolver.java
index 36de016c63396..9b6167c673d8d 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchSetDomainResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/BatchSetDomainResolver.java
@@ -65,7 +65,7 @@ private void validateInputResources(List<ResourceRefInput> resources, QueryConte
 
   private void validateInputResource(ResourceRefInput resource, QueryContext context) {
     final Urn resourceUrn = UrnUtils.getUrn(resource.getResourceUrn());
-    if (!LabelUtils.isAuthorizedToUpdateTerms(context, resourceUrn, resource.getSubResource())) {
+    if (!DomainUtils.isAuthorizedToUpdateDomainsForEntity(context, resourceUrn)) {
       throw new AuthorizationException("Unauthorized to perform this action. Please contact your DataHub administrator.");
     }
     LabelUtils.validateResource(resourceUrn, resource.getSubResource(), resource.getSubResourceType(), _entityService);
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java
index 07ea2edc9ab6f..a19a9c6e9a9c0 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/LabelUtils.java
@@ -249,6 +249,7 @@ public static void validateLabel(Urn labelUrn, String labelEntityType, EntitySer
     }
   }
 
+  // TODO: Move this out into a separate utilities class.
   public static void validateResource(Urn resourceUrn, String subResource, SubResourceType subResourceType, EntityService entityService) {
     if (!entityService.exists(resourceUrn)) {
       throw new IllegalArgumentException(String.format("Failed to update resource with urn %s. Entity does not exist.", resourceUrn));
diff --git a/datahub-web-react/src/app/entity/shared/components/styled/search/action/DomainsDropdown.tsx b/datahub-web-react/src/app/entity/shared/components/styled/search/action/DomainsDropdown.tsx
index 79b8702bd5310..7ce8222aad379 100644
--- a/datahub-web-react/src/app/entity/shared/components/styled/search/action/DomainsDropdown.tsx
+++ b/datahub-web-react/src/app/entity/shared/components/styled/search/action/DomainsDropdown.tsx
@@ -15,23 +15,24 @@ export default function DomainsDropdown({ urns, disabled = false, refetch }: Pro
     const [isEditModalVisible, setIsEditModalVisible] = useState(false);
     const [batchSetDomainMutation] = useBatchSetDomainMutation();
 
-    const batchUnsetDomains = async () => {
-        try {
-            await batchSetDomainMutation({
-                variables: {
-                    input: {
-                        resources: [...urns.map((urn) => ({ resourceUrn: urn }))],
-                    },
+    const batchUnsetDomains = () => {
+        batchSetDomainMutation({
+            variables: {
+                input: {
+                    resources: [...urns.map((urn) => ({ resourceUrn: urn }))],
                 },
+            },
+        })
+            .then(({ errors }) => {
+                if (!errors) {
+                    message.success({ content: 'Removed Domain!', duration: 2 });
+                    refetch?.();
+                }
+            })
+            .catch((e) => {
+                message.destroy();
+                message.error({ content: `Failed to remove assets from Domain: \n ${e.message || ''}`, duration: 3 });
             });
-            message.success({ content: 'Removed Domain!', duration: 2 });
-        } catch (e: unknown) {
-            message.destroy();
-            if (e instanceof Error) {
-                message.error({ content: `Failed to re,pve Domain: \n ${e.message || ''}`, duration: 3 });
-            }
-        }
-        refetch?.();
     };
 
     return (
diff --git a/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Domain/SetDomainModal.tsx b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Domain/SetDomainModal.tsx
index d1a053d1e19cc..fc5ae773b19ea 100644
--- a/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Domain/SetDomainModal.tsx
+++ b/datahub-web-react/src/app/entity/shared/containers/profile/sidebar/Domain/SetDomainModal.tsx
@@ -102,25 +102,25 @@ export const SetDomainModal = ({ urns, onCloseModal, refetch }: Props) => {
         if (!selectedDomain) {
             return;
         }
-        try {
-            await batchSetDomainMutation({
-                variables: {
-                    input: {
-                        resources: [...urns.map((urn) => ({ resourceUrn: urn }))],
-                        domainUrn: selectedDomain.urn,
-                    },
+        batchSetDomainMutation({
+            variables: {
+                input: {
+                    resources: [...urns.map((urn) => ({ resourceUrn: urn }))],
                 },
+            },
+        })
+            .then(({ errors }) => {
+                if (!errors) {
+                    message.success({ content: 'Updated Domain!', duration: 2 });
+                    refetch?.();
+                    onModalClose();
+                    setSelectedDomain(undefined);
+                }
+            })
+            .catch((e) => {
+                message.destroy();
+                message.error({ content: `Failed to add assets to Domain: \n ${e.message || ''}`, duration: 3 });
             });
-            message.success({ content: 'Updated Domain!', duration: 2 });
-        } catch (e: unknown) {
-            message.destroy();
-            if (e instanceof Error) {
-                message.error({ content: `Failed to set Domain: \n ${e.message || ''}`, duration: 3 });
-            }
-        }
-        setSelectedDomain(undefined);
-        refetch?.();
-        onModalClose();
     };
 
     const selectValue = (selectedDomain && [selectedDomain?.displayName]) || undefined;