Fix increment tab counters when linking nodes

Fix #171

Signed-off-by: Sacha Al Himdani <sacha@getprobo.com>

apps/console/src/components/tasks/TaskFormDialog.tsx

@@ -15,7 +15,7 @@ import type { ReactNode } from "react";
 import { useTranslate } from "@probo/i18n";
 import { Breadcrumb } from "@probo/ui";
 import { graphql } from "relay-runtime";
-import { useFragment } from "react-relay";
+import { useFragment, useRelayEnvironment } from "react-relay";
 import { z } from "zod";
 import { useFormWithSchema } from "/hooks/useFormWithSchema";
 import { useMutationWithToasts } from "/hooks/useMutationWithToasts";
@@ -24,6 +24,7 @@ import { PeopleSelectField } from "/components/form/PeopleSelectField";
 import type { TaskFormDialogFragment$key } from "./__generated__/TaskFormDialogFragment.graphql";
 import { MeasureSelectField } from "/components/form/MeasureSelectField";
 import { Controller } from "react-hook-form";
+import { updateStoreCounter } from "/hooks/useMutationWithIncrement";
 
 const taskFragment = graphql`
   fragment TaskFormDialogFragment on Task {
@@ -93,6 +94,7 @@ export default function TaskFormDialog(props: Props) {
   const dialogRef = props.ref ?? useDialogRef();
   const organizationId = useOrganizationId();
   const task = useFragment(taskFragment, props.task);
+  const relayEnv = useRelayEnvironment();
   const [mutate] = task
     ? useMutationWithToasts(taskUpdateMutation, {
         successMessage: __("Task updated successfully."),
@@ -142,6 +144,9 @@ export default function TaskFormDialog(props: Props) {
           },
           connections: [props.connection!],
         },
+        onCompleted: () => {
+          updateStoreCounter(relayEnv, data.measureId, "tasks(first:0)", 1);
+        },
       });
       reset();
     }

apps/console/src/components/tasks/TasksCard.tsx

@@ -16,7 +16,7 @@ import {
   useDialogRef,
 } from "@probo/ui";
 import { Fragment } from "react";
-import { graphql, useMutation } from "react-relay";
+import { graphql, useMutation, useRelayEnvironment } from "react-relay";
 import { useTranslate } from "@probo/i18n";
 import { usePageTitle } from "@probo/hooks";
 import type { ItemOf } from "/types";
@@ -24,9 +24,10 @@ import TaskFormDialog, {
   taskUpdateMutation,
 } from "/components/tasks/TaskFormDialog";
 import { useOrganizationId } from "/hooks/useOrganizationId";
-import { Link, useLocation } from "react-router";
+import { Link, useLocation, useParams } from "react-router";
 import { promisifyMutation } from "@probo/helpers";
 import type { TaskFormDialogFragment$key } from "./__generated__/TaskFormDialogFragment.graphql";
+import { updateStoreCounter } from "/hooks/useMutationWithIncrement";
 
 type Props = {
   tasks: ({
@@ -94,15 +95,13 @@ export default function TasksCard({ tasks, connectionId }: Props) {
                         <TaskStateIcon state={h.state!} />
                         {h.label}
                       </h2>
-                      {tasksPerHash
-                        .get(h.hash)
-                        ?.map((task) => (
-                          <TaskRow
-                            key={task.id}
-                            task={task}
-                            connectionId={connectionId}
-                          />
-                        ))}
+                      {tasksPerHash.get(h.hash)?.map((task) => (
+                        <TaskRow
+                          key={task.id}
+                          task={task}
+                          connectionId={connectionId}
+                        />
+                      ))}
                     </Fragment>
                   ))
               : // Todo and Done tab simply list todos
@@ -142,7 +141,9 @@ function TaskRow(props: TaskRowProps) {
   const { __ } = useTranslate();
   const confirm = useConfirm();
   const [deleteTask] = useMutation(deleteMutation);
+  const params = useParams<{ measureId?: string }>();
 
+  const relayEnv = useRelayEnvironment();
   const [updateTask, isUpdating] = useMutation(taskUpdateMutation);
 
   const onToggle = () => {
@@ -164,6 +165,16 @@ function TaskRow(props: TaskRowProps) {
             input: { taskId: props.task.id },
             connections: [props.connectionId],
           },
+          onCompleted: () => {
+            if (params.measureId) {
+              updateStoreCounter(
+                relayEnv,
+                params.measureId,
+                "tasks(first:0)",
+                -1
+              );
+            }
+          },
         }),
       {
         message: "Are you sure you want to delete this task?",

apps/console/src/hooks/useMutationWithIncrement.ts

@@ -0,0 +1,70 @@
+import { useCallback } from "react";
+import {
+  useMutation,
+  type UseMutationConfig,
+  useRelayEnvironment,
+} from "react-relay";
+import {
+  commitLocalUpdate,
+  type GraphQLTaggedNode,
+  type MutationParameters,
+} from "relay-runtime";
+import type RelayModernEnvironment from "relay-runtime/lib/store/RelayModernEnvironment";
+
+const defaultOptions = {
+  field: "totalCount",
+  value: 1,
+};
+
+/**
+ * A decorated useMutation hook that increments the store on complete.
+ */
+export function useMutationWithIncrement<T extends MutationParameters>(
+  query: GraphQLTaggedNode,
+  baseOptions: {
+    id: string;
+    node: string;
+    field?: string;
+    value?: 1 | -1;
+  },
+) {
+  const [mutate, isLoading] = useMutation<T>(query);
+  const relayEnv = useRelayEnvironment();
+  const options = { ...defaultOptions, ...baseOptions };
+  const mutateAndIncrement = useCallback(
+    (queryOptions: UseMutationConfig<T>) => {
+      return mutate({
+        ...queryOptions,
+        onCompleted: (response, error) => {
+          updateStoreCounter(
+            relayEnv,
+            options.id,
+            options.node,
+            options.value,
+            options.field,
+          );
+          queryOptions.onCompleted?.(response, error);
+        },
+      });
+    },
+    [mutate, options.id, options.node, options.field, options.value, relayEnv],
+  );
+
+  return [mutateAndIncrement, isLoading] as const;
+}
+
+export function updateStoreCounter(
+  relayEnv: RelayModernEnvironment,
+  recordId: string,
+  nodeName: string,
+  value: number = 1,
+  fieldName: string = "totalCount",
+) {
+  commitLocalUpdate(relayEnv, (store) => {
+    const node = store?.get(recordId)?.getLinkedRecord(nodeName);
+    const previousValue = node?.getValue(fieldName);
+    if (node && typeof previousValue === "number") {
+      node.setValue(previousValue + value, fieldName);
+    }
+  });
+}

apps/console/src/pages/organizations/documents/tabs/DocumentControlsTab.tsx

@@ -1,8 +1,9 @@
 import { LinkedControlsCard } from "/components/controls/LinkedControlsCard";
 import { useOutletContext } from "react-router";
 import { graphql } from "relay-runtime";
-import { useMutation, useRefetchableFragment } from "react-relay";
+import { useRefetchableFragment } from "react-relay";
 import type { DocumentControlsTabFragment$key } from "./__generated__/DocumentControlsTabFragment.graphql";
+import { useMutationWithIncrement } from "/hooks/useMutationWithIncrement.ts";
 
 export const controlsFragment = graphql`
   fragment DocumentControlsTabFragment on Document
