feat(waiting-room): choose relations when associated endpoint [EE-5187] (#8720)

app/react/edge/edge-devices/WaitingRoomView/Datatable/AssignmentDialog/AssignmentDialog.tsx

@@ -0,0 +1,166 @@
+import { Form, Formik } from 'formik';
+
+import { addPlural } from '@/portainer/helpers/strings';
+import { useUpdateEnvironmentsRelationsMutation } from '@/react/portainer/environments/queries/useUpdateEnvironmentsRelationsMutation';
+import { notifySuccess } from '@/portainer/services/notifications';
+
+import { Checkbox } from '@@/form-components/Checkbox';
+import { FormControl } from '@@/form-components/FormControl';
+import { OnSubmit, Modal } from '@@/modals';
+import { TextTip } from '@@/Tip/TextTip';
+import { Button, LoadingButton } from '@@/buttons';
+
+import { WaitingRoomEnvironment } from '../../types';
+
+import { GroupSelector, EdgeGroupsSelector, TagSelector } from './Selectors';
+import { FormValues } from './types';
+import { isAssignedToGroup } from './utils';
+import { createPayload } from './createPayload';
+
+export function AssignmentDialog({
+  onSubmit,
+  environments,
+}: {
+  onSubmit: OnSubmit<boolean>;
+  environments: Array<WaitingRoomEnvironment>;
+}) {
+  const assignRelationsMutation = useUpdateEnvironmentsRelationsMutation();
+
+  const initialValues: FormValues = {
+    group: 1,
+    overrideGroup: false,
+    edgeGroups: [],
+    overrideEdgeGroups: false,
+    tags: [],
+    overrideTags: false,
+  };
+
+  const hasPreAssignedEdgeGroups = environments.some(
+    (e) => e.EdgeGroups?.length > 0
+  );
+  const hasPreAssignedTags = environments.some((e) => e.TagIds.length > 0);
+  const hasPreAssignedGroup = environments.some((e) => isAssignedToGroup(e));
+
+  return (
+    <Modal
+      aria-label="Associate and assignment"
+      onDismiss={() => onSubmit()}
+      size="lg"
+    >
+      <Modal.Header
+        title={`Associate with assignment (${addPlural(
+          environments.length,
+          'selected edge environment'
+        )})`}
+      />
+      <Formik onSubmit={handleSubmit} initialValues={initialValues}>
+        {({ values, setFieldValue, errors }) => (
+          <Form noValidate>
+            <Modal.Body>
+              <div>
+                <FormControl
+                  size="vertical"
+                  label="Group"
+                  tooltip="For managing RBAC with user access"
+                  errors={errors.group}
+                >
+                  <GroupSelector />
+
+                  {hasPreAssignedGroup && (
+                    <div className="mt-2">
+                      <Checkbox
+                        label="Override pre-assigned group"
+                        id="overrideGroup"
+                        bold={false}
+                        checked={values.overrideGroup}
+                        onChange={(e) =>
+                          setFieldValue('overrideGroup', e.target.checked)
+                        }
+                      />
+                    </div>
+                  )}
+                </FormControl>
+
+                <FormControl
+                  size="vertical"
+                  label="Edge Groups"
+                  tooltip="Required to manage edge job and edge stack deployments"
+                  errors={errors.edgeGroups}
+                >
+                  <EdgeGroupsSelector />
+
+                  {hasPreAssignedEdgeGroups && (
+                    <div className="mt-2">
+                      <Checkbox
+                        label="Override pre-assigned edge groups"
+                        bold={false}
+                        id="overrideEdgeGroups"
+                        checked={values.overrideEdgeGroups}
+                        onChange={(e) =>
+                          setFieldValue('overrideEdgeGroups', e.target.checked)
+                        }
+                      />
+                    </div>
+                  )}
+                </FormControl>
+
+                <div className="mb-3">
+                  <TextTip color="blue">
+                    Edge group(s) created here are static only, use tags to
+                    assign to dynamic edge groups
+                  </TextTip>
+                </div>
+
+                <FormControl
+                  size="vertical"
+                  label="Tags"
+                  tooltip="Assigning tags will auto populate environments to dynamic edge groups that these tags are assigned to and any ege jobs or stacks that are deployed to that edge group"
+                  errors={errors.tags}
+                >
+                  <TagSelector />
+
+                  {hasPreAssignedTags && (
+                    <div className="mt-2">
+                      <Checkbox
+                        label="Override pre-assigned tags"
+                        bold={false}
+                        id="overrideTags"
+                        checked={values.overrideTags}
+                        onChange={(e) =>
+                          setFieldValue('overrideTags', e.target.checked)
+                        }
+                      />
+                    </div>
+                  )}
+                </FormControl>
+              </div>
+            </Modal.Body>
+            <Modal.Footer>
+              <Button onClick={() => onSubmit()} color="default">
+                Cancel
+              </Button>
+              <LoadingButton
+                isLoading={assignRelationsMutation.isLoading}
+                loadingText="Associating..."
+              >
+                Associate
+              </LoadingButton>
+            </Modal.Footer>
+          </Form>
+        )}
+      </Formik>
+    </Modal>
+  );
+
+  function handleSubmit(values: FormValues) {
+    assignRelationsMutation.mutate(
+      Object.fromEntries(environments.map((e) => createPayload(e, values))),
+      {
+        onSuccess: () => {
+          notifySuccess('Success', 'Edge environments assigned successfully');
+          onSubmit(true);
+        },
+      }
+    );
+  }
+}

app/react/edge/edge-devices/WaitingRoomView/Datatable/AssignmentDialog/Selectors/CreatableSelector.tsx

@@ -0,0 +1,60 @@
+import { useField } from 'formik';
+import _ from 'lodash';
+
+import { Select } from '@@/form-components/ReactSelect';
+import {
+  Option,
+  Option as OptionType,
+} from '@@/form-components/PortainerSelect';
+
+export function CreatableSelector({
+  name,
+  options,
+  onCreate,
+  isLoading,
+}: {
+  name: string;
+  options: Array<OptionType<number>>;
+  onCreate: (label: string) => Promise<number>;
+  isLoading: boolean;
+}) {
+  const [{ onBlur, value }, , { setValue }] = useField<Array<number>>(name);
+
+  const selectedValues = value.reduce(
+    (acc: Array<OptionType<number>>, cur) =>
+      _.compact([...acc, findOption(cur, options)]),
+    []
+  );
+
+  return (
+    <Select
+      isCreatable
+      options={options}
+      value={
+        isLoading
+          ? [...selectedValues, { label: 'Creating...', value: 0 }]
+          : selectedValues
+      }
+      isMulti
+      onCreateOption={handleCreate}
+      onChange={handleChange}
+      onBlur={onBlur}
+      isLoading={isLoading}
+      isDisabled={isLoading}
+      closeMenuOnSelect={false}
+    />
+  );
+
+  async function handleCreate(label: string) {
+    const id = await onCreate(label);
+    setValue([...value, id]);
+  }
+
+  function handleChange(value: ReadonlyArray<{ value: number }>) {
+    setValue(value.map((v) => v.value));
+  }
+}
+
+function findOption<T>(option: T, options: Array<Option<T>>) {
+  return options.find((t) => t.value === option);
+}

app/react/edge/edge-devices/WaitingRoomView/Datatable/AssignmentDialog/Selectors/EdgeGroupSelector.tsx

@@ -0,0 +1,41 @@
+import { notifySuccess } from '@/portainer/services/notifications';
+import { useCreateGroupMutation } from '@/react/edge/edge-groups/queries/useCreateEdgeGroupMutation';
+import { useEdgeGroups } from '@/react/edge/edge-groups/queries/useEdgeGroups';
+
+import { CreatableSelector } from './CreatableSelector';
+
+export function EdgeGroupsSelector() {
+  const createMutation = useCreateGroupMutation();
+
+  const edgeGroupsQuery = useEdgeGroups({
+    select: (edgeGroups) =>
+      edgeGroups
+        .filter((g) => !g.Dynamic)
+        .map((opt) => ({ label: opt.Name, value: opt.Id })),
+  });
+
+  if (!edgeGroupsQuery.data) {
+    return null;
+  }
+
+  const edgeGroups = edgeGroupsQuery.data;
+
+  return (
+    <CreatableSelector
+      name="edgeGroups"
+      options={edgeGroups}
+      onCreate={handleCreate}
+      isLoading={createMutation.isLoading}
+    />
+  );
+
+  async function handleCreate(newGroup: string) {
+    const group = await createMutation.mutateAsync({
+      name: newGroup,
+      dynamic: false,
+    });
+
+    notifySuccess('Edge group created', `Group ${group.Name} created`);
+    return group.Id;
+  }
+}

app/react/edge/edge-devices/WaitingRoomView/Datatable/AssignmentDialog/Selectors/GroupSelector.tsx

@@ -0,0 +1,123 @@
+import { useCallback, useState } from 'react';
+import { useField } from 'formik';
+
+import { useGroups } from '@/react/portainer/environments/environment-groups/queries';
+import { EnvironmentGroupId } from '@/react/portainer/environments/environment-groups/types';
+import { useCreateGroupMutation } from '@/react/portainer/environments/environment-groups/queries/useCreateGroupMutation';
+import { notifySuccess } from '@/portainer/services/notifications';
+
+import { Select } from '@@/form-components/ReactSelect';
+import { Option } from '@@/form-components/PortainerSelect';
+
+import { FormValues } from '../types';
+
+export function GroupSelector() {
+  const [{ value, onBlur }, , { setValue }] =
+    useField<FormValues['group']>('group');
+  const createMutation = useCreateGroupMutation();
+
+  const groupsQuery = useGroups({
+    select: (groups) =>
+      groups
+        .filter((g) => g.Id !== 1)
+        .map((opt) => ({ label: opt.Name, value: opt.Id })),
+  });
+
+  const { onInputChange, clearInputValue } = useCreateOnBlur({
+    options: groupsQuery.data || [],
+    setValue,
+    createValue: handleCreate,
+  });
+
+  if (!groupsQuery.data) {
+    return null;
+  }
+
+  const options = groupsQuery.data;
+  const selectedValue = value ? options.find((g) => g.value === value) : null;
+
+  return (
+    <Select
+      isCreatable
+      options={options}
+      value={
+        createMutation.isLoading
+          ? { label: 'Creating...', value: 0 }
+          : selectedValue
+      }
+      onCreateOption={handleCreate}
+      onChange={handleChange}
+      onInputChange={onInputChange}
+      onBlur={onBlur}
+      isLoading={createMutation.isLoading}
+      isDisabled={createMutation.isLoading}
+      placeholder="Select a group"
+      isClearable
+    />
+  );
+
+  function handleCreate(newGroup: string) {
+    createMutation.mutate(
+      { name: newGroup },
+      {
+        onSuccess: (data) => {
+          setValue(data.Id);
+          notifySuccess('Group created', `Group ${data.Name} created`);
+        },
+      }
+    );
+  }
+
+  function handleChange(value: { value: EnvironmentGroupId } | null) {
+    setValue(value ? value.value : 1);
+    clearInputValue();
+  }
+}
+
+function useCreateOnBlur({
+  options,
+  setValue,
+  createValue,
+}: {
+  options: Option<number>[];
+  setValue: (value: number) => void;
+  createValue: (value: string) => void;
+}) {
+  const [inputValue, setInputValue] = useState('');
+
+  const handleBlur = useCallback(() => {
+    const label = inputValue?.trim() || '';
+    if (!label) {
+      return;
+    }
+
+    const option = options.find((opt) => opt.label === label);
+    if (option) {
+      setValue(option.value);
+    } else {
+      createValue(label);
+    }
+    setInputValue('');
+  }, [createValue, inputValue, options, setValue]);
+
+  const handleInputChange = useCallback(
+    (inputValue, { action }) => {
+      if (action === 'input-change') {
+        setInputValue(inputValue);
+      }
+      if (action === 'input-blur') {
+        handleBlur();
+      }
+    },
+    [handleBlur]
+  );
+
+  const clearInputValue = useCallback(() => {
+    setInputValue('');
+  }, []);
+
+  return {
+    onInputChange: handleInputChange,
+    clearInputValue,
+  };
+}

app/react/edge/edge-devices/WaitingRoomView/Datatable/AssignmentDialog/Selectors/TagSelector.tsx

@@ -0,0 +1,35 @@
+import { notifySuccess } from '@/portainer/services/notifications';
+import { useCreateTagMutation, useTags } from '@/portainer/tags/queries';
+
+import { CreatableSelector } from './CreatableSelector';
+
+export function TagSelector() {
+  const createMutation = useCreateTagMutation();
+
+  const tagsQuery = useTags({
+    select: (tags) => tags.map((opt) => ({ label: opt.Name, value: opt.ID })),
+  });
+
+  if (!tagsQuery.data) {
+    return null;
+  }
+
+  const tags = tagsQuery.data;
+
+  return (
+    <CreatableSelector
+      name="tags"
+      options={tags}
+      onCreate={handleCreate}
+      isLoading={createMutation.isLoading}
+    />
+  );
+
+  async function handleCreate(newTag: string) {
+    const tag = await createMutation.mutateAsync(newTag);
+
+    notifySuccess('Tag created', `Tag ${tag.Name} created`);
+
+    return tag.ID;
+  }
+}

app/react/edge/edge-devices/WaitingRoomView/Datatable/AssignmentDialog/Selectors/index.ts

@@ -0,0 +1,3 @@
+export { EdgeGroupsSelector } from './EdgeGroupSelector';
+export { GroupSelector } from './GroupSelector';
+export { TagSelector } from './TagSelector';

app/react/edge/edge-devices/WaitingRoomView/Datatable/AssignmentDialog/createPayload.tsx

@@ -0,0 +1,30 @@
+import { EnvironmentRelationsPayload } from '@/react/portainer/environments/queries/useUpdateEnvironmentsRelationsMutation';
+
+import { WaitingRoomEnvironment } from '../../types';
+
+import { FormValues } from './types';
+import { isAssignedToGroup } from './utils';
+
+export function createPayload(
+  environment: WaitingRoomEnvironment,
+  values: FormValues
+) {
+  const relations: Partial<EnvironmentRelationsPayload> = {};
+
+  if (environment.TagIds.length === 0 || values.overrideTags) {
+    relations.tags = values.tags;
+  }
+
+  if (environment.EdgeGroups.length === 0 || values.overrideEdgeGroups) {
+    relations.edgeGroups = values.edgeGroups;
+  }
+
+  if (
+    (!isAssignedToGroup(environment) || values.overrideGroup) &&
+    values.group
+  ) {
+    relations.group = values.group;
+  }
+
+  return [environment.Id, relations];
+}

app/react/edge/edge-devices/WaitingRoomView/Datatable/AssignmentDialog/types.ts

@@ -0,0 +1,12 @@
+import { TagId } from '@/portainer/tags/types';
+import { EdgeGroup } from '@/react/edge/edge-groups/types';
+import { EnvironmentGroupId } from '@/react/portainer/environments/environment-groups/types';
+
+export interface FormValues {
+  group: EnvironmentGroupId | null;
+  overrideGroup: boolean;
+  edgeGroups: Array<EdgeGroup['Id']>;
+  overrideEdgeGroups: boolean;
+  tags: Array<TagId>;
+  overrideTags: boolean;
+}

app/react/edge/edge-devices/WaitingRoomView/Datatable/AssignmentDialog/utils.ts

@@ -0,0 +1,5 @@
+import { Environment } from '@/react/portainer/environments/types';
+
+export function isAssignedToGroup(environment: Environment) {
+  return ![0, 1].includes(environment.GroupId);
+}

app/react/edge/edge-devices/WaitingRoomView/Datatable/Datatable.tsx

@@ -1,22 +1,10 @@
-import { Trash2 } from 'lucide-react';
-
-import { Environment } from '@/react/portainer/environments/types';
-import { notifySuccess } from '@/portainer/services/notifications';
-import { useDeleteEnvironmentsMutation } from '@/react/portainer/environments/queries/useDeleteEnvironmentsMutation';
-
 import { Datatable as GenericDatatable } from '@@/datatables';
-import { Button } from '@@/buttons';
 import { createPersistedStore } from '@@/datatables/types';
 import { useTableState } from '@@/datatables/useTableState';
-import { confirm } from '@@/modals/confirm';
-import { buildConfirmButton } from '@@/modals/utils';
-import { ModalType } from '@@/modals';
-import { TooltipWithChildren } from '@@/Tip/TooltipWithChildren';
-
-import { useAssociateDeviceMutation, useLicenseOverused } from '../queries';
 
 import { columns } from './columns';
 import { Filter } from './Filter';
+import { TableActions } from './TableActions';
 import { useEnvironments } from './useEnvironments';
 
 const storageKey = 'edge-devices-waiting-room';
@@ -24,9 +12,6 @@ const storageKey = 'edge-devices-waiting-room';
 const settingsStore = createPersistedStore(storageKey, 'Name');
 
 export function Datatable() {
-  const associateMutation = useAssociateDeviceMutation();
-  const removeMutation = useDeleteEnvironmentsMutation();
-  const { willExceed } = useLicenseOverused();
   const tableState = useTableState(settingsStore, storageKey);
   const { data: environments, totalCount, isLoading } = useEnvironments();
 
@@ -38,76 +23,11 @@ export function Datatable() {
       title="Edge Devices Waiting Room"
       emptyContentLabel="No Edge Devices found"
       renderTableActions={(selectedRows) => (
-        <>
-          <Button
-            onClick={() => handleRemoveDevice(selectedRows)}
-            disabled={selectedRows.length === 0}
-            color="dangerlight"
-            icon={Trash2}
-          >
-            Remove Device
-          </Button>
-
-          <TooltipWithChildren
-            message={
-              willExceed(selectedRows.length) && (
-                <>
-                  Associating devices is disabled as your node count exceeds
-                  your license limit
-                </>
-              )
-            }
-          >
-            <span>
-              <Button
-                onClick={() => handleAssociateDevice(selectedRows)}
-                disabled={
-                  selectedRows.length === 0 || willExceed(selectedRows.length)
-                }
-              >
-                Associate Device
-              </Button>
-            </span>
-          </TooltipWithChildren>
-        </>
+        <TableActions selectedRows={selectedRows} />
       )}
       isLoading={isLoading}
       totalCount={totalCount}
       description={<Filter />}
     />
   );
-
-  function handleAssociateDevice(devices: Environment[]) {
-    associateMutation.mutate(
-      devices.map((d) => d.Id),
-      {
-        onSuccess() {
-          notifySuccess('Success', 'Edge devices associated successfully');
-        },
-      }
-    );
-  }
-
-  async function handleRemoveDevice(devices: Environment[]) {
-    const confirmed = await confirm({
-      title: 'Are you sure?',
-      message:
-        "You're about to remove edge device(s) from waiting room, which will not be shown until next agent startup.",
-      confirmButton: buildConfirmButton('Remove', 'danger'),
-      modalType: ModalType.Destructive,
-    });
-
-    if (!confirmed) {
-      return;
-    }
-
-    removeMutation.mutate(
-      devices.map((d) => d.Id),
-      {
-        onSuccess() {
-          notifySuccess('Success', 'Edge devices were hidden successfully');
-        },
-      }
-    );
-  }
 }

app/react/edge/edge-devices/WaitingRoomView/Datatable/TableActions.tsx

@@ -0,0 +1,142 @@
+import { Check, CheckCircle, Trash2 } from 'lucide-react';
+
+import { notifySuccess } from '@/portainer/services/notifications';
+import { useDeleteEnvironmentsMutation } from '@/react/portainer/environments/queries/useDeleteEnvironmentsMutation';
+import { Environment } from '@/react/portainer/environments/types';
+import { withReactQuery } from '@/react-tools/withReactQuery';
+
+import { Button } from '@@/buttons';
+import { ModalType, openModal } from '@@/modals';
+import { confirm } from '@@/modals/confirm';
+import { buildConfirmButton } from '@@/modals/utils';
+import { TooltipWithChildren } from '@@/Tip/TooltipWithChildren';
+
+import { useAssociateDeviceMutation, useLicenseOverused } from '../queries';
+import { WaitingRoomEnvironment } from '../types';
+
+import { AssignmentDialog } from './AssignmentDialog/AssignmentDialog';
+
+const overusedTooltip = (
+  <>
+    Associating devices is disabled as your node count exceeds your license
+    limit
+  </>
+);
+
+export function TableActions({
+  selectedRows,
+}: {
+  selectedRows: WaitingRoomEnvironment[];
+}) {
+  const associateMutation = useAssociateDeviceMutation();
+  const removeMutation = useDeleteEnvironmentsMutation();
+  const licenseOverused = useLicenseOverused(selectedRows.length);
+
+  return (
+    <>
+      <Button
+        onClick={() => handleRemoveDevice(selectedRows)}
+        disabled={selectedRows.length === 0}
+        color="dangerlight"
+        icon={Trash2}
+      >
+        Remove Device
+      </Button>
+
+      <TooltipWithChildren
+        message={
+          licenseOverused ? (
+            overusedTooltip
+          ) : (
+            <>
+              Associate device(s) and assigning edge groups, group and tags with
+              overriding options
+            </>
+          )
+        }
+      >
+        <span>
+          <Button
+            onClick={() => handleAssociateAndAssign(selectedRows)}
+            disabled={selectedRows.length === 0 || licenseOverused}
+            color="secondary"
+            icon={CheckCircle}
+          >
+            Associate and assignment
+          </Button>
+        </span>
+      </TooltipWithChildren>
+
+      <TooltipWithChildren
+        message={
+          licenseOverused ? (
+            overusedTooltip
+          ) : (
+            <>
+              Associate device(s) based on their pre-assigned edge groups, group
+              and tags
+            </>
+          )
+        }
+      >
+        <span>
+          <Button
+            onClick={() => handleAssociateDevice(selectedRows)}
+            disabled={selectedRows.length === 0 || licenseOverused}
+            icon={Check}
+          >
+            Associate Device
+          </Button>
+        </span>
+      </TooltipWithChildren>
+    </>
+  );
+
+  async function handleAssociateAndAssign(
+    environments: WaitingRoomEnvironment[]
+  ) {
+    const assigned = await openModal(withReactQuery(AssignmentDialog), {
+      environments,
+    });
+
+    if (!assigned) {
+      return;
+    }
+
+    handleAssociateDevice(environments);
+  }
+
+  function handleAssociateDevice(devices: Environment[]) {
+    associateMutation.mutate(
+      devices.map((d) => d.Id),
+      {
+        onSuccess() {
+          notifySuccess('Success', 'Edge devices associated successfully');
+        },
+      }
+    );
+  }
+
+  async function handleRemoveDevice(devices: Environment[]) {
+    const confirmed = await confirm({
+      title: 'Are you sure?',
+      message:
+        "You're about to remove edge device(s) from waiting room, which will not be shown until next agent startup.",
+      confirmButton: buildConfirmButton('Remove', 'danger'),
+      modalType: ModalType.Destructive,
+    });
+
+    if (!confirmed) {
+      return;
+    }
+
+    removeMutation.mutate(
+      devices.map((d) => d.Id),
+      {
+        onSuccess() {
+          notifySuccess('Success', 'Edge devices were hidden successfully');
+        },
+      }
+    );
+  }
+}

app/react/edge/edge-devices/WaitingRoomView/Datatable/useEnvironments.ts

@@ -55,7 +55,7 @@ export function useEnvironments() {
   const envs: Array<WaitingRoomEnvironment> =
     environmentsQuery.environments.map((env) => ({
       ...env,
-      Group: groupsQuery.data?.[env.GroupId] || '',
+      Group: (env.GroupId !== 1 && groupsQuery.data?.[env.GroupId]) || '',
       EdgeGroups:
         environmentEdgeGroupsQuery.data?.[env.Id]?.map((env) => env.group) ||
         [],

app/react/edge/edge-devices/WaitingRoomView/WaitingRoomView.tsx

@@ -13,7 +13,7 @@ export default withLimitToBE(WaitingRoomView);
 
 function WaitingRoomView() {
   const untrustedCount = useUntrustedCount();
-  const { willExceed } = useLicenseOverused();
+  const licenseOverused = useLicenseOverused(untrustedCount);
   return (
     <>
       <PageHeader
@@ -32,7 +32,7 @@ function WaitingRoomView() {
         </TextTip>
       </InformationPanel>
 
-      {willExceed(untrustedCount) && (
+      {licenseOverused && (
         <div className="row">
           <div className="col-sm-12">
             <Alert color="warn">

app/react/edge/edge-devices/WaitingRoomView/queries.ts

@@ -5,6 +5,14 @@ import axios, { parseAxiosError } from '@/portainer/services/axios';
 import { promiseSequence } from '@/portainer/helpers/promise-utils';
 import { useIntegratedLicenseInfo } from '@/react/portainer/licenses/use-license.service';
 import { useEnvironmentList } from '@/react/portainer/environments/queries';
+import {
+  mutationOptions,
+  withError,
+  withInvalidate,
+} from '@/react-tools/react-query';
+import { queryKey as nodesCountQueryKey } from '@/react/portainer/system/useNodesCount';
+import { LicenseType } from '@/react/portainer/licenses/types';
+import { queryKeys } from '@/react/portainer/environments/queries/query-keys';
 
 export function useAssociateDeviceMutation() {
   const queryClient = useQueryClient();
@@ -12,17 +20,10 @@ export function useAssociateDeviceMutation() {
   return useMutation(
     (ids: EnvironmentId[]) =>
       promiseSequence(ids.map((id) => () => associateDevice(id))),
-    {
-      onSuccess: () => {
-        queryClient.invalidateQueries(['environments']);
-      },
-      meta: {
-        error: {
-          title: 'Failure',
-          message: 'Failed to associate devices',
-        },
-      },
-    }
+    mutationOptions(
+      withError('Failed to associate devices'),
+      withInvalidate(queryClient, [queryKeys.base(), nodesCountQueryKey])
+    )
   );
 }
 
@@ -34,19 +35,14 @@ async function associateDevice(environmentId: EnvironmentId) {
   }
 }
 
-export function useLicenseOverused() {
+export function useLicenseOverused(moreNodes: number) {
   const integratedInfo = useIntegratedLicenseInfo();
-  return {
-    willExceed,
-    isOverused: willExceed(0),
-  };
 
-  function willExceed(moreNodes: number) {
-    return (
-      !!integratedInfo &&
-      integratedInfo.usedNodes + moreNodes >= integratedInfo.licenseInfo.nodes
-    );
-  }
+  return (
+    !!integratedInfo &&
+    integratedInfo.licenseInfo.type === LicenseType.Essentials &&
+    integratedInfo.usedNodes + moreNodes > integratedInfo.licenseInfo.nodes
+  );
 }
 
 export function useUntrustedCount() {