Add gRPC broker support and examples (#57)

Add gRPC client transport, protobuf definitions, compression tests, and a
broker example alongside docs and config wiring.

Add gRPC transport and proto tests, plus injectable helpers to cover

Expose the broker gRPC port, add gRPC integration tests with compression
checks, and pass the port through CI integration runs.

Author: xiangfu0

.github/workflows/tests.yml

@@ -83,20 +83,29 @@ jobs:
       - name: Set up Pinot Cluster
         run: make run-pinot-dist
         env:
-          PINOT_VERSION: 1.0.0
+          PINOT_VERSION: 1.4.0
           PINOT_HOME: /tmp/pinot
           CONTROLLER_PORT_FORWARD: 9000
           BROKER_PORT_FORWARD: 8000
-
-      - name: GORM Example
-        run: go run ./examples/gorm-example
+          BROKER_GRPC_PORT_FORWARD: 8010
+
+      - name: All Examples
+        run: |
+          set -euo pipefail
+          go run ./examples/batch-quickstart
+          go run ./examples/gorm-example
+          go run ./examples/grpc-broker-client
+          go run ./examples/json-batch-quickstart
+          go run ./examples/multistage-quickstart
+          go run ./examples/prepared-statement-example
 
       - name: Integration Test
         run: make integration-test
         env:
           ZOOKEEPER_PORT: 2123
           CONTROLLER_PORT: 9000
           BROKER_PORT: 8000
+          BROKER_GRPC_PORT: 8010
 
       - name: Publish to coveralls.io
         uses: coverallsapp/github-action@v2

AGENTS.md

@@ -0,0 +1,38 @@
+# AGENTS.md
+
+Repo-specific instructions for coding agents working on this project.
+
+## Repo layout
+- `pinot/`: core client library (HTTP + gRPC).
+- `gormpinot/`: GORM dialector for Pinot.
+- `examples/`: runnable examples (not part of default tests).
+- `integration-tests/`: integration tests that require a running Pinot quickstart.
+- `scripts/`: helper scripts to start/stop Pinot quickstart.
+
+## Tooling & style
+- Go version is defined in `go.mod`; use that version for builds/tests.
+- Run `gofmt` on modified Go files and keep `go vet` clean.
+- Update `go.mod`/`go.sum` only when dependencies change (run `go mod tidy` if needed).
+
+## Commit requirements
+- Enable repo hooks with `make hooks` (uses `.githooks/pre-commit`).
+- Pre-commit runs `make lint`, `make test`, and `scripts/check-coverage.sh`.
+- Update the coverage baseline with `make coverage-baseline` after legitimate coverage improvements.
+
+## Common commands (from `Makefile`)
+- Build: `make build` (or `go build ./...`).
+- Unit tests: `make test` (excludes `examples/` and `integration-tests/`).
+- Lint: `make lint` (gofmt, go vet, golangci-lint v2).
+- Integration tests: `make integration-test` (requires Pinot running).
+- Start Pinot quickstart (local dist): `make run-pinot-dist`.
+- Stop Pinot quickstart: `make stop-pinot-dist`.
+- Start Pinot quickstart (Docker): `make run-pinot-docker`.
+- Stop Pinot quickstart (Docker): `make stop-pinot-docker`.
+
+## Integration test notes
+- gRPC tests read `BROKER_GRPC_HOST` and `BROKER_GRPC_PORT` (defaults: `127.0.0.1:8010`).
+- The quickstart script supports `PINOT_VERSION`, `PINOT_HOME`, `BROKER_PORT_FORWARD`, `BROKER_GRPC_PORT_FORWARD`.
+- `scripts/start-pinot-quickstart.sh` uses `curl` and `jq` to verify the cluster is ready.
+
+## When changing public APIs
+- Update examples and tests that reference exported types or behavior changes.

Makefile

@@ -49,7 +49,7 @@ stop-pinot-dist:
 
 .PHONY: run-pinot-docker
 run-pinot-docker:
-	docker run --name pinot-quickstart -p 2123:2123 -p 9000:9000 -p 8000:8000 apachepinot/pinot:latest QuickStart -type MULTI_STAGE
+	docker run --name pinot-quickstart -p 2123:2123 -p 9000:9000 -p 8000:8000 -p 8010:8010 apachepinot/pinot:latest QuickStart -type MULTI_STAGE
 
 .PHONY: stop-pinot-docker
 stop-pinot-docker:
@@ -59,3 +59,15 @@ stop-pinot-docker:
 .PHONY: integration-test
 integration-test: build
 	go test -timeout 500s -v -race -covermode atomic -coverprofile=coverage.out $(INTEGRATION_TESTS_PACKAGES)
+
+.PHONY: coverage-check
+coverage-check: test
+	./scripts/check-coverage.sh
+
+.PHONY: coverage-baseline
+coverage-baseline: test
+	./scripts/update-coverage-baseline.sh
+
+.PHONY: hooks
+hooks:
+	git config core.hooksPath .githooks

README.md

@@ -1,6 +1,6 @@
 # Pinot Client GO
 
-[![Go 1.19](https://img.shields.io/badge/go-1.19-blue.svg)](https://golang.org/dl/#go1.19)
+[![Go 1.24](https://img.shields.io/badge/go-1.24-blue.svg)](https://golang.org/dl/#go1.24)
 [![GoDoc](https://img.shields.io/badge/go.dev-reference-007d9c?logo=go&logoColor=white)](https://pkg.go.dev/github.com/startreedata/pinot-client-go)
 [![Build Status](https://github.com/startreedata/pinot-client-go/actions/workflows/tests.yml/badge.svg)](https://github.com/startreedata/pinot-client-go/actions/workflows/tests.yml)
 [![Coverage Status](https://coveralls.io/repos/github/startreedata/pinot-client-go/badge.svg?branch=master)](https://coveralls.io/github/startreedata/pinot-client-go?branch=master)
@@ -149,6 +149,26 @@ pinotClient, err := pinot.NewWithConfig(&pinot.ClientConfig{
 })
 ```
 
+### gRPC broker queries
+
+Pinot brokers expose a gRPC query endpoint when configured with `pinot.broker.grpc.port`. See the Pinot gRPC broker API docs for configuration details and TLS setup: https://docs.pinot.apache.org/users/api/broker-grpc-api  
+
+```go
+pinotClient, err := pinot.NewWithConfig(&pinot.ClientConfig{
+	BrokerList: []string{"localhost:8010"},
+	GrpcConfig: &pinot.GrpcConfig{
+		Encoding:     "JSON",  // or "ARROW"
+		Compression:  "ZSTD",
+		BlockRowSize: 10000,
+		Timeout:      5 * time.Second,
+		// TLSConfig: &pinot.GrpcTLSConfig{
+		// 	Enabled:    true,
+		// 	CACertPath: "/path/to/ca.pem",
+		// },
+	},
+})
+```
+
 ## Query Pinot
 
 Please see this [example](https://github.com/startreedata/pinot-client-go/blob/master/examples/batch-quickstart/main.go) for your reference.
@@ -534,4 +554,3 @@ Go to [Github Release](https://github.com/startreedata/pinot-client-go/releases)
 The published Release will be available in [Go Modules](https://pkg.go.dev/github.com/startreedata/pinot-client-go).
 
 If not available, go to the corresponding new version page (https://pkg.go.dev/github.com/startreedata/[email protected]) and click on the "Request New Version" button.
-

examples/grpc-broker-client/main.go

@@ -0,0 +1,62 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	pinot "github.com/startreedata/pinot-client-go/pinot"
+)
+
+func main() {
+	pinotClient, err := pinot.NewWithConfig(&pinot.ClientConfig{
+		BrokerList: []string{"localhost:8010"},
+		GrpcConfig: &pinot.GrpcConfig{
+			Encoding:     "JSON",
+			Compression:  "ZSTD",
+			BlockRowSize: 10000,
+			Timeout:      5 * time.Second,
+			// TLSConfig: &pinot.GrpcTLSConfig{
+			// 	Enabled:    true,
+			// 	CACertPath: "/path/to/ca.pem",
+			// },
+		},
+	})
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	table := "baseballStats"
+	query := "select * from baseballStats where teamID = 'OAK' and yearID = 2004 order by homeRuns desc limit 5"
+	brokerResp, err := pinotClient.ExecuteSQL(table, query)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	printBrokerResp(brokerResp)
+}
+
+func printBrokerResp(brokerResp *pinot.BrokerResponse) {
+	log.Infof("Query Stats: response time - %d ms, scanned docs - %d, total docs - %d", brokerResp.TimeUsedMs, brokerResp.NumDocsScanned, brokerResp.TotalDocs)
+	if len(brokerResp.Exceptions) > 0 {
+		jsonBytes, _ := json.Marshal(brokerResp.Exceptions)
+		log.Infof("brokerResp.Exceptions:\n%s\n", jsonBytes)
+		return
+	}
+	if brokerResp.ResultTable != nil {
+		jsonBytes, _ := json.Marshal(brokerResp.ResultTable)
+		log.Infof("brokerResp.ResultTable:\n%s\n", jsonBytes)
+		line := ""
+		for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ {
+			line += fmt.Sprintf("%s(%s)\t", brokerResp.ResultTable.GetColumnName(c), brokerResp.ResultTable.GetColumnDataType(c))
+		}
+		line += "\n"
+		for r := 0; r < brokerResp.ResultTable.GetRowCount(); r++ {
+			for c := 0; c < brokerResp.ResultTable.GetColumnCount(); c++ {
+				line += fmt.Sprintf("%v\t", brokerResp.ResultTable.Get(r, c))
+			}
+			line += "\n"
+		}
+		log.Infof("ResultTable:\n%s", line)
+	}
+}

go.mod

@@ -1,24 +1,44 @@
 module github.com/startreedata/pinot-client-go
 
 require (
+	github.com/apache/arrow/go/v15 v15.0.2
 	github.com/go-zookeeper/zk v1.0.4
+	github.com/golang/snappy v1.0.0
+	github.com/klauspost/compress v1.18.3
+	github.com/pierrec/lz4/v4 v4.1.25
 	github.com/sirupsen/logrus v1.9.4
 	github.com/stretchr/testify v1.11.1
+	google.golang.org/grpc v1.78.0
+	google.golang.org/protobuf v1.36.11
 )
 
 require (
+	github.com/goccy/go-json v0.10.2 // indirect
+	github.com/google/flatbuffers v23.5.26+incompatible // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
-	golang.org/x/text v0.20.0 // indirect
+	github.com/klauspost/cpuid/v2 v2.2.5 // indirect
+	github.com/zeebo/xxh3 v1.0.2 // indirect
+	golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
+	golang.org/x/mod v0.29.0 // indirect
+	golang.org/x/net v0.47.0 // indirect
+	golang.org/x/sync v0.18.0 // indirect
+	golang.org/x/telemetry v0.0.0-20251008203120-078029d740a8 // indirect
+	golang.org/x/text v0.31.0 // indirect
+	golang.org/x/tools v0.38.0 // indirect
+	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20251029180050-ab9386a59fda // indirect
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
-	golang.org/x/sys v0.13.0 // indirect
+	golang.org/x/sys v0.38.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	gorm.io/gorm v1.31.1
 )
 
-go 1.19
+go 1.24.0
+
+toolchain go1.24.4

go.sum

@@ -1,11 +1,35 @@
+github.com/apache/arrow/go/v15 v15.0.2 h1:60IliRbiyTWCWjERBCkO1W4Qun9svcYoZrSLcyOsMLE=
+github.com/apache/arrow/go/v15 v15.0.2/go.mod h1:DGXsR3ajT524njufqf95822i+KTh+yea1jass9YXgjA=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
+github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
+github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
+github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
 github.com/go-zookeeper/zk v1.0.4 h1:DPzxraQx7OrPyXq2phlGlNSIyWEsAox0RJmjTseMV6I=
 github.com/go-zookeeper/zk v1.0.4/go.mod h1:nOB03cncLtlp4t+UAkGSV+9beXP/akpekBwL+UX1Qcw=
+github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
+github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
+github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
+github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
+github.com/golang/snappy v1.0.0 h1:Oy607GVXHs7RtbggtPBnr2RmDArIsAefDwvrdWvRhGs=
+github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
+github.com/google/flatbuffers v23.5.26+incompatible h1:M9dgRyhJemaM4Sw8+66GHBu8ioaQmyPLg1b8VwK5WJg=
+github.com/google/flatbuffers v23.5.26+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8=
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
 github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
 github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
 github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
+github.com/klauspost/compress v1.18.3 h1:9PJRvfbmTabkOX8moIpXPbMMbYN60bWImDDU7L+/6zw=
+github.com/klauspost/compress v1.18.3/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4=
+github.com/klauspost/cpuid/v2 v2.2.5 h1:0E5MSMDEoAulmXNFquVs//DdoomxaoTY1kUhbc/qbZg=
+github.com/klauspost/cpuid/v2 v2.2.5/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
+github.com/pierrec/lz4/v4 v4.1.25 h1:kocOqRffaIbU5djlIBr7Wh+cx82C0vtFb0fOurZHqD0=
+github.com/pierrec/lz4/v4 v4.1.25/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w=
@@ -14,10 +38,49 @@ github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
 github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
-golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
-golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/text v0.20.0 h1:gK/Kv2otX8gz+wn7Rmb3vT96ZwuoxnQlY+HlJVj7Qug=
-golang.org/x/text v0.20.0/go.mod h1:D4IsuqiFMhST5bX19pQ9ikHC2GsaKyk/oF+pn3ducp4=
+github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ=
+github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0=
+github.com/zeebo/xxh3 v1.0.2 h1:xZmwmqxHZA8AI603jOQ0tMqmBr9lPeFwGg6d+xy9DC0=
+github.com/zeebo/xxh3 v1.0.2/go.mod h1:5NWz9Sef7zIDm2JHfFlcQvNekmcEl9ekUZQQKCYaDcA=
+go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
+go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
+go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8=
+go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM=
+go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA=
+go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI=
+go.opentelemetry.io/otel/sdk v1.38.0 h1:l48sr5YbNf2hpCUj/FoGhW9yDkl+Ma+LrVl8qaM5b+E=
+go.opentelemetry.io/otel/sdk v1.38.0/go.mod h1:ghmNdGlVemJI3+ZB5iDEuk4bWA3GkTpW+DOoZMYBVVg=
+go.opentelemetry.io/otel/sdk/metric v1.38.0 h1:aSH66iL0aZqo//xXzQLYozmWrXxyFkBJ6qT5wthqPoM=
+go.opentelemetry.io/otel/sdk/metric v1.38.0/go.mod h1:dg9PBnW9XdQ1Hd6ZnRz689CbtrUp0wMMs9iPcgT9EZA=
+go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE=
+go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs=
+golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
+golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
+golang.org/x/mod v0.29.0 h1:HV8lRxZC4l2cr3Zq1LvtOsi/ThTgWnUk/y64QSs8GwA=
+golang.org/x/mod v0.29.0/go.mod h1:NyhrlYXJ2H4eJiRy/WDBO6HMqZQ6q9nk4JzS3NuCK+w=
+golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
+golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU=
+golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I=
+golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
+golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
+golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/telemetry v0.0.0-20251008203120-078029d740a8 h1:LvzTn0GQhWuvKH/kVRS3R3bVAsdQWI7hvfLHGgh9+lU=
+golang.org/x/telemetry v0.0.0-20251008203120-078029d740a8/go.mod h1:Pi4ztBfryZoJEkyFTI5/Ocsu2jXyDr6iSdgJiYE/uwE=
+golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM=
+golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM=
+golang.org/x/tools v0.38.0 h1:Hx2Xv8hISq8Lm16jvBZ2VQf+RLmbd7wVUsALibYI/IQ=
+golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs=
+golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 h1:H2TDz8ibqkAF6YGhCdN3jS9O0/s90v0rJh3X/OLHEUk=
+golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2/go.mod h1:K8+ghG5WaK9qNqU5K3HdILfMLy1f3aNYFI/wnl100a8=
+gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
+gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251029180050-ab9386a59fda h1:i/Q+bfisr7gq6feoJnS/DlpdwEL4ihp41fvRiM3Ork0=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251029180050-ab9386a59fda/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
+google.golang.org/grpc v1.78.0 h1:K1XZG/yGDJnzMdd/uZHAkVqJE+xIDOcmdSFZkBUicNc=
+google.golang.org/grpc v1.78.0/go.mod h1:I47qjTo4OKbMkjA/aOOwxDIiPSBofUtQUI5EfpWvW7U=
+google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
+google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

integration-tests/batch_quickstart_test.go

@@ -25,11 +25,12 @@ func getEnv(key, defaultValue string) string {
 }
 
 var (
+	brokerHost = getEnv("BROKER_HOST", "127.0.0.1")
 	brokerPort = getEnv("BROKER_PORT", "8000")
 )
 
 func getPinotClientFromBroker(useMultistageEngine bool) *pinot.Connection {
-	pinotClient, err := pinot.NewFromBrokerList([]string{"localhost:" + brokerPort})
+	pinotClient, err := pinot.NewFromBrokerList([]string{brokerHost + ":" + brokerPort})
 	if err != nil {
 		log.Fatalln(err)
 	}
@@ -55,7 +56,7 @@ func getCustomHTTPClient() *http.Client {
 }
 
 func getPinotClientFromBrokerAndCustomHTTPClient(useMultistageEngine bool) *pinot.Connection {
-	pinotClient, err := pinot.NewFromBrokerListAndClient([]string{"localhost:" + brokerPort}, getCustomHTTPClient())
+	pinotClient, err := pinot.NewFromBrokerListAndClient([]string{brokerHost + ":" + brokerPort}, getCustomHTTPClient())
 	if err != nil {
 		log.Fatalln(err)
 	}
@@ -65,7 +66,7 @@ func getPinotClientFromBrokerAndCustomHTTPClient(useMultistageEngine bool) *pino
 
 func getPinotClientFromConfig(useMultistageEngine bool) *pinot.Connection {
 	pinotClient, err := pinot.NewWithConfig(&pinot.ClientConfig{
-		BrokerList:      []string{"localhost:" + brokerPort},
+		BrokerList:      []string{brokerHost + ":" + brokerPort},
 		HTTPTimeout:     10 * time.Second,
 		ExtraHTTPHeader: map[string]string{},
 	})
@@ -78,7 +79,7 @@ func getPinotClientFromConfig(useMultistageEngine bool) *pinot.Connection {
 
 func getPinotClientFromConfigAndCustomHTTPClient(useMultistageEngine bool) *pinot.Connection {
 	pinotClient, err := pinot.NewWithConfigAndClient(&pinot.ClientConfig{
-		BrokerList:      []string{"localhost:" + brokerPort},
+		BrokerList:      []string{brokerHost + ":" + brokerPort},
 		HTTPTimeout:     10 * time.Second,
 		ExtraHTTPHeader: map[string]string{},
 	}, getCustomHTTPClient())