Use YAML config instead

Signed-off-by: Dimitar Dimitrov <dimitar.dimitrov@grafana.com>

CHANGELOG.md

@@ -140,7 +140,7 @@
 ### Query-tee
 
 * [FEATURE] Added `-proxy.compare-skip-samples-before` to skip samples before the given time when comparing responses. The time can be in RFC3339 format (or) RFC3339 without the timezone and seconds (or) date only. #9515
-* [FEATURE] Add `-backend.config` JSON configuration for per-backend options. Currently, it only supports additional HTTP request headers. #10081
+* [FEATURE] Add `-backend.config-file` for a YAML configuration file for per-backend options. Currently, it only supports additional HTTP request headers. #10081
 * [ENHANCEMENT] Added human-readable timestamps to comparison failure messages. #9665
 
 ### Documentation

tools/querytee/proxy.go

@@ -12,6 +12,7 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -24,6 +25,7 @@ import (
 	"github.com/grafana/dskit/spanlogger"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
+	"gopkg.in/yaml.v3"
 )
 
 type ProxyConfig struct {
@@ -35,7 +37,7 @@ type ProxyConfig struct {
 	BackendEndpoints                    string
 	PreferredBackend                    string
 	BackendReadTimeout                  time.Duration
-	BackendConfigStr                    string
+	BackendConfigFile                   string
 	parsedBackendConfig                 map[string]*BackendConfig
 	CompareResponses                    bool
 	LogSlowQueryResponseThreshold       time.Duration
@@ -51,10 +53,10 @@ type ProxyConfig struct {
 }
 
 type BackendConfig struct {
-	RequestHeaders http.Header `json:"request_headers"`
+	RequestHeaders http.Header `json:"request_headers" yaml:"request_headers"`
 }
 
-func exampleBackendConfig() string {
+func exampleJSONBackendConfig() string {
 	cfg := BackendConfig{
 		RequestHeaders: http.Header{
 			"Cache-Control": {"no-store"},
@@ -82,7 +84,7 @@ func (cfg *ProxyConfig) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&cfg.BackendSkipTLSVerify, "backend.skip-tls-verify", false, "Skip TLS verification on backend targets.")
 	f.StringVar(&cfg.PreferredBackend, "backend.preferred", "", "The hostname of the preferred backend when selecting the response to send back to the client. If no preferred backend is configured then the query-tee will send back to the client the first successful response received without waiting for other backends.")
 	f.DurationVar(&cfg.BackendReadTimeout, "backend.read-timeout", 150*time.Second, "The timeout when reading the response from a backend.")
-	f.StringVar(&cfg.BackendConfigStr, "backend.config", "{}", "JSON object with backend configuration. Each key is the backend hostname. This is an example value: "+exampleBackendConfig())
+	f.StringVar(&cfg.BackendConfigFile, "backend.config-file", "", "Path to a file with YAML or JSON configuration for each backend. Each key in the YAML/JSON document is a backend hostname. This is an example configuration value for a backend in JSON: "+exampleJSONBackendConfig())
 	f.BoolVar(&cfg.CompareResponses, "proxy.compare-responses", false, "Compare responses between preferred and secondary endpoints for supported routes.")
 	f.DurationVar(&cfg.LogSlowQueryResponseThreshold, "proxy.log-slow-query-response-threshold", 10*time.Second, "The minimum difference in response time between slowest and fastest back-end over which to log the query. 0 to disable.")
 	f.Float64Var(&cfg.ValueComparisonTolerance, "proxy.value-comparison-tolerance", 0.000001, "The tolerance to apply when comparing floating point values in the responses. 0 to disable tolerance and require exact match (not recommended).")
@@ -140,10 +142,14 @@ func NewProxy(cfg ProxyConfig, logger log.Logger, routes []Route, registerer pro
 		return nil, errors.New("preferred backend must be set when secondary backends request proportion is not 1")
 	}
 
-	if len(cfg.BackendConfigStr) > 0 {
-		err := json.Unmarshal([]byte(cfg.BackendConfigStr), &cfg.parsedBackendConfig)
+	if len(cfg.BackendConfigFile) > 0 {
+		configBytes, err := os.ReadFile(cfg.BackendConfigFile)
 		if err != nil {
-			return nil, fmt.Errorf("failed to parse backend JSON config: %w", err)
+			return nil, fmt.Errorf("failed to read backend config file (%s): %w", cfg.BackendConfigFile, err)
+		}
+		err = yaml.Unmarshal(configBytes, &cfg.parsedBackendConfig)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse backend YAML config: %w", err)
 		}
 	}
 

tools/querytee/proxy_test.go

@@ -11,6 +11,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"strconv"
 	"strings"
 	"testing"
@@ -696,6 +697,101 @@ func TestProxyHTTPGRPC(t *testing.T) {
 	})
 }
 
+func Test_NewProxy_BackendConfigPath(t *testing.T) {
+	// Helper to create a temporary file with content
+	createTempFile := func(t *testing.T, content string) string {
+		tmpfile, err := os.CreateTemp("", "backend-config-*.yaml")
+		require.NoError(t, err)
+
+		defer tmpfile.Close()
+
+		_, err = tmpfile.Write([]byte(content))
+		require.NoError(t, err)
+
+		return tmpfile.Name()
+	}
+
+	tests := map[string]struct {
+		configContent  string
+		createFile     bool
+		expectedError  string
+		expectedConfig map[string]*BackendConfig
+	}{
+		"missing file": {
+			createFile:    false,
+			expectedError: "failed to read backend config file (/nonexistent/path): open /nonexistent/path: no such file or directory",
+		},
+		"empty file": {
+			createFile:     true,
+			configContent:  "",
+			expectedConfig: map[string]*BackendConfig(nil),
+		},
+		"invalid YAML structure (not a map)": {
+			createFile:    true,
+			configContent: "- item1\n- item2",
+			expectedError: "failed to parse backend YAML config:",
+		},
+		"valid configuration": {
+			createFile: true,
+			configContent: `
+              backend1:
+                request_headers:
+                  X-Custom-Header: ["value1", "value2"]
+                  Cache-Control: ["no-store"]
+              backend2:
+                request_headers:
+                  Authorization: ["Bearer token123"]
+            `,
+			expectedConfig: map[string]*BackendConfig{
+				"backend1": {
+					RequestHeaders: http.Header{
+						"X-Custom-Header": {"value1", "value2"},
+						"Cache-Control":   {"no-store"},
+					},
+				},
+				"backend2": {
+					RequestHeaders: http.Header{
+						"Authorization": {"Bearer token123"},
+					},
+				},
+			},
+		},
+	}
+
+	for testName, testCase := range tests {
+		t.Run(testName, func(t *testing.T) {
+			// Base config that's valid except for the backend config path
+			cfg := ProxyConfig{
+				BackendEndpoints:                   "http://backend1:9090,http://backend2:9090",
+				ServerHTTPServiceAddress:           "localhost",
+				ServerHTTPServicePort:              0,
+				ServerGRPCServiceAddress:           "localhost",
+				ServerGRPCServicePort:              0,
+				SecondaryBackendsRequestProportion: 1.0,
+			}
+
+			if !testCase.createFile {
+				cfg.BackendConfigFile = "/nonexistent/path"
+			} else {
+				tmpPath := createTempFile(t, testCase.configContent)
+				cfg.BackendConfigFile = tmpPath
+				defer os.Remove(tmpPath)
+			}
+
+			p, err := NewProxy(cfg, log.NewNopLogger(), testRoutes, nil)
+
+			if testCase.expectedError != "" {
+				assert.ErrorContains(t, err, testCase.expectedError)
+				assert.Nil(t, p)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, p)
+				assert.Equal(t, testCase.expectedConfig, p.cfg.parsedBackendConfig)
+			}
+		})
+	}
+}
+
 func mockQueryResponse(path string, status int, res string) http.HandlerFunc {
 	return mockQueryResponseWithExpectedBody(path, "", status, res)
 }