Add code review suggestions

hop-dev · Sep 17, 2024 · 1b94ce7 · 1b94ce7
1 parent 7064282
commit 1b94ce7
Show file tree

Hide file tree

Showing 8 changed files with 26 additions and 21 deletions.
diff --git a/...plugins/security_solution/common/api/entity_analytics/entity_store/entities/common.gen.ts b/...plugins/security_solution/common/api/entity_analytics/entity_store/entities/common.gen.ts
@@ -72,3 +72,6 @@ export const HostEntity = z.object({
     })
     .optional(),
 });
+
+export type Entity = z.infer<typeof Entity>;
+export const Entity = z.union([UserEntity, HostEntity]);
diff --git a/...ns/security_solution/common/api/entity_analytics/entity_store/entities/common.schema.yaml b/...ns/security_solution/common/api/entity_analytics/entity_store/entities/common.schema.yaml
@@ -152,3 +152,8 @@ components:
             - type
             - firstSeenTimestamp
             - definitionId
+
+    Entity:
+      oneOf:
+        - $ref: '#/components/schemas/UserEntity'
+        - $ref: '#/components/schemas/HostEntity'
diff --git a/.../security_solution/common/api/entity_analytics/entity_store/entities/list_entities.gen.ts b/.../security_solution/common/api/entity_analytics/entity_store/entities/list_entities.gen.ts
@@ -18,7 +18,7 @@ import { z } from '@kbn/zod';
 import { ArrayFromString } from '@kbn/zod-helpers';
 
 import { EntityType, InspectQuery } from '../common.gen';
