[Fleet] bulk reassign kuery optimize (#134673)

* reassign kuery optimize

* fix test

* renamed to withoutManaged, added time measurement

* try catch to fix test

* unenroll improvement

* removed logging

* refactored to filter hosted agents in memory

* fixed tests

* removed withoutManaged

* added unit test

* revert plugin.ts changes

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

x-pack/plugins/fleet/server/services/agents/agent_service.test.ts

@@ -149,7 +149,9 @@ function expectApisToCallServicesSuccessfully(
     await expect(agentClient.listAgents({ showInactive: true })).resolves.toEqual(
       'getAgentsByKuery success'
     );
-    expect(mockGetAgentsByKuery).toHaveBeenCalledWith(mockEsClient, { showInactive: true });
+    expect(mockGetAgentsByKuery).toHaveBeenCalledWith(mockEsClient, {
+      showInactive: true,
+    });
   });
 
   test('client.getAgent calls getAgentById and returns results', async () => {

x-pack/plugins/fleet/server/services/agents/hosted_agent.test.ts

@@ -0,0 +1,63 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { savedObjectsClientMock } from '@kbn/core/server/mocks';
+
+import type { Agent } from '../../types';
+
+import { getHostedPolicies, isHostedAgent } from './hosted_agent';
+
+jest.mock('../agent_policy', () => {
+  return {
+    agentPolicyService: {
+      getByIDs: jest.fn().mockResolvedValue([
+        { id: 'hosted-policy', is_managed: true },
+        { id: 'regular-policy', is_managed: false },
+      ]),
+    },
+  };
+});
+
+describe('hosted agent helpers', () => {
+  const soClientMock = savedObjectsClientMock.create();
+  const expectedHostedPolicies = {
+    'hosted-policy': true,
+    'regular-policy': false,
+  };
+
+  it('should query unique managed policies', async () => {
+    const result = await getHostedPolicies(soClientMock, [
+      { policy_id: 'hosted-policy' } as Agent,
+      { policy_id: 'hosted-policy' } as Agent,
+      { policy_id: 'regular-policy' } as Agent,
+      { policy_id: 'regular-policy' } as Agent,
+    ]);
+    expect(result).toEqual(expectedHostedPolicies);
+  });
+
+  it('should return true for hosted policy', () => {
+    const isHosted = isHostedAgent(expectedHostedPolicies, { policy_id: 'hosted-policy' } as Agent);
+    expect(isHosted).toBeTruthy();
+  });
+
+  it('should return false for regular policy', () => {
+    const isHosted = isHostedAgent(expectedHostedPolicies, {
+      policy_id: 'regular-policy',
+    } as Agent);
+    expect(isHosted).toBeFalsy();
+  });
+
+  it('should return false for missing policy_id', () => {
+    const isHosted = isHostedAgent(expectedHostedPolicies, {} as Agent);
+    expect(isHosted).toBeFalsy();
+  });
+
+  it('should return false for non existing policy', () => {
+    const isHosted = isHostedAgent(expectedHostedPolicies, { policy_id: 'dummy-policy' } as Agent);
+    expect(isHosted).toBeFalsy();
+  });
+});

x-pack/plugins/fleet/server/services/agents/hosted_agent.ts

@@ -0,0 +1,36 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { SavedObjectsClientContract } from '@kbn/core/server';
+
+import type { Agent } from '../../types';
+import { agentPolicyService } from '../agent_policy';
+
+export async function getHostedPolicies(
+  soClient: SavedObjectsClientContract,
+  agents: Agent[]
+): Promise<{ [key: string]: boolean }> {
+  // get any policy ids from upgradable agents
+  const policyIdsToGet = new Set(
+    agents.filter((agent) => agent.policy_id).map((agent) => agent.policy_id!)
+  );
+
+  // get the agent policies for those ids
+  const agentPolicies = await agentPolicyService.getByIDs(soClient, Array.from(policyIdsToGet), {
+    fields: ['is_managed'],
+  });
+  const hostedPolicies = agentPolicies.reduce<Record<string, boolean>>((acc, policy) => {
+    acc[policy.id] = policy.is_managed;
+    return acc;
+  }, {});
+
+  return hostedPolicies;
+}
+
+export function isHostedAgent(hostedPolicies: { [key: string]: boolean }, agent: Agent) {
+  return agent.policy_id && hostedPolicies[agent.policy_id];
+}

x-pack/plugins/fleet/server/services/agents/reassign.ts

@@ -22,6 +22,7 @@ import {
 import type { GetAgentsOptions } from '.';
 import { createAgentAction } from './actions';
 import { searchHitToAgent } from './helpers';
+import { getHostedPolicies, isHostedAgent } from './hosted_agent';
 
 export async function reassignAgent(
   soClient: SavedObjectsClientContract,
@@ -80,10 +81,15 @@ export async function reassignAgents(
   options: ({ agents: Agent[] } | GetAgentsOptions) & { force?: boolean },
   newAgentPolicyId: string
 ): Promise<{ items: BulkActionResult[] }> {
-  const agentPolicy = await agentPolicyService.get(soClient, newAgentPolicyId);
-  if (!agentPolicy) {
+  const newAgentPolicy = await agentPolicyService.get(soClient, newAgentPolicyId);
+  if (!newAgentPolicy) {
     throw Boom.notFound(`Agent policy not found: ${newAgentPolicyId}`);
   }
+  if (newAgentPolicy.is_managed) {
+    throw new HostedAgentPolicyRestrictionRelatedError(
+      `Cannot reassign an agent to hosted agent policy ${newAgentPolicy.id}`
+    );
+  }
 
   const outgoingErrors: Record<Agent['id'], Error> = {};
   let givenAgents: Agent[] = [];
@@ -106,6 +112,8 @@ export async function reassignAgents(
   const givenOrder =
     'agentIds' in options ? options.agentIds : givenAgents.map((agent) => agent.id);
 
+  const hostedPolicies = await getHostedPolicies(soClient, givenAgents);
+
   // which are allowed to unenroll
   const agentResults = await Promise.allSettled(
     givenAgents.map(async (agent, index) => {
@@ -113,16 +121,13 @@ export async function reassignAgents(
         throw new AgentReassignmentError(`${agent.id} is already assigned to ${newAgentPolicyId}`);
       }
 
-      const isAllowed = await reassignAgentIsAllowed(
-        soClient,
-        esClient,
-        agent.id,
-        newAgentPolicyId
-      );
-      if (isAllowed) {
-        return agent;
+      if (isHostedAgent(hostedPolicies, agent)) {
+        throw new HostedAgentPolicyRestrictionRelatedError(
+          `Cannot reassign an agent from hosted agent policy ${agent.policy_id}`
+        );
       }
-      throw new AgentReassignmentError(`${agent.id} may not be reassigned to ${newAgentPolicyId}`);
+
+      return agent;
     })
   );
 

x-pack/plugins/fleet/server/services/agents/unenroll.ts

@@ -20,6 +20,7 @@ import {
   getAgentPolicyForAgent,
   bulkUpdateAgents,
 } from './crud';
+import { getHostedPolicies, isHostedAgent } from './hosted_agent';
 
 async function unenrollAgentIsAllowed(
   soClient: SavedObjectsClientContract,
@@ -80,21 +81,22 @@ export async function unenrollAgents(
     }
     return !agent.unenrollment_started_at && !agent.unenrolled_at;
   });
-  // And which are allowed to unenroll
-  const agentResults = await Promise.allSettled(
-    agentsEnrolled.map((agent) =>
-      unenrollAgentIsAllowed(soClient, esClient, agent.id).then((_) => agent)
-    )
-  );
+
+  const hostedPolicies = await getHostedPolicies(soClient, agentsEnrolled);
+
   const outgoingErrors: Record<Agent['id'], Error> = {};
+
+  // And which are allowed to unenroll
   const agentsToUpdate = options.force
     ? agentsEnrolled
-    : agentResults.reduce<Agent[]>((agents, result, index) => {
-        if (result.status === 'fulfilled') {
-          agents.push(result.value);
-        } else {
+    : agentsEnrolled.reduce<Agent[]>((agents, agent, index) => {
+        if (isHostedAgent(hostedPolicies, agent)) {
           const id = givenAgents[index].id;
-          outgoingErrors[id] = result.reason;
+          outgoingErrors[id] = new HostedAgentPolicyRestrictionRelatedError(
+            `Cannot unenroll ${agent.id} from a hosted agent policy ${agent.policy_id}`
+          );
+        } else {
+          agents.push(agent);
         }
         return agents;
       }, []);

x-pack/plugins/fleet/server/services/agents/upgrade.ts

@@ -10,7 +10,6 @@ import moment from 'moment';
 import pMap from 'p-map';
 
 import type { Agent, BulkActionResult, FleetServerAgentAction, CurrentUpgrade } from '../../types';
-import { agentPolicyService } from '..';
 import {
   AgentReassignmentError,
   HostedAgentPolicyRestrictionRelatedError,
@@ -30,6 +29,7 @@ import {
   getAgentPolicyForAgent,
 } from './crud';
 import { searchHitToAgent } from './helpers';
+import { getHostedPolicies, isHostedAgent } from './hosted_agent';
 
 const MINIMUM_EXECUTION_DURATION_SECONDS = 1800; // 30m
 
@@ -107,26 +107,13 @@ export async function sendUpgradeAgentsActions(
     givenAgents = await getAgents(esClient, options);
   }
 
-  // get any policy ids from upgradable agents
-  const policyIdsToGet = new Set(
-    givenAgents.filter((agent) => agent.policy_id).map((agent) => agent.policy_id!)
-  );
-
-  // get the agent policies for those ids
-  const agentPolicies = await agentPolicyService.getByIDs(soClient, Array.from(policyIdsToGet), {
-    fields: ['is_managed'],
-  });
-  const hostedPolicies = agentPolicies.reduce<Record<string, boolean>>((acc, policy) => {
-    acc[policy.id] = policy.is_managed;
-    return acc;
-  }, {});
-  const isHostedAgent = (agent: Agent) => agent.policy_id && hostedPolicies[agent.policy_id];
+  const hostedPolicies = await getHostedPolicies(soClient, givenAgents);
 
   // results from getAgents with options.kuery '' (or even 'active:false') may include hosted agents
   // filter them out unless options.force
   const agentsToCheckUpgradeable =
     'kuery' in options && !options.force
-      ? givenAgents.filter((agent: Agent) => !isHostedAgent(agent))
+      ? givenAgents.filter((agent: Agent) => !isHostedAgent(hostedPolicies, agent))
       : givenAgents;
 
   const kibanaVersion = appContextService.getKibanaVersion();
@@ -139,7 +126,7 @@ export async function sendUpgradeAgentsActions(
         throw new IngestManagerError(`${agent.id} is not upgradeable`);
       }
 
-      if (!options.force && isHostedAgent(agent)) {
+      if (!options.force && isHostedAgent(hostedPolicies, agent)) {
         throw new HostedAgentPolicyRestrictionRelatedError(
           `Cannot upgrade agent in hosted agent policy ${agent.policy_id}`
         );

x-pack/test/fleet_api_integration/apis/agents/reassign.ts

@@ -16,7 +16,7 @@ export default function (providerContext: FtrProviderContext) {
   const supertest = getService('supertest');
   const supertestWithoutAuth = getService('supertestWithoutAuth');
 
-  describe('reassign agent(s)', () => {
+  describe('fleet_reassign_agent', () => {
     before(async () => {
       await esArchiver.load('x-pack/test/functional/es_archives/fleet/empty_fleet_server');
     });
@@ -190,6 +190,7 @@ export default function (providerContext: FtrProviderContext) {
             policy_id: 'policy2',
           })
           .expect(200);
+
         const { body } = await supertest.get(`/api/fleet/agents`).set('kbn-xsrf', 'xxx');
         expect(body.total).to.eql(4);
         body.items.forEach((agent: any) => {