Azure: Improve Azure Prometheus exemplars UI/UX (#97198)

* Add better error messaging

- Ensure it's clear to users what action needs to be taken if an operation ID cannot be found
- Add test

* Display resource picker for trace exemplar queries

* Remove unneeded test

* Update tests

pkg/tsdb/azuremonitor/loganalytics/azure-log-analytics-datasource.go

@@ -625,6 +625,9 @@ func getCorrelationWorkspaces(ctx context.Context, baseResource string, resource
 		}()
 
 		if res.StatusCode/100 != 2 {
+			if res.StatusCode == 404 {
+				return AzureCorrelationAPIResponse{}, backend.DownstreamError(fmt.Errorf("requested trace not found by Application Insights indexing. Select the relevant Application Insights resource to search for the Operation ID directly"))
+			}
 			return AzureCorrelationAPIResponse{}, utils.CreateResponseErrorFromStatusCode(res.StatusCode, res.Status, body)
 		}
 		var data AzureCorrelationAPIResponse

pkg/tsdb/azuremonitor/loganalytics/traces_test.go

@@ -26,6 +26,10 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 	ctx := context.Background()
 	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
+		if strings.Contains(r.URL.Path, "missing-op-id") {
+			w.WriteHeader(http.StatusNotFound)
+			return
+		}
 		w.WriteHeader(http.StatusOK)
 		var correlationRes AzureCorrelationAPIResponse
 		if strings.Contains(r.URL.Path, "test-op-id") {
@@ -95,7 +99,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 	tests := []struct {
 		name                   string
 		queryModel             backend.DataQuery
-		azureLogAnalyticsQuery AzureLogAnalyticsQuery
+		azureLogAnalyticsQuery *AzureLogAnalyticsQuery
 		Err                    require.ErrorAssertionFunc
 	}{
 		{
@@ -114,7 +118,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTable,
 				URL:          "v1/apps/r1/query",
@@ -190,7 +194,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTable,
 				URL:          "v1/apps/r1/query",
@@ -264,7 +268,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTable,
 				URL:          "v1/apps/r1/query",
@@ -337,7 +341,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTable,
 				URL:          "v1/apps/r1/query",
@@ -413,7 +417,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTable,
 				URL:          "v1/apps/r1/query",
@@ -493,7 +497,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTable,
 				URL:          "v1/apps/r1/query",
@@ -573,7 +577,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTable,
 				URL:          "v1/apps/r1/query",
@@ -651,7 +655,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTrace,
 				URL:          "v1/apps/r1/query",
@@ -724,7 +728,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTrace,
 				URL:          "v1/apps/r1/query",
@@ -800,7 +804,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTrace,
 				URL:          "v1/apps/r1/query",
@@ -842,7 +846,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTrace,
 				URL:          "v1/apps/r1/query",
@@ -920,7 +924,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTrace,
 				URL:          "v1/apps/r1/query",
@@ -997,7 +1001,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTrace,
 				URL:          "v1/apps/r1/query",
@@ -1076,7 +1080,7 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 				TimeRange: timeRange,
 				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
 			},
-			azureLogAnalyticsQuery: AzureLogAnalyticsQuery{
+			azureLogAnalyticsQuery: &AzureLogAnalyticsQuery{
 				RefID:        "A",
 				ResultFormat: dataquery.ResultFormatTrace,
 				URL:          "v1/apps/r1/query",
@@ -1147,13 +1151,34 @@ func TestBuildAppInsightsQuery(t *testing.T) {
 			},
 			Err: require.NoError,
 		},
+		{
+			name: "trace query with missing operation ID",
+			queryModel: backend.DataQuery{
+				JSON: []byte(fmt.Sprintf(`{
+							"queryType": "Azure Traces",
+							"azureTraces": {
+								"resources":     ["/subscriptions/test-sub/resourceGroups/test-rg/providers/Microsoft.Insights/components/r1"],
+								"resultFormat": "%s",
+								"traceTypes":	["trace"],
+								"operationId":	"missing-op-id"
+							}
+						}`, dataquery.ResultFormatTable)),
+				RefID:     "A",
+				TimeRange: timeRange,
+				QueryType: string(dataquery.AzureQueryTypeAzureTraces),
+			},
+			azureLogAnalyticsQuery: nil,
+			Err: func(tt require.TestingT, err error, i ...interface{}) {
+				require.ErrorContains(tt, err, "requested trace not found by Application Insights indexing. Select the relevant Application Insights resource to search for the Operation ID directly")
+			},
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			query, err := buildAppInsightsQuery(ctx, tt.queryModel, dsInfo, appInsightsRegExp, log.NewNullLogger())
 			tt.Err(t, err)
-			if diff := cmp.Diff(&tt.azureLogAnalyticsQuery, query); diff != "" {
+			if diff := cmp.Diff(tt.azureLogAnalyticsQuery, query); diff != "" {
 				t.Errorf("Result mismatch (-want +got): \n%s", diff)
 			}
 		})

public/app/plugins/datasource/azuremonitor/components/QueryEditor/QueryEditor.test.tsx

@@ -8,8 +8,9 @@ import createMockDatasource from '../../__mocks__/datasource';
 import { invalidNamespaceError } from '../../__mocks__/errors';
 import createMockQuery from '../../__mocks__/query';
 import { selectors } from '../../e2e/selectors';
-import { AzureQueryType } from '../../types';
+import { AzureQueryType, ResultFormat } from '../../types';
 import { selectOptionInTest } from '../../utils/testUtils';
+import { createMockResourcePickerData } from '../MetricsQueryEditor/MetricsQueryEditor.test';
 
 import QueryEditor from './QueryEditor';
 
@@ -204,4 +205,34 @@ describe('Azure Monitor QueryEditor', () => {
       expect(screen.getByTestId(selectors.components.queryEditor.userAuthFallbackAlert)).toBeInTheDocument()
     );
   });
+
+  it('should display the default subscription for exemplar type queries', async () => {
+    const mockDatasource = createMockDatasource({ resourcePickerData: createMockResourcePickerData() });
+    const defaultSubscriptionId = 'default-subscription-id';
+    mockDatasource.azureLogAnalyticsDatasource.getDefaultOrFirstSubscription = jest
+      .fn()
+      .mockResolvedValue(defaultSubscriptionId);
+    const query = createMockQuery();
+    delete query?.subscription;
+    delete query?.azureTraces;
+    query.queryType = AzureQueryType.TraceExemplar;
+    query.query = 'test-operation-id';
+    const onChange = jest.fn();
+
+    render(<QueryEditor query={query} datasource={mockDatasource} onChange={onChange} onRunQuery={() => {}} />);
+
+    await waitFor(() =>
+      expect(onChange).toHaveBeenCalledWith(
+        expect.objectContaining({
+          azureTraces: {
+            operationId: query.query,
+            resultFormat: ResultFormat.Trace,
+            resources: [`/subscriptions/${defaultSubscriptionId}`],
+          },
+        })
+      )
+    );
+    await waitFor(() => expect(screen.getByText(defaultSubscriptionId)).toBeInTheDocument());
+    expect(await screen.getByDisplayValue('test-operation-id')).toBeInTheDocument();
+  });
 });

public/app/plugins/datasource/azuremonitor/components/QueryEditor/QueryEditor.tsx

@@ -1,6 +1,7 @@
 import { css } from '@emotion/css';
 import { debounce } from 'lodash';
 import { useCallback, useMemo, useState } from 'react';
+import { useEffectOnce } from 'react-use';
 
 import { CoreApp, QueryEditorProps } from '@grafana/data';
 import { config, reportInteraction } from '@grafana/runtime';
@@ -43,6 +44,7 @@ const QueryEditor = ({
   const [errorMessage, setError] = useLastError();
   const onRunQuery = useMemo(() => debounce(baseOnRunQuery, 500), [baseOnRunQuery]);
   const [azureLogsCheatSheetModalOpen, setAzureLogsCheatSheetModalOpen] = useState(false);
+  const [defaultSubscriptionId, setDefaultSubscriptionId] = useState('');
 
   const onQueryChange = useCallback(
     (newQuery: AzureMonitorQuery) => {
@@ -52,7 +54,15 @@ const QueryEditor = ({
     [onChange, onRunQuery]
   );
 
-  const query = usePreparedQuery(baseQuery, onQueryChange);
+  useEffectOnce(() => {
+    if (baseQuery.queryType === AzureQueryType.TraceExemplar) {
+      datasource.azureLogAnalyticsDatasource.getDefaultOrFirstSubscription().then((subscription) => {
+        setDefaultSubscriptionId(subscription || '');
+      });
+    }
+  });
+
+  const query = usePreparedQuery(baseQuery, onQueryChange, defaultSubscriptionId);
 
   const subscriptionId = query.subscription || datasource.azureMonitorDatasource.defaultSubscriptionId;
   const basicLogsEnabled =

public/app/plugins/datasource/azuremonitor/components/QueryEditor/usePreparedQuery.ts

@@ -9,21 +9,22 @@ const DEFAULT_QUERY = {
   queryType: AzureQueryType.AzureMonitor,
 };
 
-const transformExemplarQuery = (query: AzureMonitorQuery) => {
+const transformExemplarQuery = (query: AzureMonitorQuery, defaultSubscriptionId: string) => {
   if (query.queryType === AzureQueryType.TraceExemplar && query.query !== '' && !query.azureTraces) {
     query.azureTraces = {
       operationId: query.query,
       resultFormat: ResultFormat.Trace,
+      resources: [`/subscriptions/${defaultSubscriptionId}`],
     };
   }
 
   return query;
 };
 
-const prepareQuery = (query: AzureMonitorQuery) => {
+const prepareQuery = (query: AzureMonitorQuery, defaultSubscriptionId: string) => {
   // Note: _.defaults does not apply default values deeply.
   const withDefaults = defaults({}, query, DEFAULT_QUERY);
-  const transformedQuery = transformExemplarQuery(withDefaults);
+  const transformedQuery = transformExemplarQuery(withDefaults, defaultSubscriptionId);
   const migratedQuery = migrateQuery(transformedQuery);
 
   // If we didn't make any changes to the object, then return the original object to keep the
@@ -34,8 +35,12 @@ const prepareQuery = (query: AzureMonitorQuery) => {
 /**
  * Returns queries with some defaults + migrations, and calls onChange function to notify if it changes
  */
-const usePreparedQuery = (query: AzureMonitorQuery, onChangeQuery: (newQuery: AzureMonitorQuery) => void) => {
+const usePreparedQuery = (
-  const preparedQuery = useMemo(() => prepareQuery(query), [query]);
+  query: AzureMonitorQuery,
+  onChangeQuery: (newQuery: AzureMonitorQuery) => void,
+  defaultSubscriptionId: string
+) => {
+  const preparedQuery = useMemo(() => prepareQuery(query, defaultSubscriptionId), [query, defaultSubscriptionId]);
 
   useEffect(() => {
     if (preparedQuery !== query) {

public/app/plugins/datasource/azuremonitor/components/TracesQueryEditor/TracesQueryEditor.test.tsx

@@ -3,7 +3,6 @@ import userEvent from '@testing-library/user-event';
 
 import createMockDatasource from '../../__mocks__/datasource';
 import createMockQuery from '../../__mocks__/query';
-import { AzureQueryType } from '../../dataquery.gen';
 import { createMockResourcePickerData } from '../MetricsQueryEditor/MetricsQueryEditor.test';
 
 import TracesQueryEditor from './TracesQueryEditor';
@@ -184,58 +183,4 @@ describe('TracesQueryEditor', () => {
       })
     );
   });
-
-  it('should not display the resource selector for exemplar type queries', async () => {
-    const mockDatasource = createMockDatasource({ resourcePickerData: createMockResourcePickerData() });
-    const query = createMockQuery();
-    delete query?.subscription;
-    delete query?.azureTraces?.resources;
-    query.queryType = AzureQueryType.TraceExemplar;
-    query.azureTraces = { operationId: 'test-operation-id' };
-    const onChange = jest.fn();
-
-    render(
-      <TracesQueryEditor
-        query={query}
-        datasource={mockDatasource}
-        variableOptionGroup={variableOptionGroup}
-        onChange={onChange}
-        setError={() => {}}
-      />
-    );
-
-    expect(await screen.queryByRole('button', { name: 'Select a resource' })).not.toBeInTheDocument();
-    expect(await screen.getByDisplayValue('test-operation-id')).toBeInTheDocument();
-  });
-
-  it('should not display the resource selector for exemplar type queries', async () => {
-    const mockDatasource = createMockDatasource({ resourcePickerData: createMockResourcePickerData() });
-    const query = createMockQuery();
-    delete query?.subscription;
-    delete query?.azureTraces?.resources;
-    query.queryType = AzureQueryType.TraceExemplar;
-    query.azureTraces = { operationId: 'test-operation-id' };
-    const onChange = jest.fn();
-
-    render(
-      <TracesQueryEditor
-        query={query}
-        datasource={mockDatasource}
-        variableOptionGroup={variableOptionGroup}
-        onChange={onChange}
-        setError={() => {}}
-      />
-    );
-
-    const operationIDInput = await screen.getByDisplayValue('test-operation-id');
-    await userEvent.clear(operationIDInput);
-
-    expect(onChange).toHaveBeenCalledWith(
-      expect.objectContaining({
-        azureTraces: undefined,
-        queryType: AzureQueryType.AzureTraces,
-        query: undefined,
-      })
-    );
-  });
 });

public/app/plugins/datasource/azuremonitor/components/TracesQueryEditor/TracesQueryEditor.tsx

@@ -88,7 +88,6 @@ const TracesQueryEditor = ({
   return (
     <span data-testid={selectors.components.queryEditor.tracesQueryEditor.container.input}>
       <EditorRows>
-        {query.queryType !== AzureQueryType.TraceExemplar && (
         <EditorRow>
           <EditorFieldGroup>
             <ResourceField
@@ -118,7 +117,6 @@ const TracesQueryEditor = ({
             />
           </EditorFieldGroup>
         </EditorRow>
-        )}
         <EditorRow>
           <EditorFieldGroup>
             <TraceTypeField