-import { UserEntity, HostEntity } from './common.gen';
+import { Entity } from './common.gen';
 
 export type ListEntitiesRequestQuery = z.infer<typeof ListEntitiesRequestQuery>;
 export const ListEntitiesRequestQuery = z.object({
@@ -36,7 +36,7 @@ export type ListEntitiesRequestQueryInput = z.input<typeof ListEntitiesRequestQu
 
 export type ListEntitiesResponse = z.infer<typeof ListEntitiesResponse>;
 export const ListEntitiesResponse = z.object({
-  records: z.array(z.union([UserEntity, HostEntity])),
+  records: z.array(Entity),
   page: z.number().int().min(1),
   per_page: z.number().int().min(1).max(1000),
   total: z.number().int().min(0),

diff --git a/...rity_solution/common/api/entity_analytics/entity_store/entities/list_entities.schema.yaml b/...rity_solution/common/api/entity_analytics/entity_store/entities/list_entities.schema.yaml
@@ -62,9 +62,7 @@ paths:
                   records:
                     type: array
                     items:
-                      oneOf:
-                        - $ref: './common.schema.yaml#/components/schemas/UserEntity'
-                        - $ref: './common.schema.yaml#/components/schemas/HostEntity'
+                      $ref: './common.schema.yaml#/components/schemas/Entity'
                   page:
                     type: integer
                     minimum: 1

diff --git a/x-pack/plugins/security_solution/server/lib/entity_analytics/entity_store/constants.ts b/x-pack/plugins/security_solution/server/lib/entity_analytics/entity_store/constants.ts
@@ -19,3 +19,5 @@ export const ENGINE_STATUS: Record<Uppercase<EngineStatus>, EngineStatus> = {
   STARTED: 'started',
   STOPPED: 'stopped',
 };
+
+export const MAX_SEARCH_RESPONSE_SIZE = 10_000;
diff --git a/x-pack/plugins/security_solution/server/lib/entity_analytics/entity_store/definition.ts b/x-pack/plugins/security_solution/server/lib/entity_analytics/entity_store/definition.ts
@@ -7,9 +7,10 @@
 
 import { entityDefinitionSchema, type EntityDefinition } from '@kbn/entities-schema';
 import { ENTITY_STORE_DEFAULT_SOURCE_INDICES } from './constants';
+import { getEntityDefinitionId } from './utils/utils';
 
 export const HOST_ENTITY_DEFINITION: EntityDefinition = entityDefinitionSchema.parse({
-  id: 'ea_host_entity_store',
+  id: getEntityDefinitionId('host'),
   name: 'EA Host Store',
   type: 'host',
   indexPatterns: ENTITY_STORE_DEFAULT_SOURCE_INDICES,
@@ -33,9 +34,8 @@ export const HOST_ENTITY_DEFINITION: EntityDefinition = entityDefinitionSchema.p
 });
 
 export const USER_ENTITY_DEFINITION: EntityDefinition = entityDefinitionSchema.parse({
-  id: 'ea_user_entity_store',
+  id: getEntityDefinitionId('user'),
   name: 'EA User Store',
-  type: 'user',
   indexPatterns: ENTITY_STORE_DEFAULT_SOURCE_INDICES,
   identityFields: ['user.name'],
   displayNameTemplate: '{{user.name}}',

diff --git a/...ns/security_solution/server/lib/entity_analytics/entity_store/entity_store_data_client.ts b/...ns/security_solution/server/lib/entity_analytics/entity_store/entity_store_data_client.ts
@@ -9,11 +9,8 @@ import type { Logger, ElasticsearchClient, SavedObjectsClientContract } from '@k
 import type { EntityClient } from '@kbn/entityManager-plugin/server/lib/entity_client';
 
 import type { SortOrder } from '@elastic/elasticsearch/lib/api/types';
+import type { Entity } from '../../../../common/api/entity_analytics/entity_store/entities/common.gen';
 import { createQueryFilterClauses } from '../../../utils/build_query';
-import type {
-  HostEntity,
-  UserEntity,
-} from '../../../../common/api/entity_analytics/entity_store/entities/common.gen';
 import type {
   InitEntityStoreRequestBody,
   InitEntityStoreResponse,
@@ -26,7 +23,7 @@ import type {
 import { entityEngineDescriptorTypeName } from './saved_object';
 import { EngineDescriptorClient } from './saved_object/engine_descriptor';
 import { getEntitiesIndexName, getEntityDefinition } from './utils/utils';
-import { ENGINE_STATUS } from './constants';
+import { ENGINE_STATUS, MAX_SEARCH_RESPONSE_SIZE } from './constants';
 
 interface EntityStoreClientOpts {
   logger: Logger;
@@ -45,8 +42,6 @@ interface SearchEntitiesParams {
   sortOrder: SortOrder;
 }
 
-const MAX_SEARCH_RESPONSE_SIZE = 10_000;
-
 export class EntityStoreDataClient {
   private engineClient: EngineDescriptorClient;
   constructor(private readonly options: EntityStoreClientOpts) {
@@ -137,7 +132,7 @@ export class EntityStoreDataClient {
   }
 
   public async searchEntities(params: SearchEntitiesParams): Promise<{
-    records: Array<UserEntity | HostEntity>;
+    records: Entity[];
     total: number;
     inspect: InspectQuery;
   }> {
@@ -154,7 +149,7 @@ export class EntityStoreDataClient {
       },
     };
 
-    const response = await this.options.esClient.search<UserEntity | HostEntity>({
+    const response = await this.options.esClient.search<Entity>({
       index,
       query,
       size: Math.min(perPage, MAX_SEARCH_RESPONSE_SIZE),
@@ -166,7 +161,7 @@ export class EntityStoreDataClient {
 
     const total = typeof hits.total === 'number' ? hits.total : hits.total?.value ?? 0;
 
-    const records = hits.hits.map((hit) => hit._source as UserEntity | HostEntity);
+    const records = hits.hits.map((hit) => hit._source as Entity);
 
     const inspect: InspectQuery = {
       dsl: [JSON.stringify({ index, body: query }, null, 2)],

diff --git a/x-pack/plugins/security_solution/server/lib/entity_analytics/entity_store/utils/utils.ts b/x-pack/plugins/security_solution/server/lib/entity_analytics/entity_store/utils/utils.ts
@@ -6,6 +6,7 @@
  */
 
 import type { SavedObjectsFindResponse } from '@kbn/core-saved-objects-api-server';
+import { ENTITY_LATEST_PREFIX_V1 } from '@kbn/entityManager-plugin/common/constants_entities';
 import type {
   EngineDescriptor,
   EntityType,
@@ -32,6 +33,7 @@ export const getByEntityTypeQuery = (entityType: EntityType) => {
   return `${entityEngineDescriptorTypeName}.attributes.type: ${entityType}`;
 };
 
-export const getEntitiesIndexName = (entityType: EntityType) => {
-  return `.entities.v1.latest.ea_${entityType}_entity_store`;
-};
+export const getEntitiesIndexName = (entityType: EntityType) =>
+  `${ENTITY_LATEST_PREFIX_V1}.${getEntityDefinitionId(entityType)}`;
+
+export const getEntityDefinitionId = (entityType: EntityType) => `ea_${entityType}_entity_store`;