clean up + reformating option rendering

x-pack/plugins/osquery/public/agents/agents_table.tsx

@@ -28,7 +28,7 @@ import {
   generateSelectedAgentsMessage,
 } from './translations';
 
-import { AGENT_GROUP_KEY, SelectedGroups, AgentOptionValue, GroupOptionValue } from './types';
+import { AGENT_GROUP_KEY, SelectedGroups, AgentOptionValue, GroupOptionValue, Group } from './types';
 
 export interface AgentsSelection {
   agents: string[];
@@ -46,6 +46,17 @@ type GroupOption = EuiComboBoxOptionOption<AgentOptionValue | GroupOptionValue>;
 
 const getColor = generateColorPicker();
 
+const generateOptions = (groupType: AGENT_GROUP_KEY, label: string, collection: Group[]) => {
+  return {
+    label,
+    options: collection.map(({ name, id, size }) => ({
+      label: name,
+      color: getColor(groupType),
+      value: { groupType, id, size },
+    })),
+  };
+}
+
 const AgentsTableComponent: React.FC<AgentsTableProps> = ({ onChange }) => {
   const osqueryPolicyData = useOsqueryPolicies();
   const { loading: groupsLoading, totalCount: totalNumAgents, groups } = useAgentGroups(
@@ -74,26 +85,12 @@ const AgentsTableComponent: React.FC<AgentsTableProps> = ({ onChange }) => {
 
     if (groups.platforms.length > 0) {
       const groupType = AGENT_GROUP_KEY.Platform;
-      opts.push({
-        label: AGENT_PLATFORMS_LABEL,
-        options: groups.platforms.map(({ name, size }) => ({
-          label: name,
-          color: getColor(groupType),
-          value: { groupType, size },
-        })),
-      });
+      opts.push(generateOptions(groupType, AGENT_PLATFORMS_LABEL, groups.platforms))
     }
 
     if (groups.policies.length > 0) {
       const groupType = AGENT_GROUP_KEY.Policy;
-      opts.push({
-        label: AGENT_POLICY_LABEL,
-        options: groups.policies.map(({ name, size }) => ({
-          label: name,
-          color: getColor(groupType),
-          value: { groupType, size },
-        })),
-      });
+      opts.push(generateOptions(groupType, AGENT_POLICY_LABEL, groups.policies))
     }
 
     if (agents && agents.length > 0) {
@@ -102,6 +99,7 @@ const AgentsTableComponent: React.FC<AgentsTableProps> = ({ onChange }) => {
         label: AGENT_SELECTION_LABEL,
         options: (agents as Agent[]).map((agent: Agent) => ({
           label: agent.local_metadata.host.hostname,
+          key: agent.local_metadata.elastic.agent.id,
           color: getColor(groupType),
           value: {
             groupType,
@@ -126,7 +124,7 @@ const AgentsTableComponent: React.FC<AgentsTableProps> = ({ onChange }) => {
         policiesSelected: [],
       };
       // parse through the selections to be able to determine how many are actually selected
-      const selectedAgents = [];
+      const selectedAgents: AgentOptionValue[] = [];
       const selectedGroups: SelectedGroups = {
         policy: {},
         platform: {},
@@ -144,26 +142,27 @@ const AgentsTableComponent: React.FC<AgentsTableProps> = ({ onChange }) => {
             value = opt.value as GroupOptionValue;
             if (!newAgentSelection.allAgentsSelected) {
               // we don't need to calculate diffs when all agents are selected
-              selectedGroups.platform[opt.label] = value.size;
+              selectedGroups.platform[opt.value?.id ?? opt.label] = value.size;
             }
             newAgentSelection.platformsSelected.push(opt.label);
             break;
           case AGENT_GROUP_KEY.Policy:
             value = opt.value as GroupOptionValue;
             if (!newAgentSelection.allAgentsSelected) {
               // we don't need to calculate diffs when all agents are selected
-              selectedGroups.policy[opt.label] = value.size ?? 0;
+              selectedGroups.policy[opt.value?.id ?? opt.label] = value.size;
             }
             newAgentSelection.policiesSelected.push(opt.label);
             break;
           case AGENT_GROUP_KEY.Agent:
             value = opt.value as AgentOptionValue;
             if (!newAgentSelection.allAgentsSelected) {
               // we don't need to count how many agents are selected if they are all selected
-              selectedAgents.push(opt.value);
+              selectedAgents.push(value);
+            }
+            if (value?.id) {
+              newAgentSelection.agents.push(value.id);
             }
-            // TODO: fix this casting by updating the opt type to be a union
-            newAgentSelection.agents.push(value.id as string);
             break;
           default:
             // this should never happen!
@@ -177,7 +176,7 @@ const AgentsTableComponent: React.FC<AgentsTableProps> = ({ onChange }) => {
         const checkAgent = generateAgentCheck(selectedGroups);
         setNumAgentsSelected(
           // filter out all the agents counted by selected policies and platforms
-          selectedAgents.filter((a) => checkAgent(a as AgentOptionValue)).length +
+          selectedAgents.filter(checkAgent).length +
             // add the number of agents added via policy and platform groups
             getNumAgentsInGrouping(selectedGroups) -
             // subtract the number of agents double counted by policy/platform selections
@@ -191,18 +190,22 @@ const AgentsTableComponent: React.FC<AgentsTableProps> = ({ onChange }) => {
   );
 
   const renderOption = useCallback((option, searchValue, contentClassName) => {
-    const { label, value } = option;
+    const { label, value, key } = option;
     return value?.groupType === AGENT_GROUP_KEY.Agent ? (
       <EuiHealth color={value?.online ? 'success' : 'danger'}>
         <span className={contentClassName}>
           <EuiHighlight search={searchValue}>{label}</EuiHighlight>
+          &nbsp;
+          <span>({key})</span>
         </span>
       </EuiHealth>
     ) : (
       <span className={contentClassName}>
+        <span>[{value?.size}]</span>
+        &nbsp;
         <EuiHighlight search={searchValue}>{label}</EuiHighlight>
         &nbsp;
-        <span>({value?.size})</span>
+        {value?.id && label !== value?.id && (<span>({value?.id})</span>)}
       </span>
     );
   }, []);

x-pack/plugins/osquery/public/agents/helpers.ts

@@ -43,11 +43,11 @@ export const processAggregations = (aggs: Record<string, Aggregate>) => {
   const platformTerms = aggs.platforms as TermsAggregate<AggregationDataPoint>;
   const policyTerms = aggs.policies as TermsAggregate<AggregationDataPoint>;
 
-  const policies = policyTerms?.buckets.map((o) => ({ name: o.key, size: o.doc_count })) ?? [];
+  const policies = policyTerms?.buckets.map((o) => ({ name: o.key, id: o.key, size: o.doc_count })) ?? [];
 
   if (platformTerms?.buckets) {
     for (const { key, doc_count: size, policies: platformPolicies } of platformTerms.buckets) {
-      platforms.push({ name: key, size });
+      platforms.push({ name: key, id: key, size });
       if (platformPolicies?.buckets && policies.length > 0) {
         overlap[key] = platformPolicies.buckets.reduce((acc: { [key: string]: number }, pol) => {
           acc[pol.key] = pol.doc_count;

x-pack/plugins/osquery/public/agents/types.ts

@@ -17,6 +17,7 @@ export type AggregationDataPoint = BaseDataPoint & {
 };
 
 export interface Group {
+  id: string;
   name: string;
   size: number;
 }
@@ -29,13 +30,13 @@ export interface SelectedGroups {
 }
 
 interface BaseGroupOption {
+  id?: string;
   groupType: AGENT_GROUP_KEY;
 }
 
 export type AgentOptionValue = BaseGroupOption & {
   groups: { [groupType: string]: string };
   online: boolean;
-  id: string;
 };
 
 export type GroupOptionValue = BaseGroupOption & {

x-pack/plugins/osquery/public/agents/use_agent_groups.ts

@@ -7,6 +7,7 @@
 import { useState } from 'react';
 import { useQuery } from 'react-query';
 import { useKibana } from '../common/lib/kibana';
+import { useAgentPolicies } from './use_agent_policies';
 
 import {
   OsqueryQueries,
@@ -25,6 +26,7 @@ interface UseAgentGroups {
 export const useAgentGroups = ({ osqueryPolicies, osqueryPoliciesLoading }: UseAgentGroups) => {
   const { data } = useKibana().services;
 
+  const { agentPoliciesLoading, agentPolicyById } = useAgentPolicies(osqueryPolicies)
   const [platforms, setPlatforms] = useState<Group[]>([]);
   const [policies, setPolicies] = useState<Group[]>([]);
   const [loading, setLoading] = useState(true);
@@ -76,16 +78,22 @@ export const useAgentGroups = ({ osqueryPolicies, osqueryPoliciesLoading }: UseA
           policies: newPolicies,
         } = processAggregations(responseData.rawResponse.aggregations);
 
-        setPlatforms(newPlatforms);
+        setPlatforms(newPlatforms)
         setOverlap(newOverlap);
-        setPolicies(newPolicies);
+        setPolicies(newPolicies.map(p => {
+          const name = agentPolicyById[p.id]?.name ?? p.name
+          return {
+            ...p,
+            name
+          }
+        }));
       }
 
       setLoading(false);
       setTotalCount(responseData.totalCount);
     },
     {
-      enabled: !osqueryPoliciesLoading,
+      enabled: !osqueryPoliciesLoading && !agentPoliciesLoading,
     }
   );
 

x-pack/plugins/osquery/public/agents/use_agent_policies.ts

@@ -0,0 +1,34 @@
+/*
+ * 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 { useQueries, UseQueryResult } from 'react-query';
+import { useKibana } from '../common/lib/kibana';
+import { AgentPolicy, agentPolicyRouteService, GetOneAgentPolicyResponse } from '../../../fleet/common';
+
+export const useAgentPolicies = (policyIds: string[] = []) => {
+  const { http } = useKibana().services;
+
+  const agentResponse = useQueries(
+    policyIds.map((policyId) => ({
+      queryKey: ['agentPolicy', policyId],
+      queryFn: () => http.get(agentPolicyRouteService.getInfoPath(policyId)),
+      options: {enabled: policyIds.length > 0}
+    })),
+  ) as Array<UseQueryResult<GetOneAgentPolicyResponse>>;
+
+  const agentPoliciesLoading = agentResponse.some(p => p.isLoading) 
+  const agentPolicies = agentResponse.map(p => p.data?.item)
+  const agentPolicyById = agentPolicies.reduce((acc, p) => {
+    if (!p) {
+      return acc
+    }
+    acc[p.id] = p
+    return acc
+  }, {} as {[key: string]: AgentPolicy})
+
+  return { agentPoliciesLoading, agentPolicies, agentPolicyById };
+};

x-pack/plugins/osquery/public/agents/use_osquery_policies.ts

@@ -15,11 +15,13 @@ export const useOsqueryPolicies = () => {
   const { isLoading: osqueryPoliciesLoading, data: osqueryPolicies } = useQuery(
     ['osqueryPolicies'],
     async () => {
-      return await http.get('/api/fleet/package_policies', {
+      const d = await http.get('/api/fleet/package_policies', {
         query: {
           kuery: `${PACKAGE_POLICY_SAVED_OBJECT_TYPE}.package.name:osquery_manager`,
         },
       });
+      console.log('init', d)
+      return d
     },
     { select: (data) => data.items.map((p: { policy_id: string }) => p.policy_id) }
   );