Tempo: Switch out Select with AsyncSelect component to get loading state in Tempo Search (#45110)

* Replace Select with AsyncSelect to get loading state

public/app/plugins/datasource/tempo/QueryEditor/NativeSearch.test.tsx

@@ -0,0 +1,109 @@
+import NativeSearch from './NativeSearch';
+import React from 'react';
+import { act, render, screen } from '@testing-library/react';
+import { TempoDatasource, TempoQuery } from '../datasource';
+import userEvent from '@testing-library/user-event';
+
+const getOptions = jest.fn().mockImplementation(() => {
+  return Promise.resolve([
+    {
+      value: 'customer',
+      label: 'customer',
+    },
+    {
+      value: 'driver',
+      label: 'driver',
+    },
+  ]);
+});
+
+jest.mock('../language_provider', () => {
+  return jest.fn().mockImplementation(() => {
+    return { getOptions };
+  });
+});
+
+const mockQuery = {
+  refId: 'A',
+  queryType: 'nativeSearch',
+  key: 'Q-595a9bbc-2a25-49a7-9249-a52a0a475d83-0',
+  serviceName: 'driver',
+} as TempoQuery;
+
+describe('NativeSearch', () => {
+  it('should call the `onChange` function on click of the Input', async () => {
+    const promise = Promise.resolve();
+    const handleOnChange = jest.fn(() => promise);
+    const fakeOptionChoice = {
+      key: 'Q-595a9bbc-2a25-49a7-9249-a52a0a475d83-0',
+      queryType: 'nativeSearch',
+      refId: 'A',
+      serviceName: 'driver',
+      spanName: 'driver',
+    };
+
+    render(
+      <NativeSearch
+        datasource={{} as TempoDatasource}
+        query={mockQuery}
+        onChange={handleOnChange}
+        onRunQuery={() => {}}
+      />
+    );
+
+    const asyncServiceSelect = await screen.findByRole('combobox', { name: 'select-span-name' });
+
+    expect(asyncServiceSelect).toBeInTheDocument();
+    userEvent.click(asyncServiceSelect);
+
+    const driverOption = await screen.findByText('driver');
+    userEvent.click(driverOption);
+
+    expect(handleOnChange).toHaveBeenCalledWith(fakeOptionChoice);
+  });
+});
+
+describe('TempoLanguageProvider with delay', () => {
+  const getOptions2 = jest.fn().mockImplementation(() => {
+    return Promise.resolve([
+      {
+        value: 'customer',
+        label: 'customer',
+      },
+      {
+        value: 'driver',
+        label: 'driver',
+      },
+    ]);
+  });
+
+  jest.mock('../language_provider', () => {
+    return jest.fn().mockImplementation(() => {
+      setTimeout(() => {
+        return { getOptions2 };
+      }, 3000);
+    });
+  });
+
+  it('should show loader', async () => {
+    const promise = Promise.resolve();
+    const handleOnChange = jest.fn(() => promise);
+
+    render(
+      <NativeSearch
+        datasource={{} as TempoDatasource}
+        query={mockQuery}
+        onChange={handleOnChange}
+        onRunQuery={() => {}}
+      />
+    );
+
+    const asyncServiceSelect = screen.getByRole('combobox', { name: 'select-span-name' });
+
+    userEvent.click(asyncServiceSelect);
+    const loader = screen.getByText('Loading options...');
+
+    expect(loader).toBeInTheDocument();
+    await act(() => promise);
+  });
+});

public/app/plugins/datasource/tempo/QueryEditor/NativeSearch.tsx

@@ -1,4 +1,4 @@
-import React, { useState, useEffect, useMemo } from 'react';
+import React, { useCallback, useState, useEffect, useMemo } from 'react';
 import {
   InlineFieldRow,
   InlineField,
@@ -8,7 +8,7 @@ import {
   BracesPlugin,
   TypeaheadInput,
   TypeaheadOutput,
-  Select,
+  AsyncSelect,
   Alert,
   useStyles2,
 } from '@grafana/ui';
@@ -48,50 +48,59 @@ const NativeSearch = ({ datasource, query, onChange, onBlur, onRunQuery }: Props
   const styles = useStyles2(getStyles);
   const languageProvider = useMemo(() => new TempoLanguageProvider(datasource), [datasource]);
   const [hasSyntaxLoaded, setHasSyntaxLoaded] = useState(false);
-  const [autocomplete, setAutocomplete] = useState<{
-    serviceNameOptions: Array<SelectableValue<string>>;
-    spanNameOptions: Array<SelectableValue<string>>;
-  }>({
-    serviceNameOptions: [],
-    spanNameOptions: [],
+  const [asyncServiceNameValue, setAsyncServiceNameValue] = useState<SelectableValue<any>>({
+    value: '',
+  });
+  const [asyncSpanNameValue, setAsyncSpanNameValue] = useState<SelectableValue<any>>({
+    value: '',
   });
   const [error, setError] = useState(null);
   const [inputErrors, setInputErrors] = useState<{ [key: string]: boolean }>({});
+  const [isLoading, setIsLoading] = useState<{
+    serviceName: boolean;
+    spanName: boolean;
+  }>({
+    serviceName: false,
+    spanName: false,
+  });
+
+  async function fetchOptionsCallback(nameType: string, lp: TempoLanguageProvider) {
+    try {
+      const res = await lp.getOptions(nameType === 'serviceName' ? 'service.name' : 'name');
+      setIsLoading((prevValue) => ({ ...prevValue, [nameType]: false }));
+      return res;
+    } catch (error) {
+      if (error?.status === 404) {
+        setIsLoading((prevValue) => ({ ...prevValue, [nameType]: false }));
+      } else {
+        dispatch(notifyApp(createErrorNotification('Error', error)));
+        setIsLoading((prevValue) => ({ ...prevValue, [nameType]: false }));
+      }
+      setError(error);
+      return [];
+    }
+  }
 
-  const fetchServiceNameOptions = useMemo(
-    () =>
-      debounce(
-        async () => {
-          const res = await languageProvider.getOptions('service.name');
-          setAutocomplete((prev) => ({ ...prev, serviceNameOptions: res }));
-        },
-        500,
-        { leading: true, trailing: true }
-      ),
+  const loadOptionsOfType = useCallback(
+    (nameType: string) => {
+      setIsLoading((prevValue) => ({ ...prevValue, [nameType]: true }));
+      return fetchOptionsCallback(nameType, languageProvider);
+    },
     [languageProvider]
   );
 
-  const fetchSpanNameOptions = useMemo(
-    () =>
-      debounce(
-        async () => {
-          const res = await languageProvider.getOptions('name');
-          setAutocomplete((prev) => ({ ...prev, spanNameOptions: res }));
-        },
-        500,
-        { leading: true, trailing: true }
-      ),
-    [languageProvider]
+  const fetchOptionsOfType = useCallback(
+    (nameType: string) => debounce(() => loadOptionsOfType(nameType), 500, { leading: true, trailing: true }),
+    [loadOptionsOfType]
   );
 
   useEffect(() => {
-    const fetchAutocomplete = async () => {
+    const fetchOptions = async () => {
       try {
         await languageProvider.start();
-        const serviceNameOptions = await languageProvider.getOptions('service.name');
-        const spanNameOptions = await languageProvider.getOptions('name');
+        fetchOptionsCallback('serviceName', languageProvider);
+        fetchOptionsCallback('spanName', languageProvider);
         setHasSyntaxLoaded(true);
-        setAutocomplete({ serviceNameOptions, spanNameOptions });
       } catch (error) {
         // Display message if Tempo is connected but search 404's
         if (error?.status === 404) {
@@ -99,10 +108,11 @@ const NativeSearch = ({ datasource, query, onChange, onBlur, onRunQuery }: Props
         } else {
           dispatch(notifyApp(createErrorNotification('Error', error)));
         }
+        setHasSyntaxLoaded(true);
       }
     };
-    fetchAutocomplete();
-  }, [languageProvider, fetchServiceNameOptions, fetchSpanNameOptions]);
+    fetchOptions();
+  }, [languageProvider, fetchOptionsOfType]);
 
   const onTypeahead = async (typeahead: TypeaheadInput): Promise<TypeaheadOutput> => {
     return await languageProvider.provideCompletionItems(typeahead);
@@ -127,41 +137,53 @@ const NativeSearch = ({ datasource, query, onChange, onBlur, onRunQuery }: Props
       <div className={styles.container}>
         <InlineFieldRow>
           <InlineField label="Service Name" labelWidth={14} grow>
-            <Select
+            <AsyncSelect
               inputId="service"
               menuShouldPortal
-              options={autocomplete.serviceNameOptions}
-              value={query.serviceName || ''}
+              cacheOptions={false}
+              loadOptions={fetchOptionsOfType('serviceName')}
+              onOpenMenu={fetchOptionsOfType('serviceName')}
+              isLoading={isLoading.serviceName}
+              value={asyncServiceNameValue.value}
               onChange={(v) => {
+                setAsyncServiceNameValue({
+                  value: v,
+                });
                 onChange({
                   ...query,
                   serviceName: v?.value || undefined,
                 });
               }}
               placeholder="Select a service"
-              onOpenMenu={fetchServiceNameOptions}
               isClearable
+              defaultOptions
               onKeyDown={onKeyDown}
+              aria-label={'select-service-name'}
             />
           </InlineField>
         </InlineFieldRow>
         <InlineFieldRow>
           <InlineField label="Span Name" labelWidth={14} grow>
-            <Select
+            <AsyncSelect
               inputId="spanName"
               menuShouldPortal
-              options={autocomplete.spanNameOptions}
-              value={query.spanName || ''}
+              cacheOptions={false}
+              loadOptions={fetchOptionsOfType('spanName')}
+              onOpenMenu={fetchOptionsOfType('spanName')}
+              isLoading={isLoading.spanName}
+              value={asyncSpanNameValue.value}
               onChange={(v) => {
+                setAsyncSpanNameValue({ value: v });
                 onChange({
                   ...query,
                   spanName: v?.value || undefined,
                 });
               }}
               placeholder="Select a span"
-              onOpenMenu={fetchSpanNameOptions}
               isClearable
+              defaultOptions
               onKeyDown={onKeyDown}
+              aria-label={'select-span-name'}
             />
           </InlineField>
         </InlineFieldRow>