feat: support for client middleware in Flight configuration and testing

Author: NguyenHoangSon96

influxdb3/config.go

@@ -31,6 +31,7 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/apache/arrow-go/v18/arrow/flight"
 	"github.com/influxdata/line-protocol/v2/lineprotocol"
 )
 
@@ -149,6 +150,9 @@ type ClientConfig struct {
 
 	// Proxy URL
 	Proxy string
+
+	// Flight client middleware
+	Middleware []flight.ClientMiddleware
 }
 
 // validate validates the config.

influxdb3/query.go

@@ -79,7 +79,7 @@ func (c *Client) initializeQueryClient(hostPortURL string, secure bool, proxyURL
 		}
 	}
 
-	client, err := flight.NewClientWithMiddleware(hostPortURL, nil, nil, opts...)
+	client, err := flight.NewClientWithMiddleware(hostPortURL, nil, c.config.Middleware, opts...)
 	if err != nil {
 		return fmt.Errorf("flight: %w", err)
 	}

influxdb3/query_test.go

@@ -213,6 +213,60 @@ func TestQueryWithLargeResponsePass(t *testing.T) {
 	assert.Nil(t, qIter.Err())
 }
 
+func TestQueryWithMiddlewareSuccess(t *testing.T) {
+	s := *testutil.StartCheckMessageFromMiddlewareFlightServer(t)
+	defer func() {
+		s.Shutdown()
+	}()
+
+	middlewares := []flight.ClientMiddleware{
+		flight.CreateClientMiddleware(&testutil.ClientTestMiddleware{}),
+	}
+
+	client, err := New(ClientConfig{
+		Host:       "http://" + s.Addr().String(),
+		Token:      "my_secret_token",
+		Database:   "explore",
+		Middleware: middlewares,
+	})
+	require.NoError(t, err)
+	defer func(client *Client) {
+		err := client.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}(client)
+	_, qErr := client.Query(context.Background(),
+		"SELECT name FROM examples",
+	)
+	require.NotContains(t, qErr.Error(), "invalid value from middleware")
+}
+
+func TestQueryWithMiddlewareFail(t *testing.T) {
+	s := *testutil.StartCheckMessageFromMiddlewareFlightServer(t)
+	defer func() {
+		s.Shutdown()
+	}()
+
+	client, err := New(ClientConfig{
+		Host:       "http://" + s.Addr().String(),
+		Token:      "my_secret_token",
+		Database:   "explore",
+		Middleware: nil,
+	})
+	require.NoError(t, err)
+	defer func(client *Client) {
+		err := client.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}(client)
+	_, qErr := client.Query(context.Background(),
+		"SELECT name FROM examples",
+	)
+	require.Contains(t, qErr.Error(), "invalid value from middleware")
+}
+
 func TestQueryWithQueryTimeoutDeadlineExpired(t *testing.T) {
 	s := flight.NewServerWithMiddleware(nil)
 	err := s.Init("localhost:0")

influxdb3/testutil/mocks.go

@@ -3,6 +3,7 @@ package testutil
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -20,6 +21,7 @@ import (
 	"github.com/apache/arrow-go/v18/arrow/memory"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
@@ -461,3 +463,45 @@ func arrayOf(mem memory.Allocator, a any, valids []bool) arrow.Array {
 		panic(fmt.Errorf("arrdata: invalid data slice type %T", a))
 	}
 }
+
+type ClientTestMiddleware struct {
+}
+
+func (c *ClientTestMiddleware) StartCall(ctx context.Context) context.Context {
+	return metadata.AppendToOutgoingContext(ctx, "sent-from-middleware", "some-value")
+}
+
+type CheckMessageFromMiddlewareFlightServer struct {
+	flight.BaseFlightServer
+}
+
+func (f *CheckMessageFromMiddlewareFlightServer) DoGet(tkt *flight.Ticket, fs flight.FlightService_DoGetServer) error {
+	ctx := fs.Context()
+	md, _ := metadata.FromIncomingContext(ctx)
+	values := md.Get("sent-from-middleware")
+	if values == nil || values[0] != "some-value" {
+		return status.Errorf(codes.Internal, "invalid value from middleware")
+	}
+
+	return nil
+}
+
+//nolint:all
+func StartCheckMessageFromMiddlewareFlightServer(t *testing.T) *flight.Server {
+	mockServer := CheckMessageFromMiddlewareFlightServer{}
+	s := flight.NewServerWithMiddleware([]flight.ServerMiddleware{})
+	err := s.Init("localhost:0")
+	if err != nil {
+		assert.Fail(t, err.Error())
+	}
+	s.RegisterFlightService(&mockServer)
+
+	go func() {
+		err := s.Serve()
+		if err != nil {
+			assert.Fail(t, err.Error())
+		}
+	}()
+
+	return &s
+}