[Fleet] Add test/fix for invalid/missing ids in bulk agent reassign

x-pack/plugins/fleet/server/plugin.ts

@@ -80,8 +80,8 @@ import {
 import {
   getAgentStatusById,
   authenticateAgentWithAccessToken,
-  listAgents,
-  getAgent,
+  getAgentsByKuery,
+  getAgentById,
 } from './services/agents';
 import { agentCheckinState } from './services/agents/checkin/state';
 import { registerFleetUsageCollector } from './collectors/register';
@@ -322,8 +322,8 @@ export class FleetPlugin
         },
       },
       agentService: {
-        getAgent,
-        listAgents,
+        getAgent: getAgentById,
+        listAgents: getAgentsByKuery,
         getAgentStatusById,
         authenticateAgentWithAccessToken,
       },

x-pack/plugins/fleet/server/routes/agent/handlers.ts

@@ -44,8 +44,7 @@ export const getAgentHandler: RequestHandler<
   const esClient = context.core.elasticsearch.client.asCurrentUser;
 
   try {
-    const agent = await AgentService.getAgent(esClient, request.params.agentId);
-
+    const agent = await AgentService.getAgentById(esClient, request.params.agentId);
     const body: GetOneAgentResponse = {
       item: {
         ...agent,
@@ -134,8 +133,7 @@ export const updateAgentHandler: RequestHandler<
     await AgentService.updateAgent(esClient, request.params.agentId, {
       user_provided_metadata: request.body.user_provided_metadata,
     });
-    const agent = await AgentService.getAgent(esClient, request.params.agentId);
-
+    const agent = await AgentService.getAgentById(esClient, request.params.agentId);
     const body = {
       item: {
         ...agent,
@@ -245,7 +243,7 @@ export const getAgentsHandler: RequestHandler<
   const esClient = context.core.elasticsearch.client.asCurrentUser;
 
   try {
-    const { agents, total, page, perPage } = await AgentService.listAgents(esClient, {
+    const { agents, total, page, perPage } = await AgentService.getAgentsByKuery(esClient, {
       page: request.query.page,
       perPage: request.query.perPage,
       showInactive: request.query.showInactive,
@@ -310,6 +308,7 @@ export const postBulkAgentsReassignHandler: RequestHandler<
 
   const soClient = context.core.savedObjects.client;
   const esClient = context.core.elasticsearch.client.asInternalUser;
+
   try {
     const results = await AgentService.reassignAgents(
       soClient,

x-pack/plugins/fleet/server/routes/agent/index.ts

@@ -125,7 +125,7 @@ export const registerAPIRoutes = (router: IRouter, config: FleetConfigType) => {
       options: { tags: [`access:${PLUGIN_ID}-all`] },
     },
     postNewAgentActionHandlerBuilder({
-      getAgent: AgentService.getAgent,
+      getAgent: AgentService.getAgentById,
       createAgentAction: AgentService.createAgentAction,
     })
   );

x-pack/plugins/fleet/server/routes/agent/upgrade_handler.ts

@@ -15,7 +15,7 @@ import * as AgentService from '../../services/agents';
 import { appContextService } from '../../services';
 import { defaultIngestErrorHandler } from '../../errors';
 import { isAgentUpgradeable } from '../../../common/services';
-import { getAgent } from '../../services/agents';
+import { getAgentById } from '../../services/agents';
 
 export const postAgentUpgradeHandler: RequestHandler<
   TypeOf<typeof PostAgentUpgradeRequestSchema.params>,
@@ -36,7 +36,7 @@ export const postAgentUpgradeHandler: RequestHandler<
       },
     });
   }
-  const agent = await getAgent(esClient, request.params.agentId);
+  const agent = await getAgentById(esClient, request.params.agentId);
   if (agent.unenrollment_started_at || agent.unenrolled_at) {
     return response.customError({
       statusCode: 400,

x-pack/plugins/fleet/server/routes/agent_policy/handlers.ts

@@ -11,7 +11,7 @@ import bluebird from 'bluebird';
 
 import { fullAgentPolicyToYaml } from '../../../common/services';
 import { appContextService, agentPolicyService, packagePolicyService } from '../../services';
-import { listAgents } from '../../services/agents';
+import { getAgentsByKuery } from '../../services/agents';
 import { AGENT_SAVED_OBJECT_TYPE } from '../../constants';
 import type {
   GetAgentPoliciesRequestSchema,
@@ -58,7 +58,7 @@ export const getAgentPoliciesHandler: RequestHandler<
     await bluebird.map(
       items,
       (agentPolicy: GetAgentPoliciesResponseItem) =>
-        listAgents(esClient, {
+        getAgentsByKuery(esClient, {
           showInactive: false,
           perPage: 0,
           page: 1,

x-pack/plugins/fleet/server/services/agent_policy.ts

@@ -43,7 +43,7 @@ import {
 } from '../errors';
 import { getFullAgentPolicyKibanaConfig } from '../../common/services/full_agent_policy_kibana_config';
 
-import { createAgentPolicyAction, listAgents } from './agents';
+import { createAgentPolicyAction, getAgentsByKuery } from './agents';
 import { packagePolicyService } from './package_policy';
 import { outputService } from './output';
 import { agentPolicyUpdateEventHandler } from './agent_policy_update';
@@ -520,7 +520,7 @@ class AgentPolicyService {
       throw new Error('The default agent policy cannot be deleted');
     }
 
-    const { total } = await listAgents(esClient, {
+    const { total } = await getAgentsByKuery(esClient, {
       showInactive: false,
       perPage: 0,
       page: 1,

x-pack/plugins/fleet/server/services/agents/checkin/state_new_actions.ts

@@ -39,7 +39,7 @@ import {
   getAgentPolicyActionByIds,
 } from '../actions';
 import { appContextService } from '../../app_context';
-import { getAgent, updateAgent } from '../crud';
+import { getAgentById, updateAgent } from '../crud';
 
 import { toPromiseAbortable, AbortError, createRateLimiter } from './rxjs_utils';
 
@@ -266,7 +266,7 @@ export function agentCheckinStateNewActionsFactory() {
           (action) => action.type === 'INTERNAL_POLICY_REASSIGN'
         );
         if (hasConfigReassign) {
-          return from(getAgent(esClient, agent.id)).pipe(
+          return from(getAgentById(esClient, agent.id)).pipe(
             concatMap((refreshedAgent) => {
               if (!refreshedAgent.policy_id) {
                 throw new Error('Agent does not have a policy assigned');

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

@@ -6,15 +6,14 @@
  */
 
 import Boom from '@hapi/boom';
-import type { SearchResponse } from 'elasticsearch';
+import type { SearchResponse, MGetResponse, GetResponse } from 'elasticsearch';
 import type { SavedObjectsClientContract, ElasticsearchClient } from 'src/core/server';
 
 import type { AgentSOAttributes, Agent, ListWithKuery } from '../../types';
 import { appContextService, agentPolicyService } from '../../services';
 import type { FleetServerAgent } from '../../../common';
 import { isAgentUpgradeable, SO_SEARCH_LIMIT } from '../../../common';
 import { AGENT_SAVED_OBJECT_TYPE, AGENTS_INDEX } from '../../constants';
-import type { ESSearchHit } from '../../../../../../typings/elasticsearch';
 import { escapeSearchQueryPhrase, normalizeKuery } from '../saved_object';
 import type { KueryNode } from '../../../../../../src/plugins/data/server';
 import { esKuery } from '../../../../../../src/plugins/data/server';
@@ -59,7 +58,35 @@ export function removeSOAttributes(kuery: string) {
   return kuery.replace(/attributes\./g, '').replace(/fleet-agents\./g, '');
 }
 
-export async function listAgents(
+export type GetAgentsOptions =
+  | {
+      agentIds: string[];
+    }
+  | {
+      kuery: string;
+      showInactive?: boolean;
+    };
+
+export async function getAgents(esClient: ElasticsearchClient, options: GetAgentsOptions) {
+  let initialResults = [];
+
+  if ('agentIds' in options) {
+    initialResults = await getAgentsById(esClient, options.agentIds);
+  } else if ('kuery' in options) {
+    initialResults = (
+      await getAllAgentsByKuery(esClient, {
+        kuery: options.kuery,
+        showInactive: options.showInactive ?? false,
+      })
+    ).agents;
+  } else {
+    throw new IngestManagerError('Cannot get agents');
+  }
+
+  return initialResults;
+}
+
+export async function getAgentsByKuery(
   esClient: ElasticsearchClient,
   options: ListWithKuery & {
     showInactive: boolean;
@@ -91,8 +118,7 @@ export async function listAgents(
 
   const kueryNode = _joinFilters(filters);
   const body = kueryNode ? { query: esKuery.toElasticsearchQuery(kueryNode) } : {};
-
-  const res = await esClient.search({
+  const res = await esClient.search<SearchResponse<FleetServerAgent>>({
     index: AGENTS_INDEX,
     from: (page - 1) * perPage,
     size: perPage,
@@ -101,27 +127,24 @@ export async function listAgents(
     body,
   });
 
-  let agentResults: Agent[] = res.body.hits.hits.map(searchHitToAgent);
-  let total = res.body.hits.total.value;
-
+  let agents = res.body.hits.hits.map(searchHitToAgent);
   // filtering for a range on the version string will not work,
   // nor does filtering on a flattened field (local_metadata), so filter here
   if (showUpgradeable) {
-    agentResults = agentResults.filter((agent) =>
+    agents = agents.filter((agent) =>
       isAgentUpgradeable(agent, appContextService.getKibanaVersion())
     );
-    total = agentResults.length;
   }
 
   return {
-    agents: res.body.hits.hits.map(searchHitToAgent),
-    total,
+    agents,
+    total: agents.length,
     page,
     perPage,
   };
 }
 
-export async function listAllAgents(
+export async function getAllAgentsByKuery(
   esClient: ElasticsearchClient,
   options: Omit<ListWithKuery, 'page' | 'perPage'> & {
     showInactive: boolean;
@@ -130,7 +153,7 @@ export async function listAllAgents(
   agents: Agent[];
   total: number;
 }> {
-  const res = await listAgents(esClient, { ...options, page: 1, perPage: SO_SEARCH_LIMIT });
+  const res = await getAgentsByKuery(esClient, { ...options, page: 1, perPage: SO_SEARCH_LIMIT });
 
   return {
     agents: res.agents,
@@ -161,34 +184,51 @@ export async function countInactiveAgents(
   return res.body.hits.total.value;
 }
 
-export async function getAgent(esClient: ElasticsearchClient, agentId: string) {
+export async function getAgentById(esClient: ElasticsearchClient, agentId: string) {
+  const agentNotFoundError = new AgentNotFoundError(`Agent ${agentId} not found`);
   try {
-    const agentHit = await esClient.get<ESSearchHit<FleetServerAgent>>({
+    const agentHit = await esClient.get<GetResponse<FleetServerAgent>>({
       index: AGENTS_INDEX,
       id: agentId,
     });
+
+    if (agentHit.body.found === false) {
+      throw agentNotFoundError;
+    }
     const agent = searchHitToAgent(agentHit.body);
 
     return agent;
   } catch (err) {
     if (isESClientError(err) && err.meta.statusCode === 404) {
-      throw new AgentNotFoundError(`Agent ${agentId} not found`);
+      throw agentNotFoundError;
     }
     throw err;
   }
 }
 
-export async function getAgents(
+async function getAgentDocuments(
   esClient: ElasticsearchClient,
   agentIds: string[]
-): Promise<Agent[]> {
-  const body = { docs: agentIds.map((_id) => ({ _id })) };
-
-  const res = await esClient.mget({
-    body,
+): Promise<Array<GetResponse<FleetServerAgent>>> {
+  const res = await esClient.mget<MGetResponse<FleetServerAgent>>({
     index: AGENTS_INDEX,
+    body: { docs: agentIds.map((_id) => ({ _id })) },
   });
-  const agents = res.body.docs.map(searchHitToAgent);
+
+  return res.body.docs || [];
+}
+
+export async function getAgentsById(
+  esClient: ElasticsearchClient,
+  agentIds: string[],
+  options: { includeMissing?: boolean } = { includeMissing: false }
+): Promise<Agent[]> {
+  const allDocs = await getAgentDocuments(esClient, agentIds);
+  const agentDocs = options.includeMissing
+    ? allDocs
+    : allDocs.filter((res) => res._id && res._source);
+  const agents = agentDocs.map((doc) => searchHitToAgent(doc));
+
   return agents;
 }
 
@@ -201,7 +241,7 @@ export async function getAgentByAccessAPIKeyId(
     q: `access_api_key_id:${escapeSearchQueryPhrase(accessAPIKeyId)}`,
   });
 
-  const [agent] = res.body.hits.hits.map(searchHitToAgent);
+  const agent = searchHitToAgent(res.body.hits.hits[0]);
 
   if (!agent) {
     throw new AgentNotFoundError('Agent not found');
@@ -288,7 +328,7 @@ export async function getAgentPolicyForAgent(
   esClient: ElasticsearchClient,
   agentId: string
 ) {
-  const agent = await getAgent(esClient, agentId);
+  const agent = await getAgentById(esClient, agentId);
   if (!agent.policy_id) {
     return;
   }

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

@@ -5,10 +5,15 @@
  * 2.0.
  */
 
-import type { ESSearchHit } from '../../../../../../typings/elasticsearch';
+import type { GetResponse, SearchResponse } from 'elasticsearch';
+
 import type { Agent, AgentSOAttributes, FleetServerAgent } from '../../types';
 
-export function searchHitToAgent(hit: ESSearchHit<FleetServerAgent>): Agent {
+type FleetServerAgentESResponse =
+  | GetResponse<FleetServerAgent>
+  | SearchResponse<FleetServerAgent>['hits']['hits'][0];
+
+export function searchHitToAgent(hit: FleetServerAgentESResponse): Agent {
   return {
     id: hit._id,
     ...hit._source,

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

@@ -121,7 +121,7 @@ function createClientsMock() {
       case unmanagedAgentPolicySO2.id:
         return unmanagedAgentPolicySO2;
       default:
-        throw new Error('Not found');
+        throw new Error(`${id} not found`);
     }
   });
   soClientMock.bulkGet.mockImplementation(async (options) => {
@@ -147,7 +147,7 @@ function createClientsMock() {
       case agentInUnmanagedDoc._id:
         return { body: agentInUnmanagedDoc };
       default:
-        throw new Error('Not found');
+        throw new Error(`${id} not found`);
     }
   });
   // @ts-expect-error