feat: review changes

src/core_modules/capture-core/components/WidgetRelatedStages/WidgetRelatedStages.component.js

@@ -1,7 +1,7 @@
 // @flow
 import React, { forwardRef, useCallback, useEffect, useImperativeHandle, useState } from 'react';
-import { useDataQuery } from '@dhis2/app-runtime';
 import { useRelatedStages } from './useRelatedStages';
+import { useOrgUnitAutoSelect } from './hooks/useAutoSelctOrgUnitRelatedStage';
 import type { Props, RelatedStageDataValueStates } from './WidgetRelatedStages.types';
 import { RelatedStagesActions } from './RelatedStagesActions';
 import { relatedStageStatus } from './constants';
@@ -29,16 +29,6 @@ const WidgetRelatedStagesPlain = ({
         occurredLabel,
         enrollmentId,
     });
-
-    const { loading: orgUnitLoading, data: orgUnitData } = useDataQuery({
-        orgUnits: {
-            resource: 'me',
-            params: {
-                fields: ['organisationUnits[id,path,displayName,children::isNotEmpty]'],
-            },
-        },
-    });
-
     const [saveAttempted, setSaveAttempted] = useState(false);
     const [errorMessages, setErrorMessages] = useState({});
     const [relatedStageDataValues, setRelatedStageDataValues] = useState<RelatedStageDataValueStates>({
@@ -47,17 +37,7 @@ const WidgetRelatedStagesPlain = ({
         orgUnit: undefined,
         linkedEventId: undefined,
     });
-
-    useEffect(() => {
-        const orgUnits = orgUnitData?.orgUnits?.organisationUnits || [];
-        if (orgUnits.length === 1 && !orgUnits[0]?.children) {
-            const { displayName, ...rest } = orgUnits[0];
-            setRelatedStageDataValues(prev => ({
-                ...prev,
-                orgUnit: { ...rest, name: displayName },
-            }));
-        }
-    }, [orgUnitData]);
+    const { isLoading: orgUnitLoading } = useOrgUnitAutoSelect(setRelatedStageDataValues);
 
     const addErrorMessage = (message: ErrorMessagesForRelatedStages) => {
         setErrorMessages((prevMessages: Object) => ({

src/core_modules/capture-core/components/WidgetRelatedStages/hooks/useAutoSelctOrgUnitRelatedStage.js

@@ -0,0 +1,35 @@
+// @flow
+import { useEffect } from 'react';
+import { useApiMetadataQuery } from '../../../utils/reactQueryHelpers';
+
+export const useOrgUnitAutoSelect = (setRelatedStageDataValues: any) => {
+    const queryKey = ['organisationUnits'];
+    const queryFn = {
+        resource: 'organisationUnits',
+        params: {
+            fields: 'id,path,displayName,children::isNotEmpty',
+            withinUserSearchHierarchy: true,
+            paging: 2,
+        },
+    };
+    const { data, isLoading } = useApiMetadataQuery(queryKey, queryFn);
+
+    useEffect(() => {
+        const orgUnits = data?.organisationUnits;
+        if (isLoading || !orgUnits || orgUnits.length !== 1) return;
+
+        const [orgUnit] = orgUnits;
+        if (orgUnit && !orgUnit.children) {
+            const { displayName, ...rest } = orgUnit;
+            setRelatedStageDataValues(prev => ({
+                ...prev,
+                orgUnit: { ...rest, name: displayName },
+            }));
+        }
+    }, [data, isLoading, setRelatedStageDataValues]);
+
+
+    return {
+        isLoading,
+    };
+};