From d1760760780b99e5bd0ba57e38a52ad28c3eb5e4 Mon Sep 17 00:00:00 2001
From: Adam Yeats <16296989+adamyeats@users.noreply.github.com>
Date: Wed, 11 Dec 2024 13:02:51 +0000
Subject: [PATCH] Azure Monitor: Add safety around usage of frame.Meta.Custom
 struct (#97766)

---
 .../azure-log-analytics-datasource.go         | 53 ++++++++++++-------
 1 file changed, 34 insertions(+), 19 deletions(-)

diff --git a/pkg/tsdb/azuremonitor/loganalytics/azure-log-analytics-datasource.go b/pkg/tsdb/azuremonitor/loganalytics/azure-log-analytics-datasource.go
index dee69d9fe5c..93db9a8ebbd 100644
--- a/pkg/tsdb/azuremonitor/loganalytics/azure-log-analytics-datasource.go
+++ b/pkg/tsdb/azuremonitor/loganalytics/azure-log-analytics-datasource.go
@@ -320,14 +320,22 @@ func (e *AzureLogAnalyticsDatasource) executeQuery(ctx context.Context, query *A
 	}
 
 	logLimitDisabled := backend.GrafanaConfigFromContext(ctx).FeatureToggles().IsEnabled("azureMonitorDisableLogLimit")
+
 	frame, err := ResponseTableToFrame(t, query.RefID, query.Query, query.QueryType, query.ResultFormat, logLimitDisabled)
 	if err != nil {
 		return nil, err
 	}
+
 	frame = appendErrorNotice(frame, logResponse.Error)
 	if frame == nil {
-		dataResponse := backend.DataResponse{}
-		return &dataResponse, nil
+		return &backend.DataResponse{}, nil
+	}
+
+	// Ensure Meta.Custom is initialized
+	if frame.Meta.Custom == nil {
+		frame.Meta.Custom = &LogAnalyticsMeta{
+			ColumnTypes: make([]string, 0),
+		}
 	}
 
 	queryUrl, err := getQueryUrl(query.Query, query.Resources, dsInfo.Routes["Azure Portal"].URL, query.TimeRange)
@@ -335,30 +343,37 @@ func (e *AzureLogAnalyticsDatasource) executeQuery(ctx context.Context, query *A
 		return nil, err
 	}
 
-	if (query.QueryType == dataquery.AzureQueryTypeAzureTraces || query.QueryType == dataquery.AzureQueryTypeTraceql) && query.ResultFormat == dataquery.ResultFormatTrace {
-		frame.Meta.PreferredVisualization = data.VisTypeTrace
-	}
-
-	if query.ResultFormat == dataquery.ResultFormatTable {
-		frame.Meta.PreferredVisualization = data.VisTypeTable
-	}
-
-	if query.ResultFormat == dataquery.ResultFormatLogs {
-		frame.Meta.PreferredVisualization = data.VisTypeLogs
-		frame.Meta.Custom = &LogAnalyticsMeta{
-			ColumnTypes:     frame.Meta.Custom.(*LogAnalyticsMeta).ColumnTypes,
-			AzurePortalLink: queryUrl,
+	// Set the preferred visualization
+	switch query.ResultFormat {
+	case dataquery.ResultFormatTrace:
+		if query.QueryType == dataquery.AzureQueryTypeAzureTraces || query.QueryType == dataquery.AzureQueryTypeTraceql {
+			frame.Meta.PreferredVisualization = data.VisTypeTrace
 		}
-	}
-
-	if query.ResultFormat == dataquery.ResultFormatTimeSeries {
+	case dataquery.ResultFormatTable:
+		frame.Meta.PreferredVisualization = data.VisTypeTable
+	case dataquery.ResultFormatLogs:
+		frame.Meta.PreferredVisualization = data.VisTypeLogs
+		if logMeta, ok := frame.Meta.Custom.(*LogAnalyticsMeta); ok {
+			frame.Meta.Custom = &LogAnalyticsMeta{
+				ColumnTypes:     logMeta.ColumnTypes,
+				AzurePortalLink: queryUrl,
+			}
+		} else {
+			frame.Meta.Custom = &LogAnalyticsMeta{
+				AzurePortalLink: queryUrl,
+			}
+		}
+	case dataquery.ResultFormatTimeSeries:
 		tsSchema := frame.TimeSeriesSchema()
 		if tsSchema.Type == data.TimeSeriesTypeLong {
 			wideFrame, err := data.LongToWide(frame, nil)
 			if err == nil {
 				frame = wideFrame
 			} else {
-				frame.AppendNotices(data.Notice{Severity: data.NoticeSeverityWarning, Text: "could not convert frame to time series, returning raw table: " + err.Error()})
+				frame.AppendNotices(data.Notice{
+					Severity: data.NoticeSeverityWarning,
+					Text:     "could not convert frame to time series, returning raw table: " + err.Error(),
+				})
 			}
 		}
 	}