1 files changed, 24 insertions, 5 deletions
diff --git a/app/assets/javascripts/work_items/components/work_item_links/work_item_children_wrapper.vue b/app/assets/javascripts/work_items/components/work_item_links/work_item_children_wrapper.vue
index 9d9414b5399..f4de7c1dddc 100644
--- a/app/assets/javascripts/work_items/components/work_item_links/work_item_children_wrapper.vue
+++ b/app/assets/javascripts/work_items/components/work_item_links/work_item_children_wrapper.vue
@@ -13,6 +13,7 @@ import { findHierarchyWidgets } from '../../utils';
 import { addHierarchyChild, removeHierarchyChild } from '../../graphql/cache_utils';
 import reorderWorkItem from '../../graphql/reorder_work_item.mutation.graphql';
 import updateWorkItemMutation from '../../graphql/update_work_item.mutation.graphql';
+import groupWorkItemByIidQuery from '../../graphql/group_work_item_by_iid.query.graphql';
 import workItemByIidQuery from '../../graphql/work_item_by_iid.query.graphql';
 import WorkItemLinkChild from './work_item_link_child.vue';
 
@@ -20,8 +21,12 @@ export default {
   components: {
     WorkItemLinkChild,
   },
-  inject: ['fullPath'],
+  inject: ['isGroup'],
   props: {
+    fullPath: {
+      type: String,
+      required: true,
+    },
     workItemType: {
       type: String,
       required: false,
@@ -83,7 +88,14 @@ export default {
         const { data } = await this.$apollo.mutate({
           mutation: updateWorkItemMutation,
           variables: { input: { id: child.id, hierarchyWidget: { parentId: null } } },
-          update: (cache) => removeHierarchyChild(cache, this.fullPath, this.workItemIid, child),
+          update: (cache) =>
+            removeHierarchyChild({
+              cache,
+              fullPath: this.fullPath,
+              iid: this.workItemIid,
+              isGroup: this.isGroup,
+              workItem: child,
+            }),
         });
 
         if (data.workItemUpdate.errors.length) {
@@ -109,7 +121,14 @@ export default {
         const { data } = await this.$apollo.mutate({
           mutation: updateWorkItemMutation,
           variables: { input: { id: child.id, hierarchyWidget: { parentId: this.workItemId } } },
-          update: (cache) => addHierarchyChild(cache, this.fullPath, this.workItemIid, child),
+          update: (cache) =>
+            addHierarchyChild({
+              cache,
+              fullPath: this.fullPath,
+              iid: this.workItemIid,
+              isGroup: this.isGroup,
+              workItem: child,
+            }),
         });
 
         if (data.workItemUpdate.errors.length) {
@@ -124,7 +143,7 @@ export default {
     },
     addWorkItemQuery({ iid }) {
       this.$apollo.addSmartQuery('prefetchedWorkItem', {
-        query: workItemByIidQuery,
+        query: this.isGroup ? groupWorkItemByIidQuery : workItemByIidQuery,
         variables: {
           fullPath: this.fullPath,
           iid,
@@ -206,7 +225,7 @@ export default {
           update: (store) => {
             store.updateQuery(
               {
-                query: workItemByIidQuery,
+                query: this.isGroup ? groupWorkItemByIidQuery : workItemByIidQuery,
                 variables: { fullPath: this.fullPath, iid: this.workItemIid },
               },
               (sourceData) =>