Add partial results to search examples demo (#96366) (#99835)

* Add partial results to search examples demo

* Add functional test

* Move types into separate file and separate custom strategies

* Review feedback

* Update test

* Try to fix test

* Attempt to fix test

* Try to fix observable error

* Sanity check

* Fix test

* Another attempt

* Remove rxjs from strategy

Co-authored-by: Lukas Olson <olson.lukas@gmail.com>

examples/search_examples/common/index.ts

@@ -6,17 +6,7 @@
  * Side Public License, v 1.
  */
 
-import { IEsSearchResponse, IEsSearchRequest } from '../../../src/plugins/data/common';
-
 export const PLUGIN_ID = 'searchExamples';
 export const PLUGIN_NAME = 'Search Examples';
 
-export interface IMyStrategyRequest extends IEsSearchRequest {
-  get_cool: boolean;
-}
-export interface IMyStrategyResponse extends IEsSearchResponse {
-  cool: string;
-  executed_at: number;
-}
-
 export const SERVER_SEARCH_ROUTE_PATH = '/api/examples/search';

examples/search_examples/common/types.ts

@@ -0,0 +1,26 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import {
+  IEsSearchRequest,
+  IEsSearchResponse,
+  IKibanaSearchRequest,
+  IKibanaSearchResponse,
+} from '../../../src/plugins/data/common';
+
+export interface IMyStrategyRequest extends IEsSearchRequest {
+  get_cool: boolean;
+}
+export interface IMyStrategyResponse extends IEsSearchResponse {
+  cool: string;
+  executed_at: number;
+}
+
+export type FibonacciRequest = IKibanaSearchRequest<{ n: number }>;
+
+export type FibonacciResponse = IKibanaSearchResponse<{ values: number[] }>;

examples/search_examples/public/search/app.tsx

@@ -26,27 +26,27 @@ import {
   EuiCode,
   EuiComboBox,
   EuiFormLabel,
+  EuiFieldNumber,
+  EuiProgress,
   EuiTabbedContent,
+  EuiTabbedContentTab,
 } from '@elastic/eui';
 
 import { CoreStart } from '../../../../src/core/public';
 import { mountReactNode } from '../../../../src/core/public/utils';
 import { NavigationPublicPluginStart } from '../../../../src/plugins/navigation/public';
 
-import {
-  PLUGIN_ID,
-  PLUGIN_NAME,
-  IMyStrategyResponse,
-  SERVER_SEARCH_ROUTE_PATH,
-} from '../../common';
+import { PLUGIN_ID, PLUGIN_NAME, SERVER_SEARCH_ROUTE_PATH } from '../../common';
 
 import {
   DataPublicPluginStart,
+  IKibanaSearchResponse,
   IndexPattern,
   IndexPatternField,
   isCompleteResponse,
   isErrorResponse,
 } from '../../../../src/plugins/data/public';
+import { IMyStrategyResponse } from '../../common/types';
 
 interface SearchExamplesAppDeps {
   notifications: CoreStart['notifications'];
@@ -88,7 +88,10 @@ export const SearchExamplesApp = ({
 }: SearchExamplesAppDeps) => {
   const { IndexPatternSelect } = data.ui;
   const [getCool, setGetCool] = useState<boolean>(false);
+  const [fibonacciN, setFibonacciN] = useState<number>(10);
   const [timeTook, setTimeTook] = useState<number | undefined>();
+  const [total, setTotal] = useState<number>(100);
+  const [loaded, setLoaded] = useState<number>(0);
   const [indexPattern, setIndexPattern] = useState<IndexPattern | null>();
   const [fields, setFields] = useState<IndexPatternField[]>();
   const [selectedFields, setSelectedFields] = useState<IndexPatternField[]>([]);
@@ -99,7 +102,15 @@ export const SearchExamplesApp = ({
     IndexPatternField | null | undefined
   >();
   const [request, setRequest] = useState<Record<string, any>>({});
-  const [response, setResponse] = useState<Record<string, any>>({});
+  const [rawResponse, setRawResponse] = useState<Record<string, any>>({});
+  const [selectedTab, setSelectedTab] = useState(0);
+
+  function setResponse(response: IKibanaSearchResponse) {
+    setRawResponse(response.rawResponse);
+    setLoaded(response.loaded!);
+    setTotal(response.total!);
+    setTimeTook(response.rawResponse.took);
+  }
 
   // Fetch the default index pattern using the `data.indexPatterns` service, as the component is mounted.
   useEffect(() => {
@@ -152,8 +163,7 @@ export const SearchExamplesApp = ({
       .subscribe({
         next: (res) => {
           if (isCompleteResponse(res)) {
-            setResponse(res.rawResponse);
-            setTimeTook(res.rawResponse.took);
+            setResponse(res);
             const avgResult: number | undefined = res.rawResponse.aggregations
               ? // @ts-expect-error @elastic/elasticsearch no way to declare a type for aggregation in the search response
                 res.rawResponse.aggregations[1].value
@@ -234,7 +244,7 @@ export const SearchExamplesApp = ({
 
       setRequest(searchSource.getSearchRequestBody());
       const { rawResponse: res } = await searchSource.fetch$().toPromise();
-      setResponse(res);
+      setRawResponse(res);
 
       const message = <EuiText>Searched {res.hits.total} documents.</EuiText>;
       notifications.toasts.addSuccess(
@@ -247,7 +257,7 @@ export const SearchExamplesApp = ({
         }
       );
     } catch (e) {
-      setResponse(e.body);
+      setRawResponse(e.body);
       notifications.toasts.addWarning(`An error has occurred: ${e.message}`);
     }
   };
@@ -260,6 +270,41 @@ export const SearchExamplesApp = ({
     doAsyncSearch('myStrategy');
   };
 
+  const onPartialResultsClickHandler = () => {
+    setSelectedTab(1);
+    const req = {
+      params: {
+        n: fibonacciN,
+      },
+    };
+
+    // Submit the search request using the `data.search` service.
+    setRequest(req.params);
+    const searchSubscription$ = data.search
+      .search(req, {
+        strategy: 'fibonacciStrategy',
+      })
+      .subscribe({
+        next: (res) => {
+          setResponse(res);
+          if (isCompleteResponse(res)) {
+            notifications.toasts.addSuccess({
+              title: 'Query result',
+              text: 'Query finished',
+            });
+            searchSubscription$.unsubscribe();
+          } else if (isErrorResponse(res)) {
+            // TODO: Make response error status clearer
+            notifications.toasts.addWarning('An error has occurred');
+            searchSubscription$.unsubscribe();
+          }
+        },
+        error: () => {
+          notifications.toasts.addDanger('Failed to run search');
+        },
+      });
+  };
+
   const onClientSideSessionCacheClickHandler = () => {
     doAsyncSearch('myStrategy', data.search.session.getSessionId());
   };
@@ -284,7 +329,7 @@ export const SearchExamplesApp = ({
     doSearchSourceSearch(withOtherBucket);
   };
 
-  const reqTabs = [
+  const reqTabs: EuiTabbedContentTab[] = [
     {
       id: 'request',
       name: <EuiText data-test-subj="requestTab">Request</EuiText>,
@@ -318,6 +363,7 @@ export const SearchExamplesApp = ({
               values={{ time: timeTook ?? 'Unknown' }}
             />
           </EuiText>
+          <EuiProgress value={loaded} max={total} size="xs" data-test-subj="progressBar" />
           <EuiCodeBlock
             language="json"
             fontSize="s"
@@ -326,7 +372,7 @@ export const SearchExamplesApp = ({
             isCopyable
             data-test-subj="responseCodeBlock"
           >
-            {JSON.stringify(response, null, 2)}
+            {JSON.stringify(rawResponse, null, 2)}
           </EuiCodeBlock>
         </>
       ),
@@ -484,6 +530,37 @@ export const SearchExamplesApp = ({
                 </EuiText>
               </EuiText>
               <EuiSpacer />
+              <EuiTitle size="xs">
+                <h3>Handling partial results</h3>
+              </EuiTitle>
+              <EuiText>
+                The observable returned from <EuiCode>data.search</EuiCode> provides partial results
+                when the response is not yet complete. These can be handled to update a chart or
+                simply a progress bar:
+                <EuiSpacer />
+                <EuiCodeBlock language="html" fontSize="s" paddingSize="s" overflowHeight={450}>
+                  &lt;EuiProgress value=&#123;response.loaded&#125; max=&#123;response.total&#125;
+                  /&gt;
+                </EuiCodeBlock>
+                Below is an example showing a custom search strategy that emits partial Fibonacci
+                sequences up to the length provided, updates the response with each partial result,
+                and updates a progress bar (see the Response tab).
+                <EuiFieldNumber
+                  id="FibonacciN"
+                  placeholder="Number of Fibonacci numbers to generate"
+                  value={fibonacciN}
+                  onChange={(event) => setFibonacciN(parseInt(event.target.value, 10))}
+                />
+                <EuiButtonEmpty
+                  size="xs"
+                  onClick={onPartialResultsClickHandler}
+                  iconType="play"
+                  data-test-subj="requestFibonacci"
+                >
+                  Request Fibonacci sequence
+                </EuiButtonEmpty>
+              </EuiText>
+              <EuiSpacer />
               <EuiTitle size="s">
                 <h3>Writing a custom search strategy</h3>
               </EuiTitle>
@@ -567,8 +644,13 @@ export const SearchExamplesApp = ({
                 </EuiButtonEmpty>
               </EuiText>
             </EuiFlexItem>
+
             <EuiFlexItem style={{ width: '60%' }}>
-              <EuiTabbedContent tabs={reqTabs} />
+              <EuiTabbedContent
+                tabs={reqTabs}
+                selectedTab={reqTabs[selectedTab]}
+                onTabClick={(tab) => setSelectedTab(reqTabs.indexOf(tab))}
+              />
             </EuiFlexItem>
           </EuiFlexGrid>
         </EuiPageContentBody>

examples/search_examples/server/fibonacci_strategy.ts

@@ -0,0 +1,52 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import uuid from 'uuid';
+import { ISearchStrategy } from '../../../src/plugins/data/server';
+import { FibonacciRequest, FibonacciResponse } from '../common/types';
+
+export const fibonacciStrategyProvider = (): ISearchStrategy<
+  FibonacciRequest,
+  FibonacciResponse
+> => {
+  const responseMap = new Map<string, [number[], number, number]>();
+  return ({
+    search: (request: FibonacciRequest) => {
+      const id = request.id ?? uuid();
+      const [sequence, total, started] = responseMap.get(id) ?? [
+        [],
+        request.params?.n ?? 0,
+        Date.now(),
+      ];
+      if (sequence.length < 2) {
+        if (total > 0) sequence.push(sequence.length);
+      } else {
+        const [a, b] = sequence.slice(-2);
+        sequence.push(a + b);
+      }
+      const loaded = sequence.length;
+      responseMap.set(id, [sequence, total, started]);
+      if (loaded >= total) {
+        responseMap.delete(id);
+      }
+
+      const isRunning = loaded < total;
+      const isPartial = isRunning;
+      const took = Date.now() - started;
+      const values = sequence.slice(0, loaded);
+
+      // Usually we'd do something like "of()" but for some reason it breaks in tests with the error
+      // "You provided an invalid object where a stream was expected." which is why we have to cast
+      // down below as well
+      return [{ id, loaded, total, isRunning, isPartial, rawResponse: { took, values } }];
+    },
+    cancel: async (id: string) => {
+      responseMap.delete(id);
+    },
+  } as unknown) as ISearchStrategy<FibonacciRequest, FibonacciResponse>;
+};

examples/search_examples/server/my_strategy.ts

@@ -8,7 +8,7 @@
 
 import { map } from 'rxjs/operators';
 import { ISearchStrategy, PluginStart } from '../../../src/plugins/data/server';
-import { IMyStrategyResponse, IMyStrategyRequest } from '../common';
+import { IMyStrategyRequest, IMyStrategyResponse } from '../common/types';
 
 export const mySearchStrategyProvider = (
   data: PluginStart

examples/search_examples/server/plugin.ts

@@ -24,6 +24,7 @@ import {
 } from './types';
 import { mySearchStrategyProvider } from './my_strategy';
 import { registerRoutes } from './routes';
+import { fibonacciStrategyProvider } from './fibonacci_strategy';
 
 export class SearchExamplesPlugin
   implements
@@ -48,7 +49,9 @@ export class SearchExamplesPlugin
 
     core.getStartServices().then(([_, depsStart]) => {
       const myStrategy = mySearchStrategyProvider(depsStart.data);
+      const fibonacciStrategy = fibonacciStrategyProvider();
       deps.data.search.registerSearchStrategy('myStrategy', myStrategy);
+      deps.data.search.registerSearchStrategy('fibonacciStrategy', fibonacciStrategy);
       registerRoutes(router);
     });
 

src/plugins/data/server/search/strategies/ese_search/ese_search_strategy.ts

@@ -39,7 +39,7 @@ export const enhancedEsSearchStrategyProvider = (
   legacyConfig$: Observable<SharedGlobalConfig>,
   logger: Logger,
   usage?: SearchUsage
-): ISearchStrategy<IEsSearchRequest> => {
+): ISearchStrategy => {
   async function cancelAsyncSearch(id: string, esClient: IScopedClusterClient) {
     try {
       await esClient.asCurrentUser.asyncSearch.delete({ id });

x-pack/test/examples/search_examples/index.ts

@@ -26,5 +26,6 @@ export default function ({ getService, loadTestFile }: PluginFunctionalProviderC
     loadTestFile(require.resolve('./search_session_example'));
     loadTestFile(require.resolve('./search_example'));
     loadTestFile(require.resolve('./search_sessions_cache'));
+    loadTestFile(require.resolve('./partial_results_example'));
   });
 }

x-pack/test/examples/search_examples/partial_results_example.ts

@@ -0,0 +1,38 @@
+/*
+ * 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 expect from '@kbn/expect';
+import { FtrProviderContext } from '../../functional/ftr_provider_context';
+
+// eslint-disable-next-line import/no-default-export
+export default function ({ getService, getPageObjects }: FtrProviderContext) {
+  const testSubjects = getService('testSubjects');
+  const PageObjects = getPageObjects(['common']);
+  const retry = getService('retry');
+
+  describe('Partial results example', () => {
+    before(async () => {
+      await PageObjects.common.navigateToApp('searchExamples');
+      await testSubjects.click('/search');
+    });
+
+    it('should update a progress bar', async () => {
+      await testSubjects.click('responseTab');
+      const progressBar = await testSubjects.find('progressBar');
+
+      const value = await progressBar.getAttribute('value');
+      expect(value).to.be('0');
+
+      await testSubjects.click('requestFibonacci');
+
+      await retry.waitFor('update progress bar', async () => {
+        const newValue = await progressBar.getAttribute('value');
+        return parseFloat(newValue) > 0;
+      });
+    });
+  });
+}