[Fleet] Fix changing space for agent actions (#203683)

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

@@ -7,6 +7,7 @@
 
 import { createAppContextStartContractMock } from '../../mocks';
 import { agentPolicyService } from '../agent_policy';
+import { getAgentsByKuery } from '../agents';
 import { appContextService } from '../app_context';
 import { packagePolicyService } from '../package_policy';
 
@@ -16,6 +17,7 @@ import { isSpaceAwarenessEnabled } from './helpers';
 jest.mock('./helpers');
 jest.mock('../agent_policy');
 jest.mock('../package_policy');
+jest.mock('../agents');
 
 describe('updateAgentPolicySpaces', () => {
   beforeEach(() => {
@@ -55,6 +57,10 @@ describe('updateAgentPolicySpaces', () => {
           } as any,
         ],
       });
+
+    jest.mocked(getAgentsByKuery).mockResolvedValue({
+      agents: [],
+    } as any);
   });
 
   it('does nothings if agent policy already in correct space', async () => {

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

@@ -6,27 +6,29 @@
  */
 
 import deepEqual from 'fast-deep-equal';
-
 import { DEFAULT_SPACE_ID } from '@kbn/spaces-plugin/common';
+import type { SortResults } from '@elastic/elasticsearch/lib/api/types';
 
 import {
   AGENTS_INDEX,
+  AGENT_ACTIONS_INDEX,
   AGENT_POLICY_SAVED_OBJECT_TYPE,
   PACKAGE_POLICY_SAVED_OBJECT_TYPE,
   SO_SEARCH_LIMIT,
   UNINSTALL_TOKENS_SAVED_OBJECT_TYPE,
 } from '../../../common/constants';
-
 import { appContextService } from '../app_context';
 import { agentPolicyService } from '../agent_policy';
 import { ENROLLMENT_API_KEYS_INDEX } from '../../constants';
 import { packagePolicyService } from '../package_policy';
 import { FleetError, HostedAgentPolicyRestrictionRelatedError } from '../../errors';
-
 import type { UninstallTokenSOAttributes } from '../security/uninstall_token_service';
+import { closePointInTime, getAgentsByKuery, openPointInTime } from '../agents';
 
 import { isSpaceAwarenessEnabled } from './helpers';
 
+const UPDATE_AGENT_BATCH_SIZE = 1000;
+
 export async function updateAgentPolicySpaces({
   agentPolicyId,
   currentSpaceId,
@@ -49,6 +51,7 @@ export async function updateAgentPolicySpaces({
   const soClient = appContextService.getInternalUserSOClientWithoutSpaceExtension();
 
   const currentSpaceSoClient = appContextService.getInternalUserSOClientForSpaceId(currentSpaceId);
+  const newSpaceSoClient = appContextService.getInternalUserSOClientForSpaceId(newSpaceIds[0]);
   const existingPolicy = await agentPolicyService.get(currentSpaceSoClient, agentPolicyId);
 
   const existingPackagePolicies = await packagePolicyService.findAllForAgentPolicy(
@@ -166,4 +169,55 @@ export async function updateAgentPolicySpaces({
     ignore_unavailable: true,
     refresh: true,
   });
+
+  const agentIndexExists = await esClient.indices.exists({
+    index: AGENTS_INDEX,
+  });
+
+  // Update agent actions
+  if (agentIndexExists) {
+    const pitId = await openPointInTime(esClient);
+
+    try {
+      let hasMore = true;
+      let searchAfter: SortResults | undefined;
+      while (hasMore) {
+        const { agents } = await getAgentsByKuery(esClient, newSpaceSoClient, {
+          kuery: `policy_id:"${agentPolicyId}"`,
+          showInactive: true,
+          perPage: UPDATE_AGENT_BATCH_SIZE,
+          pitId,
+          searchAfter,
+        });
+
+        if (agents.length === 0) {
+          hasMore = false;
+          break;
+        }
+
+        const lastAgent = agents[agents.length - 1];
+        searchAfter = lastAgent.sort;
+
+        await esClient.updateByQuery({
+          index: AGENT_ACTIONS_INDEX,
+          query: {
+            bool: {
+              must: {
+                terms: {
+                  agents: agents.map(({ id }) => id),
+                },
+              },
+            },
+          },
+          script: `ctx._source.namespaces = [${newSpaceIds
+            .map((spaceId) => `"${spaceId}"`)
+            .join(',')}]`,
+          ignore_unavailable: true,
+          refresh: true,
+        });
+      }
+    } finally {
+      await closePointInTime(esClient, pitId);
+    }
+  }
 }

x-pack/test/fleet_api_integration/apis/space_awareness/change_space_agent_policies.ts

@@ -6,7 +6,10 @@
  */
 
 import expect from '@kbn/expect';
+import { v4 as uuidV4 } from 'uuid';
+import { Client } from '@elastic/elasticsearch';
 import { CreateAgentPolicyResponse, GetOnePackagePolicyResponse } from '@kbn/fleet-plugin/common';
+import { FleetServerAgentAction } from '@kbn/fleet-plugin/common/types';
 import { FtrProviderContext } from '../../../api_integration/ftr_provider_context';
 import { skipIfNoDockerRegistry } from '../../helpers';
 import { SpaceTestApiClient } from './api_helper';
@@ -19,6 +22,33 @@ import {
 } from './helpers';
 import { testUsers, setupTestUsers } from '../test_users';
 
+export async function createFleetAction(esClient: Client, agentId: string, spaceId?: string) {
+  const actionResponse = await esClient.index({
+    index: '.fleet-actions',
+    refresh: 'wait_for',
+    body: {
+      '@timestamp': new Date().toISOString(),
+      expiration: new Date().toISOString(),
+      agents: [agentId],
+      action_id: uuidV4(),
+      data: {},
+      type: 'UPGRADE',
+      namespaces: spaceId ? [spaceId] : undefined,
+    },
+  });
+
+  return actionResponse._id;
+}
+
+async function getFleetActionDoc(esClient: Client, actionId: string) {
+  const actionResponse = await esClient.get<FleetServerAgentAction>({
+    index: '.fleet-actions',
+    id: actionId,
+  });
+
+  return actionResponse;
+}
+
 export default function (providerContext: FtrProviderContext) {
   const { getService } = providerContext;
   const supertest = getService('supertest');
@@ -39,6 +69,9 @@ export default function (providerContext: FtrProviderContext) {
     let policy1AgentId: string;
     let policy2AgentId: string;
 
+    let agent1ActionId: string;
+    let agent2ActionId: string;
+
     before(async () => {
       TEST_SPACE_1 = spaces.getDefaultTestSpace();
       await setupTestUsers(getService('security'), true);
@@ -63,6 +96,9 @@ export default function (providerContext: FtrProviderContext) {
       policy1AgentId = await createFleetAgent(esClient, defaultSpacePolicy1.item.id);
       policy2AgentId = await createFleetAgent(esClient, defaultSpacePolicy2.item.id);
 
+      agent1ActionId = await createFleetAction(esClient, policy1AgentId, 'default');
+      agent2ActionId = await createFleetAction(esClient, policy2AgentId, 'default');
+
       const packagePolicyRes = await apiClient.createPackagePolicy(undefined, {
         policy_ids: [defaultSpacePolicy1.item.id],
         name: `test-nginx-${Date.now()}`,
@@ -172,6 +208,15 @@ export default function (providerContext: FtrProviderContext) {
         }
       }
 
+      async function assertActionSpaces(actionId: string, expectedSpaces: string[]) {
+        const actionDoc = await getFleetActionDoc(esClient, actionId);
+        if (expectedSpaces.length === 1 && expectedSpaces[0] === 'default') {
+          expect(actionDoc._source?.namespaces ?? ['default']).to.eql(expectedSpaces);
+        } else {
+          expect(actionDoc._source?.namespaces).to.eql(expectedSpaces);
+        }
+      }
+
       it('should allow set policy in multiple space', async () => {
         await apiClient.putAgentPolicy(defaultSpacePolicy1.item.id, {
           name: 'tata',
@@ -189,6 +234,9 @@ export default function (providerContext: FtrProviderContext) {
         await assertAgentSpaces(policy1AgentId, ['default', TEST_SPACE_1]);
         await assertAgentSpaces(policy2AgentId, ['default']);
 
+        await assertActionSpaces(agent1ActionId, ['default', TEST_SPACE_1]);
+        await assertActionSpaces(agent2ActionId, ['default']);
+
         await assertEnrollemntApiKeysForSpace('default', [
           defaultSpacePolicy1.item.id,
           defaultSpacePolicy2.item.id,
@@ -213,6 +261,10 @@ export default function (providerContext: FtrProviderContext) {
         await assertPackagePolicyNotAvailableInSpace();
         await assertAgentSpaces(policy1AgentId, [TEST_SPACE_1]);
         await assertAgentSpaces(policy2AgentId, ['default']);
+
+        await assertActionSpaces(agent1ActionId, [TEST_SPACE_1]);
+        await assertActionSpaces(agent2ActionId, ['default']);
+
         await assertEnrollemntApiKeysForSpace('default', [defaultSpacePolicy2.item.id]);
         await assertEnrollemntApiKeysForSpace(TEST_SPACE_1, [defaultSpacePolicy1.item.id]);
         // Ensure no side effect on other policies

x-pack/test/fleet_api_integration/apis/space_awareness/helpers.ts

@@ -49,6 +49,12 @@ export async function cleanFleetIndices(esClient: Client) {
       ignore_unavailable: true,
       refresh: true,
     }),
+    esClient.deleteByQuery({
+      index: AGENT_ACTIONS_INDEX,
+      q: '*',
+      ignore_unavailable: true,
+      refresh: true,
+    }),
   ]);
 }