@@ -68,8 +69,24 @@ export default function DocumentControlsTab() {
   }>();
   const [data, refetch] = useRefetchableFragment(controlsFragment, document);
   const controls = data.controls.edges.map((edge) => edge.node);
-  const [detachControl, isDetaching] = useMutation(detachControlMutation);
-  const [attachControl, isAttaching] = useMutation(attachControlMutation);
+  const incrementOptions = {
+    id: data.id,
+    node: "controls(first:0)",
+  };
+  const [detachControl, isDetaching] = useMutationWithIncrement(
+    detachControlMutation,
+    {
+      ...incrementOptions,
+      value: -1,
+    },
+  );
+  const [attachControl, isAttaching] = useMutationWithIncrement(
+    attachControlMutation,
+    {
+      ...incrementOptions,
+      value: 1,
+    },
+  );
   const isLoading = isDetaching || isAttaching;
   return (
     <LinkedControlsCard

apps/console/src/pages/organizations/measures/dialog/CreateEvidenceDialog.tsx

@@ -12,11 +12,12 @@ import {
   type DialogRef,
 } from "@probo/ui";
 import { useTranslate } from "@probo/i18n";
-import { graphql } from "react-relay";
+import { graphql, useRelayEnvironment } from "react-relay";
 import { useState } from "react";
 import { z } from "zod";
 import { useFormWithSchema } from "/hooks/useFormWithSchema";
 import { useMutationWithToasts } from "/hooks/useMutationWithToasts";
+import { updateStoreCounter } from "/hooks/useMutationWithIncrement";
 
 const uploadEvidenceMutation = graphql`
   mutation CreateEvidenceDialogUploadMutation(
@@ -73,6 +74,7 @@ export function CreateEvidenceDialog(props: Props) {
 function EvidenceUpload({ measureId, connectionId }: Omit<Props, "ref">) {
   const { __ } = useTranslate();
 
+  const relayEnv = useRelayEnvironment();
   const [mutate, isUpdating] = useMutationWithToasts(uploadEvidenceMutation, {
     successMessage: __("Evidence uploaded successfully"),
     errorMessage: __("Failed to create evidence"),
@@ -90,6 +92,9 @@ function EvidenceUpload({ measureId, connectionId }: Omit<Props, "ref">) {
         uploadables: {
           "input.file": file,
         },
+        onSuccess: () => {
+          updateStoreCounter(relayEnv, measureId, "evidences(first:0)", 1);
+        },
       });
     }
   };

apps/console/src/pages/organizations/measures/tabs/MeasureControlsTab.tsx

@@ -1,11 +1,8 @@
-import {
-  graphql,
-  useMutation,
-  useRefetchableFragment,
-} from "react-relay";
+import { graphql, useRefetchableFragment } from "react-relay";
 import { useOutletContext } from "react-router";
 import { LinkedControlsCard } from "/components/controls/LinkedControlsCard";
 import type { MeasureControlsTabFragment$key } from "./__generated__/MeasureControlsTabFragment.graphql";
+import { useMutationWithIncrement } from "/hooks/useMutationWithIncrement.ts";
 
 export const controlsFragment = graphql`
   fragment MeasureControlsTabFragment on Measure
@@ -73,8 +70,24 @@ export default function MeasureControlsTab() {
   const connectionId = data.controls.__id;
   const controls = data.controls?.edges?.map((edge) => edge.node) ?? [];
 
-  const [detachControl, isDetaching] = useMutation(detachControlMutation);
-  const [attachControl, isAttaching] = useMutation(attachControlMutation);
+  const incrementOptions = {
+    id: data.id,
+    node: "controls(first:0)",
+  };
+  const [detachControl, isDetaching] = useMutationWithIncrement(
+    detachControlMutation,
+    {
+      ...incrementOptions,
+      value: -1,
+    },
+  );
+  const [attachControl, isAttaching] = useMutationWithIncrement(
+    attachControlMutation,
+    {
+      ...incrementOptions,
+      value: 1,
+    },
+  );
   const isLoading = isDetaching || isAttaching;
 
   return (

apps/console/src/pages/organizations/measures/tabs/MeasureEvidencesTab.tsx

@@ -18,7 +18,12 @@ import {
   useDialogRef,
 } from "@probo/ui";
 import { graphql } from "relay-runtime";
-import { useFragment, useMutation, usePaginationFragment } from "react-relay";
+import {
+  useFragment,
+  useMutation,
+  usePaginationFragment,
+  useRelayEnvironment,
+} from "react-relay";
 import { SortableTable } from "/components/SortableTable";
 import type { MeasureEvidencesTabFragment_evidence$key } from "./__generated__/MeasureEvidencesTabFragment_evidence.graphql";
 import { fileSize, fileType, promisifyMutation, sprintf } from "@probo/helpers";
@@ -27,6 +32,7 @@ import { useOrganizationId } from "/hooks/useOrganizationId";
 import { CreateEvidenceDialog } from "../dialog/CreateEvidenceDialog";
 import { useState } from "react";
 import { EvidenceDownloadDialog } from "../dialog/EvidenceDownloadDialog";
+import { updateStoreCounter } from "/hooks/useMutationWithIncrement";
 
 export const evidencesFragment = graphql`
   fragment MeasureEvidencesTabFragment on Measure
@@ -134,7 +140,7 @@ export default function MeasureEvidencesTab() {
           key={evidence?.id}
           onClose={() =>
             navigate(
-              `/organizations/${organizationId}/measures/${measure.id}/evidences`,
+              `/organizations/${organizationId}/measures/${measure.id}/evidences`
             )
           }
           evidenceId={evidence.id}
@@ -162,6 +168,7 @@ function EvidenceRow(props: {
   const [mutate, isDeleting] = useMutation(deleteEvidenceMutation);
   const confirm = useConfirm();
   const [isDownloading, setIsDownloading] = useState(false);
+  const relayEnv = useRelayEnvironment();
 
   const handleDelete = () => {
     confirm(
@@ -173,16 +180,24 @@ function EvidenceRow(props: {
               evidenceId: evidence.id,
             },
           },
+          onCompleted: () => {
+            updateStoreCounter(
+              relayEnv,
+              props.measureId,
+              "evidences(first:0)",
+              -1
+            );
+          },
         });
       },
       {
         message: sprintf(
           __(
-            'This will permanently delete the evidence "%s". This action cannot be undone.',
+            'This will permanently delete the evidence "%s". This action cannot be undone.'
           ),
-          evidence.filename,
+          evidence.filename
         ),
-      },
+      }
     );
   };
 

apps/console/src/pages/organizations/measures/tabs/MeasureRisksTab.tsx

@@ -1,7 +1,8 @@
-import { graphql, useFragment, useMutation } from "react-relay";
+import { graphql, useFragment } from "react-relay";
 import type { MeasureRisksTabFragment$key } from "./__generated__/MeasureRisksTabFragment.graphql";
 import { useOutletContext } from "react-router";
 import { LinkedRisksCard } from "/components/risks/LinkedRisksCard";
+import { useMutationWithIncrement } from "/hooks/useMutationWithIncrement";
 
 export const risksFragment = graphql`
   fragment MeasureRisksTabFragment on Measure {
@@ -53,8 +54,24 @@ export default function MeasureRisksTab() {
   const connectionId = data.risks.__id;
   const risks = data.risks?.edges?.map((edge) => edge.node) ?? [];
 
-  const [detachRisk, isDetaching] = useMutation(detachRiskMutation);
-  const [attachRisk, isAttaching] = useMutation(attachRiskMutation);
+  const incrementOptions = {
+    id: data.id,
+    node: "risks(first:0)",
+  };
+  const [detachRisk, isDetaching] = useMutationWithIncrement(
+    detachRiskMutation,
+    {
+      ...incrementOptions,
+      value: -1,
+    },
+  );
+  const [attachRisk, isAttaching] = useMutationWithIncrement(
+    attachRiskMutation,
+    {
+      ...incrementOptions,
+      value: 1,
+    },
+  );
   const isLoading = isDetaching || isAttaching;
 
   return (

apps/console/src/pages/organizations/risks/tabs/RiskDocumentsTab.tsx

@@ -1,7 +1,8 @@
-import { graphql, useFragment, useMutation } from "react-relay";
+import { graphql, useFragment } from "react-relay";
 import { useOutletContext } from "react-router";
 import { LinkedDocumentsCard } from "/components/documents/LinkedDocumentsCard";
 import type { RiskDocumentsTabFragment$key } from "./__generated__/RiskDocumentsTabFragment.graphql";
+import { useMutationWithIncrement } from "/hooks/useMutationWithIncrement.ts";
 
 export const documentsFragment = graphql`
   fragment RiskDocumentsTabFragment on Risk {
@@ -53,8 +54,24 @@ export default function RiskDocumentsTab() {
   const connectionId = data.documents.__id;
   const documents = data.documents?.edges?.map((edge) => edge.node) ?? [];
 
-  const [detachDocument, isDetaching] = useMutation(detachDocumentMutation);
-  const [attachDocument, isAttaching] = useMutation(attachDocumentMutation);
+  const incrementOptions = {
+    id: data.id,
+    node: "documents(first:0)",
+  };
+  const [detachDocument, isDetaching] = useMutationWithIncrement(
+    detachDocumentMutation,
+    {
+      ...incrementOptions,
+      value: -1,
+    },
+  );
+  const [attachDocument, isAttaching] = useMutationWithIncrement(
+    attachDocumentMutation,
+    {
+      ...incrementOptions,
+      value: 1,
+    },
+  );
   const isLoading = isDetaching || isAttaching;
 
   return (

apps/console/src/pages/organizations/risks/tabs/RiskMeasuresTab.tsx

@@ -1,7 +1,8 @@
-import { graphql, useFragment, useMutation } from "react-relay";
+import { graphql, useFragment } from "react-relay";
 import type { RiskMeasuresTabFragment$key } from "./__generated__/RiskMeasuresTabFragment.graphql";
 import { useOutletContext } from "react-router";
 import { LinkedMeasuresCard } from "/components/measures/LinkedMeasuresCard";
+import { useMutationWithIncrement } from "/hooks/useMutationWithIncrement.ts";
 
 export const measuresFragment = graphql`
   fragment RiskMeasuresTabFragment on Risk {
@@ -52,9 +53,24 @@ export default function RiskMeasuresTab() {
   const data = useFragment(measuresFragment, risk);
   const connectionId = data.measures.__id;
   const measures = data.measures?.edges?.map((edge) => edge.node) ?? [];
-
-  const [detachMeasure, isDetaching] = useMutation(detachMeasureMutation);
-  const [attachMeasure, isAttaching] = useMutation(attachMeasureMutation);
+  const incrementOptions = {
+    id: data.id,
+    node: "measures(first:0)",
+  };
+  const [detachMeasure, isDetaching] = useMutationWithIncrement(
+    detachMeasureMutation,
+    {
+      ...incrementOptions,
+      value: -1,
+    },
+  );
+  const [attachMeasure, isAttaching] = useMutationWithIncrement(
+    attachMeasureMutation,
+    {
+      ...incrementOptions,
+      value: 1,
+    },
+  );
   const isLoading = isDetaching || isAttaching;
 
   return (