[backend] include draft index in draft context queries

opencti-platform/opencti-graphql/src/database/engine.js

@@ -166,6 +166,7 @@ import { ENTITY_TYPE_DELETE_OPERATION } from '../modules/deleteOperation/deleteO
 import { buildEntityData } from './data-builder';
 import { controlUserConfidenceAgainstElement } from '../utils/confidence-level';
 import { enrichWithRemoteCredentials } from '../config/credentials';
+import { inDraftContext } from '../utils/draftContext';
 
 const ELK_ENGINE = 'elk';
 const OPENSEARCH_ENGINE = 'opensearch';
@@ -2622,6 +2623,11 @@ const completeSpecialFilterKeys = async (context, user, inputFilters) => {
   };
 };
 
+const getIndicesToQuery = (index, context, user) => {
+  const draftContext = inDraftContext(user);
+  return index + (!draftContext ? '' : (`,${READ_INDEX_DRAFT}`));
+};
+
 const elQueryBodyBuilder = async (context, user, options) => {
   // eslint-disable-next-line no-use-before-define
   const { ids = [], after, orderBy = null, orderMode = 'asc', noSize = false, noSort = false, intervalInclude = false } = options;
@@ -2714,11 +2720,39 @@ const elQueryBodyBuilder = async (context, user, options) => {
   } else { // If not ordering criteria, order by standard_id
     ordering.push({ 'standard_id.keyword': 'asc' });
   }
+  // Handle draft
+  const draftContext = inDraftContext(user);
+  const draftMust = [];
+  if (draftContext) {
+    const mustLive = {
+      bool: {
+        must_not: [
+          { term: { _index: READ_INDEX_DRAFT } },
+          { match: { draft_ids: draftContext } }
+        ]
+      }
+    };
+    const mustDraft = {
+      bool: {
+        must: [
+          { term: { _index: READ_INDEX_DRAFT } },
+          { match: { draft_ids: draftContext } }
+        ]
+      }
+    };
+    const draftBool = {
+      bool: {
+        should: [mustLive, mustDraft],
+        minimum_should_match: 1,
+      },
+    };
+    draftMust.push(draftBool);
+  }
   // Build query
   const body = {
     query: {
       bool: {
-        must: [...accessMust, ...mustFilters],
+        must: [...accessMust, ...mustFilters, ...draftMust],
         must_not: accessMustNot,
       },
     },
@@ -2750,7 +2784,7 @@ export const elRawCount = async (query) => {
 };
 export const elCount = async (context, user, indexName, options = {}) => {
   const body = await elQueryBodyBuilder(context, user, { ...options, noSize: true, noSort: true });
-  const query = { index: indexName, body };
+  const query = { index: getIndicesToQuery(indexName), body };
   logApp.debug('[SEARCH] elCount', { query });
   return elRawCount(query);
 };
@@ -2797,7 +2831,7 @@ export const elHistogramCount = async (context, user, indexName, options = {}) =
     },
   };
   const query = {
-    index: indexName,
+    index: getIndicesToQuery(indexName),
     _source_excludes: '*', // Dont need to get anything
     body,
   };
@@ -2830,7 +2864,7 @@ export const elAggregationCount = async (context, user, indexName, options = {})
     },
   };
   const query = {
-    index: indexName,
+    index: getIndicesToQuery(indexName),
     body,
   };
   logApp.debug('[SEARCH] aggregationCount', { query });
@@ -2944,7 +2978,7 @@ export const elAggregationRelationsCount = async (context, user, indexName, opti
       },
     };
   }
-  const query = { index: indexName, body };
+  const query = { index: getIndicesToQuery(indexName), body };
   logApp.debug('[SEARCH] aggregationRelationsCount', { query });
   const isIdFields = field?.endsWith('internal_id');
   return elRawSearch(context, user, types, query)
@@ -2995,7 +3029,7 @@ export const elAggregationNestedTermsWithFilter = async (context, user, indexNam
     }
   };
   const query = {
-    index: indexName,
+    index: getIndicesToQuery(indexName),
     body,
   };
   logApp.debug('[SEARCH] elAggregationNestedTermsWithFilter', { query });
@@ -3038,7 +3072,7 @@ export const elAggregationsList = async (context, user, indexName, aggregations,
     }
   }
   const query = {
-    index: indexName,
+    index: getIndicesToQuery(indexName),
     track_total_hits: true,
     _source: false,
     body,
@@ -3075,7 +3109,7 @@ export const elPaginate = async (context, user, indexName, options = {}) => {
     body.size = ES_MAX_PAGINATION;
   }
   const query = {
-    index: indexName,
+    index: getIndicesToQuery(indexName),
     track_total_hits: true,
     _source: baseData ? baseFields : true,
     body,