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
2025-02-25 18:55:37 -06:00 · 2025-01-09 11:33:47 +00:00 · 2025-01-09 11:33:47 +00:00 · 2860eb52d0
commit 2860eb52d0
parent 42f6f917c9
7 changed files with 126 additions and 109 deletions
--- a/pkg/tsdb/azuremonitor/loganalytics/azure-log-analytics-datasource.go
+++ b/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
--- a/pkg/tsdb/azuremonitor/loganalytics/traces_test.go
+++ b/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)
 			}
 		})
--- a/public/app/plugins/datasource/azuremonitor/components/QueryEditor/QueryEditor.test.tsx
+++ b/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();
+  });
 });
--- a/public/app/plugins/datasource/azuremonitor/components/QueryEditor/QueryEditor.tsx
+++ b/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 =
--- a/public/app/plugins/datasource/azuremonitor/components/QueryEditor/usePreparedQuery.ts
+++ b/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 preparedQuery = useMemo(() => prepareQuery(query), [query]);
+const usePreparedQuery = (
+  query: AzureMonitorQuery,
+  onChangeQuery: (newQuery: AzureMonitorQuery) => void,
+  defaultSubscriptionId: string
+) => {
+  const preparedQuery = useMemo(() => prepareQuery(query, defaultSubscriptionId), [query, defaultSubscriptionId]);

  useEffect(() => {
    if (preparedQuery !== query) {
--- a/public/app/plugins/datasource/azuremonitor/components/TracesQueryEditor/TracesQueryEditor.test.tsx
+++ b/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,
-      })
-    );
-  });
 });
--- a/public/app/plugins/datasource/azuremonitor/components/TracesQueryEditor/TracesQueryEditor.tsx
+++ b/public/app/plugins/datasource/azuremonitor/components/TracesQueryEditor/TracesQueryEditor.tsx
@ -88,37 +88,35 @@ const TracesQueryEditor = ({
  return (
    <span data-testid={selectors.components.queryEditor.tracesQueryEditor.container.input}>
      <EditorRows>
-        {query.queryType !== AzureQueryType.TraceExemplar && (
-          <EditorRow>
-            <EditorFieldGroup>
-              <ResourceField
-                query={query}
-                datasource={datasource}
-                subscriptionId={subscriptionId}
-                variableOptionGroup={variableOptionGroup}
-                onQueryChange={onChange}
-                setError={setError}
-                selectableEntryTypes={[
-                  ResourceRowType.Subscription,
-                  ResourceRowType.ResourceGroup,
-                  ResourceRowType.Resource,
-                  ResourceRowType.Variable,
-                ]}
-                resources={query.azureTraces?.resources ?? []}
-                queryType="traces"
-                disableRow={disableRow}
-                renderAdvanced={(resources, onChange) => (
-                  // It's required to cast resources because the resource picker
-                  // specifies the type to string | AzureMonitorResource.
-                  // eslint-disable-next-line
-                  <AdvancedResourcePicker resources={resources as string[]} onChange={onChange} />
-                )}
-                selectionNotice={() => 'You may only choose items of the same resource type.'}
-                range={range}
-              />
-            </EditorFieldGroup>
-          </EditorRow>
-        )}
+        <EditorRow>
+          <EditorFieldGroup>
+            <ResourceField
+              query={query}
+              datasource={datasource}
+              subscriptionId={subscriptionId}
+              variableOptionGroup={variableOptionGroup}
+              onQueryChange={onChange}
+              setError={setError}
+              selectableEntryTypes={[
+                ResourceRowType.Subscription,
+                ResourceRowType.ResourceGroup,
+                ResourceRowType.Resource,
+                ResourceRowType.Variable,
+              ]}
+              resources={query.azureTraces?.resources ?? []}
+              queryType="traces"
+              disableRow={disableRow}
+              renderAdvanced={(resources, onChange) => (
+                // It's required to cast resources because the resource picker
+                // specifies the type to string | AzureMonitorResource.
+                // eslint-disable-next-line
+                <AdvancedResourcePicker resources={resources as string[]} onChange={onChange} />
+              )}
+              selectionNotice={() => 'You may only choose items of the same resource type.'}
+              range={range}
+            />
+          </EditorFieldGroup>
+        </EditorRow>
        <EditorRow>
          <EditorFieldGroup>
            <TraceTypeField