Improve types

elastic · Nov 4, 2024 · a62728d · a62728d
1 parent b6bbce6
commit a62728d
Show file tree

Hide file tree

Showing 12 changed files with 158 additions and 130 deletions.
diff --git a/packages/kbn-es-types/src/search.ts b/packages/kbn-es-types/src/search.ts
@@ -674,12 +674,15 @@ export interface ESQLColumn {
 
 export type ESQLRow = unknown[];
 
-export interface ESQLSearchResponse {
-  columns: ESQLColumn[];
+export interface ESQLSearchResponse<
+  TColumn extends ESQLColumn = ESQLColumn,
+  TRow extends ESQLRow = ESQLRow
+> {
+  columns: TColumn[];
   // In case of ?drop_null_columns in the query, then
   // all_columns will have available and empty fields
   // while columns only the available ones (non nulls)
-  all_columns?: ESQLColumn[];
+  all_columns?: TRow[];
   values: ESQLRow[];
   took?: number;
   _clusters?: estypes.ClusterStatistics;

diff --git a/x-pack/packages/kbn-entities-schema/src/schema/entity.ts b/x-pack/packages/kbn-entities-schema/src/schema/entity.ts
@@ -11,7 +11,7 @@ import { arrayOfStringsSchema } from './common';
 export const entityBaseSchema = z.object({
   id: z.string(),
   type: z.string(),
-  identityFields: arrayOfStringsSchema,
+  identityFields: z.union([arrayOfStringsSchema, z.string()]),
   displayName: z.string(),
   metrics: z.optional(z.record(z.string(), z.number())),
   definitionVersion: z.string(),

diff --git a/...ck/packages/observability/observability_utils/es/client/create_observability_es_client.ts b/...ck/packages/observability/observability_utils/es/client/create_observability_es_client.ts
@@ -9,6 +9,7 @@ import type { ElasticsearchClient, Logger } from '@kbn/core/server';
 import type { ESQLSearchResponse, ESSearchRequest, InferSearchResponseOf } from '@kbn/es-types';
 import { withSpan } from '@kbn/apm-utils';
 import type { EsqlQueryRequest } from '@elastic/elasticsearch/lib/api/types';
+import { esqlResultToPlainObjects } from '../utils/esql_result_to_plain_objects';
 
 type SearchRequest = ESSearchRequest & {
   index: string | string[];
@@ -25,7 +26,10 @@ export interface ObservabilityElasticsearchClient {
     operationName: string,
     parameters: TSearchRequest
   ): Promise<InferSearchResponseOf<TDocument, TSearchRequest>>;
-  esql(operationName: string, parameters: EsqlQueryRequest): Promise<ESQLSearchResponse>;
+  esql<TDocument = unknown>(
+    operationName: string,
+    parameters: EsqlQueryRequest
+  ): Promise<TDocument[]>;
   client: ElasticsearchClient;
 }
 
@@ -40,7 +44,7 @@ export function createObservabilityEsClient({
 }): ObservabilityElasticsearchClient {
   return {
     client,
-    esql(operationName: string, parameters: EsqlQueryRequest) {
+    esql<TDocument = unknown>(operationName: string, parameters: EsqlQueryRequest) {
       logger.trace(() => `Request (${operationName}):\n${JSON.stringify(parameters, null, 2)}`);
       return withSpan({ name: operationName, labels: { plugin } }, () => {
         return client.esql.query(
@@ -54,7 +58,7 @@ export function createObservabilityEsClient({
       })
         .then((response) => {
           logger.trace(() => `Response (${operationName}):\n${JSON.stringify(response, null, 2)}`);
-          return response as unknown as ESQLSearchResponse;
+          return esqlResultToPlainObjects<TDocument>(response as unknown as ESQLSearchResponse);
         })
         .catch((error) => {
           throw error;

diff --git a/x-pack/packages/observability/observability_utils/es/utils/esql_result_to_plain_objects.ts b/x-pack/packages/observability/observability_utils/es/utils/esql_result_to_plain_objects.ts
@@ -6,17 +6,14 @@
  */
 
 import type { ESQLSearchResponse } from '@kbn/es-types';
-import { castArray } from 'lodash';
 import { unflattenObject } from '../../object/unflatten_object';
 
-export function esqlResultToPlainObjects<T extends Record<string, any>>(
-  result: ESQLSearchResponse,
-  knownArrayFields?: string[]
-): T[] {
-  const knownArrayFieldsSet = new Set(knownArrayFields);
+export function esqlResultToPlainObjects<TDocument = unknown>(
+  result: ESQLSearchResponse
+): TDocument[] {
   return result.values.map((row) => {
     return unflattenObject(
-      row.reduce<Record<string, unknown>>((acc, value, index) => {
+      row.reduce<Record<string, any>>((acc, value, index) => {
         const column = result.columns[index];
 
         if (!column) {
@@ -26,11 +23,11 @@ export function esqlResultToPlainObjects<T extends Record<string, any>>(
         // Removes the type suffix from the column name
         const name = column.name.replace(/\.(text|keyword)$/, '');
         if (!acc[name]) {
-          acc[column.name] = knownArrayFieldsSet.has(column.name) ? castArray(value) : value;
+          acc[column.name] = value;
         }
 
         return acc;
       }, {})
-    );
-  }) as T[];
+    ) as TDocument;
+  });
 }
diff --git a/x-pack/plugins/observability_solution/inventory/common/entities.ts b/x-pack/plugins/observability_solution/inventory/common/entities.ts
@@ -17,6 +17,7 @@ import {
 import { decode, encode } from '@kbn/rison';
 import { isRight } from 'fp-ts/lib/Either';
 import * as t from 'io-ts';
+import { AgentName } from '@kbn/elastic-agent-utils';
 
 export const entityColumnIdsRt = t.union([
   t.literal(ENTITY_DISPLAY_NAME),
@@ -122,3 +123,24 @@ export type EntityGroup = {
 export type InventoryEntityLatest = z.infer<typeof entityLatestSchema> & {
   alertsCount?: number;
 };
+
+export const isHostEntity = (
+  entity: InventoryEntityLatest
+): entity is InventoryEntityLatest & { cloud?: { provider: string } } => {
+  return entity.entity.type === 'host';
+};
+
+export const isContainerEntity = (
+  entity: InventoryEntityLatest
+): entity is InventoryEntityLatest & { cloud?: { provider: string } } => {
+  return entity.entity.type === 'container';
+};
+
+export const isServiceEntity = (
+  entity: InventoryEntityLatest
+): entity is InventoryEntityLatest & {
+  agent?: { name: AgentName };
+  service: { name: string; environment: string };
+} => {
+  return entity.entity.type === 'service';
+};
diff --git a/x-pack/plugins/observability_solution/inventory/public/components/entity_icon/index.tsx b/x-pack/plugins/observability_solution/inventory/public/components/entity_icon/index.tsx
@@ -6,31 +6,26 @@
  */
 
 import React from 'react';
-import { AGENT_NAME, CLOUD_PROVIDER, ENTITY_TYPES } from '@kbn/observability-shared-plugin/common';
 import { type CloudProvider, CloudProviderIcon, AgentIcon } from '@kbn/custom-icons';
 import { EuiFlexGroup, EuiFlexItem, EuiIcon } from '@elastic/eui';
-import type { AgentName } from '@kbn/elastic-agent-utils';
 import { euiThemeVars } from '@kbn/ui-theme';
-import type { InventoryEntityLatest } from '../../../common/entities';
+import {
+  isHostEntity,
+  type InventoryEntityLatest,
+  isContainerEntity,
+  isServiceEntity,
+} from '../../../common/entities';
 
 interface EntityIconProps {
   entity: InventoryEntityLatest;
 }
 
-type NotNullableCloudProvider = Exclude<CloudProvider, null>;
-
-const getSingleValue = <T,>(value?: T | T[] | null): T | undefined => {
-  return value == null ? undefined : Array.isArray(value) ? value[0] : value;
-};
-
 export function EntityIcon({ entity }: EntityIconProps) {
-  const entityType = entity.entity.type;
   const defaultIconSize = euiThemeVars.euiSizeL;
 
-  if (entityType === ENTITY_TYPES.HOST || entityType === ENTITY_TYPES.CONTAINER) {
-    const cloudProvider = getSingleValue(
-      entity[CLOUD_PROVIDER] as NotNullableCloudProvider | NotNullableCloudProvider[]
-    );
+  if (isHostEntity(entity) || isContainerEntity(entity)) {
+    const cloudProvider = entity.cloud?.provider;
+
     return (
       <EuiFlexGroup
         style={{ width: defaultIconSize, height: defaultIconSize }}
@@ -39,7 +34,7 @@ export function EntityIcon({ entity }: EntityIconProps) {
       >
         <EuiFlexItem grow={false}>
           <CloudProviderIcon
-            cloudProvider={cloudProvider}
+            cloudProvider={cloudProvider as CloudProvider | undefined}
             size="m"
             title={cloudProvider}
             role="presentation"
@@ -49,9 +44,8 @@ export function EntityIcon({ entity }: EntityIconProps) {
     );
   }
 
-  if (entityType === ENTITY_TYPES.SERVICE) {
-    const agentName = getSingleValue(entity[AGENT_NAME] as AgentName | AgentName[]);
-    return <AgentIcon agentName={agentName} role="presentation" />;
+  if (isServiceEntity(entity)) {
+    return <AgentIcon agentName={entity.agent?.name} role="presentation" />;
   }
 
   if (entityType.startsWith('kubernetes')) {

diff --git a/...ck/plugins/observability_solution/inventory/public/hooks/use_detail_view_redirect.test.ts b/...ck/plugins/observability_solution/inventory/public/hooks/use_detail_view_redirect.test.ts
@@ -8,35 +8,23 @@
 import { renderHook } from '@testing-library/react-hooks';
 import { useDetailViewRedirect } from './use_detail_view_redirect';
 import { useKibana } from './use_kibana';
-import {
-  AGENT_NAME,
-  CLOUD_PROVIDER,
-  CONTAINER_ID,
-  ENTITY_DEFINITION_ID,
-  ENTITY_DISPLAY_NAME,
-  ENTITY_ID,
-  ENTITY_IDENTITY_FIELDS,
-  ENTITY_LAST_SEEN,
-  ENTITY_TYPE,
-  HOST_NAME,
-  ENTITY_TYPES,
-  SERVICE_ENVIRONMENT,
-  SERVICE_NAME,
-} from '@kbn/observability-shared-plugin/common';
+import { CONTAINER_ID, HOST_NAME, SERVICE_NAME } from '@kbn/observability-shared-plugin/common';
 import { unflattenEntity } from '../../common/utils/unflatten_entity';
-import type { Entity } from '../../common/entities';
+import type { InventoryEntityLatest } from '../../common/entities';
 
 jest.mock('./use_kibana');
 jest.mock('../../common/utils/unflatten_entity');
 
 const useKibanaMock = useKibana as jest.Mock;
 const unflattenEntityMock = unflattenEntity as jest.Mock;
 
-const commonEntityFields: Partial<Entity> = {
-  [ENTITY_LAST_SEEN]: '2023-10-09T00:00:00Z',
-  [ENTITY_ID]: '1',
-  [ENTITY_DISPLAY_NAME]: 'entity_name',
-  [ENTITY_DEFINITION_ID]: 'entity_definition_id',
+const commonEntityFields: Partial<InventoryEntityLatest['entity']> = {
+  lastSeenTimestamp: '2023-10-09T00:00:00Z',
+  id: '1',
+  displayName: 'entity_name',
+  definitionId: 'entity_definition_id',
+  definitionVersion: '1',
+  schemaVersion: '1',
 };
 
 describe('useDetailViewRedirect', () => {
@@ -71,12 +59,18 @@ describe('useDetailViewRedirect', () => {
   });
 
   it('getEntityRedirectUrl should return the correct URL for host entity', () => {
-    const entity: Entity = {
-      ...(commonEntityFields as Entity),
-      [ENTITY_IDENTITY_FIELDS]: [HOST_NAME],
-      [ENTITY_TYPE]: 'host',
-      [HOST_NAME]: 'host-1',
-      [CLOUD_PROVIDER]: null,
+    const entity: InventoryEntityLatest = {
+      entity: {
+        ...(commonEntityFields as InventoryEntityLatest['entity']),
+        type: 'host',
+        identityFields: ['host.name'],
+      },
+      host: {
+        name: 'host-1',
+      },
+      cloud: {
+        provider: null,
+      },
     };
 
     mockGetIdentityFieldsValue.mockReturnValue({ [HOST_NAME]: 'host-1' });
@@ -90,12 +84,18 @@ describe('useDetailViewRedirect', () => {
   });
 
   it('getEntityRedirectUrl should return the correct URL for container entity', () => {
-    const entity: Entity = {
-      ...(commonEntityFields as Entity),
-      [ENTITY_IDENTITY_FIELDS]: [CONTAINER_ID],
-      [ENTITY_TYPE]: 'container',
-      [CONTAINER_ID]: 'container-1',
-      [CLOUD_PROVIDER]: null,
+    const entity: InventoryEntityLatest = {
+      entity: {
+        ...(commonEntityFields as InventoryEntityLatest['entity']),
+        type: 'container',
+        identityFields: ['container.id'],
+      },
+      container: {
+        id: 'container-1',
+      },
+      cloud: {
+        provider: null,
+      },
     };
 
     mockGetIdentityFieldsValue.mockReturnValue({ [CONTAINER_ID]: 'container-1' });
@@ -112,13 +112,19 @@ describe('useDetailViewRedirect', () => {
   });
 
   it('getEntityRedirectUrl should return the correct URL for service entity', () => {
-    const entity: Entity = {
-      ...(commonEntityFields as Entity),
-      [ENTITY_IDENTITY_FIELDS]: [SERVICE_NAME],
-      [ENTITY_TYPE]: 'service',
-      [SERVICE_NAME]: 'service-1',
-      [SERVICE_ENVIRONMENT]: 'prod',
-      [AGENT_NAME]: 'node',
+    const entity: InventoryEntityLatest = {
+      entity: {
+        ...(commonEntityFields as InventoryEntityLatest['entity']),
+        type: 'service',
+        identityFields: ['service.name'],
+      },
+      agent: {
+        name: 'node',
+      },
+      service: {
+        name: 'service-1',
+        environment: 'prod',
+      },
     };
     mockGetIdentityFieldsValue.mockReturnValue({ [SERVICE_NAME]: 'service-1' });
     mockGetRedirectUrl.mockReturnValue('service-overview-url');
@@ -145,10 +151,15 @@ describe('useDetailViewRedirect', () => {
     [ENTITY_TYPES.KUBERNETES.STATEFULSET.ecs, 'kubernetes-21694370-bcb2-11ec-b64f-7dd6e8e82013'],
   ].forEach(([entityType, dashboardId]) => {
     it(`getEntityRedirectUrl should return the correct URL for ${entityType} entity`, () => {
-      const entity: Entity = {
-        ...(commonEntityFields as Entity),
-        [ENTITY_IDENTITY_FIELDS]: ['some.field'],
-        [ENTITY_TYPE]: entityType,
+      const entity: InventoryEntityLatest = {
+        entity: {
+          ...(commonEntityFields as InventoryEntityLatest['entity']),
+          type: entityType,
+          identityFields: ['some.field'],
+        },
+        some: {
+          field: 'some-value',
+        },
       };
 
       mockAsKqlFilter.mockReturnValue('kql-query');