chore: Add S3 bucket configuration function for querycomparator (#20777)

Author: benclive

tools/querycomparator/compare.go

@@ -37,7 +37,6 @@ func addCompareCommand(app *kingpin.Application) {
 	cmd.Flag("host2", "Second Loki host").Default("localhost:3102").StringVar(&host2)
 
 	cmd.Action(func(_ *kingpin.ParseContext) error {
-		storageBucket = cfg.Bucket
 		orgID = cfg.OrgID
 
 		parsed, err := parseTimeConfig(&cfg)

tools/querycomparator/config.go

@@ -47,7 +47,6 @@ func parseTimeConfig(cfg *Config) (*ParsedConfig, error) {
 
 // Global variables for bucket and org ID (used by storage functions)
 var (
-	storageBucket      string
 	orgID              string
 	indexStoragePrefix string
 	logger             log.Logger

tools/querycomparator/execute.go

@@ -16,6 +16,7 @@ import (
 	"github.com/grafana/dskit/user"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
+	"github.com/thanos-io/objstore"
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/h2c"
 
@@ -51,7 +52,6 @@ func addExecuteCommand(app *kingpin.Application) {
 	cmd.Flag("engine", "Engine version (1 or 2)").Default("2").IntVar(&engineVersion)
 
 	cmd.Action(func(_ *kingpin.ParseContext) error {
-		storageBucket = cfg.Bucket
 		orgID = cfg.OrgID
 
 		parsed, err := parseTimeConfig(&cfg)
@@ -69,22 +69,29 @@ func addExecuteCommand(app *kingpin.Application) {
 
 		switch engineVersion {
 		case 1:
-			return doExecuteLocallyV1(params)
+			return doExecuteLocallyV1(params, cfg.Bucket)
 		case 2:
-			return doExecuteLocallyV2(params)
+			return doExecuteLocallyV2(params, MustGCSDataobjBucket(cfg.Bucket))
 		default:
 			return fmt.Errorf("unsupported engine version: %d (must be 1 or 2)", engineVersion)
 		}
 	})
 }
 
+func initV2Settings() {
+	if indexStoragePrefix == "" {
+		level.Warn(logger).Log("msg", "index storage prefix is not set. using default value.")
+		indexStoragePrefix = "index/v0"
+	}
+}
+
 // doExecuteLocallyV1 executes a query using the V1 engine
-func doExecuteLocallyV1(params logql.LiteralParams) error {
+func doExecuteLocallyV1(params logql.LiteralParams, bucketName string) error {
 	if indexStoragePrefix == "" {
 		level.Warn(logger).Log("msg", "index storage prefix is not set. v1 engine may not find any chunks.")
 	}
 	level.Info(logger).Log("msg", "executing local query with V1 engine")
-	result, err := doLocalQueryWithV1Engine(params)
+	result, err := doLocalQueryWithV1Engine(params, bucketName)
 	if err != nil {
 		level.Error(logger).Log("msg", "local query with V1 engine failed", "error", err)
 		return fmt.Errorf("V1 query execution failed: %w", err)
@@ -93,29 +100,32 @@ func doExecuteLocallyV1(params logql.LiteralParams) error {
 }
 
 // doExecuteLocallyV2 executes a query using the V2 engine
-func doExecuteLocallyV2(params logql.LiteralParams) error {
+func doExecuteLocallyV2(params logql.LiteralParams, bucket objstore.Bucket) error {
+	initV2Settings()
 	level.Info(logger).Log("msg", "executing local query with V2 engine")
-	result, err := doLocalQueryWithV2Engine(params)
+	result, err := doLocalQueryWithV2Engine(params, bucket)
 	if err != nil {
 		level.Error(logger).Log("msg", "V2 query execution failed", "error", err)
 		return fmt.Errorf("V2 query execution failed: %w", err)
 	}
 	return checkResult(result)
 }
 
-func doExecuteLocallyV2Scheduler(params logql.LiteralParams) error {
+func doExecuteLocallyV2Scheduler(params logql.LiteralParams, bucket objstore.Bucket) error {
+	initV2Settings()
 	level.Info(logger).Log("msg", "executing local query with V2 engine via local scheduler and worker")
-	result, err := doLocalQueryWithV2EngineScheduler(params)
+	result, err := doLocalQueryWithV2EngineScheduler(params, bucket)
 	if err != nil {
 		level.Error(logger).Log("msg", "v2 query execution failed", "error", err)
 		return fmt.Errorf("v2 query execution failed: %w", err)
 	}
 	return checkResult(result)
 }
 
-func doExecuteLocallyV2SchedulerRemote(params logql.LiteralParams) error {
+func doExecuteLocallyV2SchedulerRemote(params logql.LiteralParams, bucket objstore.Bucket) error {
+	initV2Settings()
 	level.Info(logger).Log("msg", "executing local query with V2 engine via remote scheduler and worker")
-	result, err := doLocalQueryWithV2EngineSchedulerRemote(params)
+	result, err := doLocalQueryWithV2EngineSchedulerRemote(params, bucket)
 	if err != nil {
 		level.Error(logger).Log("msg", "v2 query execution failed", "error", err)
 		return fmt.Errorf("v2 query execution failed: %w", err)
@@ -138,23 +148,23 @@ func checkResult(result logqlmodel.Result) error {
 }
 
 // doLocalQueryWithV2Engine executes a query using the V2 engine
-func doLocalQueryWithV2Engine(params logql.LiteralParams) (logqlmodel.Result, error) {
+func doLocalQueryWithV2Engine(params logql.LiteralParams, bucket objstore.Bucket) (logqlmodel.Result, error) {
 	logger := glog.NewLogfmtLogger(os.Stderr)
 	ms := metastore.NewObjectMetastore(
-		MustDataobjBucket(),
+		bucket,
 		metastore.Config{IndexStoragePrefix: "index/v0"},
 		logger,
 		metastore.NewObjectMetastoreMetrics(prometheus.DefaultRegisterer),
 	)
 
 	ctx := user.InjectOrgID(context.Background(), orgID)
-	qe := engine.NewBasic(engine.ExecutorConfig{BatchSize: 512}, ms, MustDataobjBucket(), logql.NoLimits, prometheus.DefaultRegisterer, logger)
+	qe := engine.NewBasic(engine.ExecutorConfig{BatchSize: 512}, ms, bucket, logql.NoLimits, prometheus.DefaultRegisterer, logger)
 	query := qe.Query(params)
 	return query.Exec(ctx)
 }
 
 // doLocalQueryWithV1Engine executes a query using the V1 engine
-func doLocalQueryWithV1Engine(params logql.LiteralParams) (logqlmodel.Result, error) {
+func doLocalQueryWithV1Engine(params logql.LiteralParams, bucketName string) (logqlmodel.Result, error) {
 	ctx := user.InjectOrgID(context.Background(), orgID)
 
 	l := &validation.Limits{}
@@ -175,7 +185,7 @@ func doLocalQueryWithV1Engine(params logql.LiteralParams) (logqlmodel.Result, er
 		ObjectStore: bucket.ConfigWithNamedStores{
 			Config: bucket.Config{
 				GCS: gcs.Config{
-					BucketName: storageBucket,
+					BucketName: bucketName,
 				},
 			},
 		},
@@ -226,7 +236,7 @@ func doLocalQueryWithV1Engine(params logql.LiteralParams) (logqlmodel.Result, er
 	return query.Exec(ctx)
 }
 
-func doLocalQueryWithV2EngineScheduler(params logql.LiteralParams) (logqlmodel.Result, error) {
+func doLocalQueryWithV2EngineScheduler(params logql.LiteralParams, bucket objstore.Bucket) (logqlmodel.Result, error) {
 	ctx := user.InjectOrgID(context.Background(), orgID)
 
 	sched, err := engine.NewScheduler(engine.SchedulerParams{
@@ -239,16 +249,14 @@ func doLocalQueryWithV2EngineScheduler(params logql.LiteralParams) (logqlmodel.R
 		return logqlmodel.Result{}, fmt.Errorf("starting scheduler service: %w", err)
 	}
 
-	b := MustDataobjBucket()
-
 	metastoreMetrics := metastore.NewObjectMetastoreMetrics(prometheus.DefaultRegisterer)
 	msConfig := metastore.Config{IndexStoragePrefix: "index/v0"}
 	workerLogger := glog.With(logger, "component", "worker")
 	worker, err := engine.NewWorker(engine.WorkerParams{
 		Logger:         workerLogger,
 		AdvertiseAddr:  nil,
-		Bucket:         b,
-		Metastore:      metastore.NewObjectMetastore(b, msConfig, workerLogger, metastoreMetrics),
+		Bucket:         bucket,
+		Metastore:      metastore.NewObjectMetastore(bucket, msConfig, workerLogger, metastoreMetrics),
 		LocalScheduler: sched,
 		Config: engine.WorkerConfig{
 			SchedulerLookupAddress:  "",
@@ -274,7 +282,7 @@ func doLocalQueryWithV2EngineScheduler(params logql.LiteralParams) (logqlmodel.R
 				BatchSize: 128,
 			},
 		},
-		Metastore: metastore.NewObjectMetastore(b, msConfig, engineLogger, metastoreMetrics),
+		Metastore: metastore.NewObjectMetastore(bucket, msConfig, engineLogger, metastoreMetrics),
 		Scheduler: sched,
 		Limits:    logql.NoLimits,
 	})
@@ -285,7 +293,8 @@ func doLocalQueryWithV2EngineScheduler(params logql.LiteralParams) (logqlmodel.R
 	return e.Execute(ctx, params)
 }
 
-func doLocalQueryWithV2EngineSchedulerRemote(params logql.LiteralParams) (logqlmodel.Result, error) {
+// doLocalQueryWithV2EngineSchedulerRemote executes a query using the V2 engine via remote scheduler and workers so it also executes the serialization logic.
+func doLocalQueryWithV2EngineSchedulerRemote(params logql.LiteralParams, bucket objstore.Bucket) (logqlmodel.Result, error) {
 	ctx := user.InjectOrgID(context.Background(), orgID)
 
 	schedSrv, schedSvc, err := newServerService("scheduler", 3101, logger, prometheus.NewRegistry())
@@ -313,19 +322,18 @@ func doLocalQueryWithV2EngineSchedulerRemote(params logql.LiteralParams) (logqlm
 		return logqlmodel.Result{}, fmt.Errorf("starting worker service: %w", err)
 	}
 
-	b := MustDataobjBucket()
 	metastoreMetrics := metastore.NewObjectMetastoreMetrics(prometheus.DefaultRegisterer)
-	msConfig := metastore.Config{IndexStoragePrefix: "index/v0"}
+	msConfig := metastore.Config{IndexStoragePrefix: indexStoragePrefix}
 	workerLogger := glog.With(logger, "component", "worker")
 	worker, err := engine.NewWorker(engine.WorkerParams{
 		Logger:         workerLogger,
 		AdvertiseAddr:  workerSrv.HTTPListenAddr(),
-		Bucket:         b,
-		Metastore:      metastore.NewObjectMetastore(b, msConfig, workerLogger, metastoreMetrics),
+		Bucket:         bucket,
+		Metastore:      metastore.NewObjectMetastore(bucket, msConfig, workerLogger, metastoreMetrics),
 		LocalScheduler: nil,
 		Config: engine.WorkerConfig{
 			SchedulerLookupAddress:  schedSrv.HTTPListenAddr().String(),
-			SchedulerLookupInterval: 60,
+			SchedulerLookupInterval: time.Minute,
 			WorkerThreads:           64,
 		},
 		Executor: engine.ExecutorConfig{
@@ -348,7 +356,7 @@ func doLocalQueryWithV2EngineSchedulerRemote(params logql.LiteralParams) (logqlm
 				BatchSize: 128,
 			},
 		},
-		Metastore: metastore.NewObjectMetastore(b, msConfig, engineLogger, metastoreMetrics),
+		Metastore: metastore.NewObjectMetastore(bucket, msConfig, engineLogger, metastoreMetrics),
 		Scheduler: sched,
 		Limits:    logql.NoLimits,
 	})

tools/querycomparator/main_test.go

@@ -9,22 +9,21 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestDebugCmd(t *testing.T) {
+func TestQueryComparatorCmd(t *testing.T) {
 	t.Skip("Test for debugging purposes only")
 
-	storageBucket = ""      // TODO: set bucket name for local engine support
-	indexStoragePrefix = "" // TODO: set index storage prefix for local engine support
-	orgID = ""              // TODO: set org ID for querying remote instances
+	bucket := MustGCSDataobjBucket("")
+	orgID = "" // TODO: set org ID for querying remote instances
 
-	start := time.Date(2025, 10, 24, 0, 0, 0, 0, time.UTC)
-	end := start.Add(45 * time.Minute)
-	query := "{namespace=\"test\"}"
+	start := time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC)
+	end := start.Add(1 * time.Minute)
+	query := `{container="distributor"}`
 
 	params, err := logql.NewLiteralParams(query, start, end, 0, 0, logproto.BACKWARD, 1000, nil, nil)
 	require.NoError(t, err)
 
 	// Run subcommands as necessary
 	require.NoError(t, doComparison(params, "localhost:3101", "localhost:3102"))
-	require.NoError(t, queryMetastore(params))
-	require.NoError(t, doExecuteLocallyV2(params))
+	require.NoError(t, queryMetastore(params, bucket))
+	require.NoError(t, doExecuteLocallyV2SchedulerRemote(params, bucket))
 }

tools/querycomparator/metastore.go

@@ -12,6 +12,7 @@ import (
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/user"
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/thanos-io/objstore"
 
 	"github.com/grafana/loki/v3/pkg/dataobj/metastore"
 	"github.com/grafana/loki/v3/pkg/logproto"
@@ -31,33 +32,34 @@ func addMetastoreCommand(app *kingpin.Application) {
 	cmd.Flag("query", "LogQL query to analyze").Required().StringVar(&cfg.Query)
 
 	cmd.Action(func(_ *kingpin.ParseContext) error {
-		storageBucket = cfg.Bucket
 		orgID = cfg.OrgID
 
 		parsed, err := parseTimeConfig(&cfg)
 		if err != nil {
 			return err
 		}
 
+		bucket := MustGCSDataobjBucket(cfg.Bucket)
+
 		params, err := logql.NewLiteralParams(cfg.Query, parsed.StartTime, parsed.EndTime, 0, 0, logproto.BACKWARD, 10, nil, nil)
 		if err != nil {
 			return err
 		}
 
-		return queryMetastore(params)
+		return queryMetastore(params, bucket)
 	})
 }
 
 // queryMetastore queries the metastore for stream sections
-func queryMetastore(params logql.LiteralParams) error {
+func queryMetastore(params logql.LiteralParams, bucket objstore.Bucket) error {
 	query := params.QueryString()
 	closeIdx := strings.Index(query, "}")
 	streamMatchers, err := syntax.ParseMatchers(query[:closeIdx+1], true)
 	if err != nil {
 		return err
 	}
 
-	sections, err := getSections(params.Start(), params.End(), streamMatchers)
+	sections, err := getSections(bucket, params.Start(), params.End(), streamMatchers)
 	if err != nil {
 		return err
 	}
@@ -70,11 +72,11 @@ func queryMetastore(params logql.LiteralParams) error {
 
 // getSections queries the metastore for dataobject sections matching the query selector
 // Currently, it does not pass structured metadata predicates
-func getSections(start, end time.Time, streamMatchers []*labels.Matcher) ([]*metastore.DataobjSectionDescriptor, error) {
+func getSections(bucket objstore.Bucket, start, end time.Time, streamMatchers []*labels.Matcher) ([]*metastore.DataobjSectionDescriptor, error) {
 	ctx := user.InjectOrgID(context.Background(), orgID)
 	ms := metastore.NewObjectMetastore(
-		MustDataobjBucket(),
-		metastore.Config{IndexStoragePrefix: "index/v0"},
+		bucket,
+		metastore.Config{IndexStoragePrefix: indexStoragePrefix},
 		log.NewLogfmtLogger(os.Stderr),
 		metastore.NewObjectMetastoreMetrics(nil),
 	)

tools/querycomparator/storage.go

@@ -5,27 +5,56 @@ import (
 	"log"
 
 	glog "github.com/go-kit/log"
+	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/loki/v3/pkg/storage/bucket/gcs"
+	"github.com/grafana/loki/v3/pkg/storage/bucket/s3"
 	"github.com/thanos-io/objstore"
 )
 
-// MustDataobjBucket creates a GCS bucket client for dataobj storage
-func MustDataobjBucket() objstore.Bucket {
+// MustS3DataobjBucket creates a S3 bucket client for dataobj storage
+// The access key id, secret access key, and session token are required for S3 dataobj bucket and must be provided.
+// The region endpoint follows the format "s3.<aws region name>.amazonaws.com" e.g. "s3.eu-south-2.amazonaws.com".
+func MustS3DataobjBucket(bucketName string, regionEndpoint string) objstore.Bucket {
+	accessKeyID := ""
+	secretAccessKey := ""
+	sessionToken := ""
+
+	if accessKeyID == "" || secretAccessKey == "" || sessionToken == "" {
+		log.Fatal("access key id, secret access key, and session token are required for S3 dataobj bucket")
+	}
+
+	bkt, err := s3.NewBucketClient(s3.Config{
+		Endpoint:        regionEndpoint,
+		BucketName:      bucketName,
+		AccessKeyID:     accessKeyID,
+		SecretAccessKey: flagext.SecretWithValue(secretAccessKey),
+		SessionToken:    flagext.SecretWithValue(sessionToken),
+	}, "querycomparator", glog.NewNopLogger(), nil)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	prefixedBkt := objstore.NewPrefixedBucket(bkt, "dataobj")
+	return prefixedBkt
+}
+
+// MustGCSDataobjBucket creates a GCS bucket client for dataobj storage
+func MustGCSDataobjBucket(bucketName string) objstore.Bucket {
 	bkt, err := gcs.NewBucketClient(context.Background(), gcs.Config{
-		BucketName: storageBucket,
-	}, "testing", glog.NewNopLogger(), nil)
+		BucketName: bucketName,
+	}, "querycomparator", glog.NewNopLogger(), nil)
 	if err != nil {
 		log.Fatal(err)
 	}
 	objBucket := objstore.NewPrefixedBucket(bkt, "dataobj")
 	return objBucket
 }
 
-// MustRawBucket creates a GCS bucket client for raw storage
-func MustRawBucket() objstore.Bucket {
+// MustRawGCSBucket creates a GCS bucket client for raw storage
+func MustRawGCSBucket(bucketName string) objstore.Bucket {
 	bkt, err := gcs.NewBucketClient(context.Background(), gcs.Config{
-		BucketName: storageBucket,
-	}, "testing", glog.NewNopLogger(), nil)
+		BucketName: bucketName,
+	}, "querycomparator", glog.NewNopLogger(), nil)
 	if err != nil {
 		log.Fatal(err)
 	}