(arrow) branch main updated: GH-39837: [Go][Flight] Allow cloning existing cookies in middleware (#39838)

[ze...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new c2ca9bcede GH-39837: [Go][Flight] Allow cloning existing cookies in middleware (#39838)
c2ca9bcede is described below

commit c2ca9bcedeb004f9d7f5d3e1aafc7b83ce6c1e3f
Author: Matt Topol <zo...@gmail.com>
AuthorDate: Mon Jan 29 15:39:10 2024 -0500

    GH-39837: [Go][Flight] Allow cloning existing cookies in middleware (#39838)
    
    
    
    ### Rationale for this change
    This is needed for https://github.com/apache/arrow-adbc/issues/1194 to facilitate better connection handling for flight clients in ADBC by copying the existing cookies over when creating a sub-client.
    
    ### What changes are included in this PR?
    Creating a `Clone` method on the `CookieMiddleware` so that a user can create and hold a reference to a specific cookie middleware instance and then create new ones on the fly that copy over the existing cookies at that moment.
    
    ### Are these changes tested?
    Yes.
    
    ### Are there any user-facing changes?
    No
    
    * Closes: #39837
    
    Authored-by: Matt Topol <zo...@gmail.com>
    Signed-off-by: Matt Topol <zo...@gmail.com>
---
 go/arrow/flight/cookie_middleware.go      | 24 +++++++++++++
 go/arrow/flight/cookie_middleware_test.go | 60 +++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/go/arrow/flight/cookie_middleware.go b/go/arrow/flight/cookie_middleware.go
index 27754a13b8..39c86d8303 100644
--- a/go/arrow/flight/cookie_middleware.go
+++ b/go/arrow/flight/cookie_middleware.go
@@ -23,6 +23,7 @@ import (
 	"sync"
 	"time"
 
+	"golang.org/x/exp/maps"
 	"google.golang.org/grpc/metadata"
 )
 
@@ -40,11 +41,34 @@ func NewClientCookieMiddleware() ClientMiddleware {
 	return CreateClientMiddleware(&clientCookieMiddleware{jar: make(map[string]http.Cookie)})
 }
 
+func NewCookieMiddleware() CookieMiddleware {
+	return &clientCookieMiddleware{jar: make(map[string]http.Cookie)}
+}
+
+// CookieMiddleware is a go-routine safe middleware for flight clients
+// which properly handles Set-Cookie headers for storing cookies.
+// This can be passed into `CreateClientMiddleware` to create a new
+// middleware object. You can also clone it to create middleware for a
+// new client which starts with the same cookies.
+type CookieMiddleware interface {
+	CustomClientMiddleware
+	// Clone creates a new CookieMiddleware that starts out with the same
+	// cookies that this one already has. This is useful when creating a
+	// new client connection for the same server.
+	Clone() CookieMiddleware
+}
+
 type clientCookieMiddleware struct {
 	jar map[string]http.Cookie
 	mx  sync.Mutex
 }
 
+func (cc *clientCookieMiddleware) Clone() CookieMiddleware {
+	cc.mx.Lock()
+	defer cc.mx.Unlock()
+	return &clientCookieMiddleware{jar: maps.Clone(cc.jar)}
+}
+
 func (cc *clientCookieMiddleware) StartCall(ctx context.Context) context.Context {
 	cc.mx.Lock()
 	defer cc.mx.Unlock()
diff --git a/go/arrow/flight/cookie_middleware_test.go b/go/arrow/flight/cookie_middleware_test.go
index 0adf492765..4007d056b2 100644
--- a/go/arrow/flight/cookie_middleware_test.go
+++ b/go/arrow/flight/cookie_middleware_test.go
@@ -239,3 +239,63 @@ func TestCookieExpiration(t *testing.T) {
 	cookieMiddleware.expectedCookies = map[string]string{}
 	makeReq(client, t)
 }
+
+func TestCookiesClone(t *testing.T) {
+	cookieMiddleware := &serverAddCookieMiddleware{}
+
+	s := flight.NewServerWithMiddleware([]flight.ServerMiddleware{
+		flight.CreateServerMiddleware(cookieMiddleware),
+	})
+	s.Init("localhost:0")
+	f := &flightServer{}
+	s.RegisterFlightService(f)
+
+	go s.Serve()
+	defer s.Shutdown()
+
+	makeReq := func(c flight.Client, t *testing.T) {
+		flightStream, err := c.ListFlights(context.Background(), &flight.Criteria{})
+		assert.NoError(t, err)
+
+		for {
+			_, err := flightStream.Recv()
+			if err != nil {
+				if errors.Is(err, io.EOF) {
+					break
+				}
+				assert.NoError(t, err)
+			}
+		}
+	}
+
+	credsOpt := grpc.WithTransportCredentials(insecure.NewCredentials())
+	cookies := flight.NewCookieMiddleware()
+	client1, err := flight.NewClientWithMiddleware(s.Addr().String(), nil,
+		[]flight.ClientMiddleware{flight.CreateClientMiddleware(cookies)}, credsOpt)
+	require.NoError(t, err)
+	defer client1.Close()
+
+	// set cookies
+	cookieMiddleware.cookies = []*http.Cookie{
+		{Name: "foo", Value: "bar"},
+		{Name: "foo2", Value: "bar2", MaxAge: 1},
+	}
+	makeReq(client1, t)
+
+	// validate set
+	cookieMiddleware.expectedCookies = map[string]string{
+		"foo": "bar", "foo2": "bar2",
+	}
+	makeReq(client1, t)
+
+	client2, err := flight.NewClientWithMiddleware(s.Addr().String(), nil,
+		[]flight.ClientMiddleware{flight.CreateClientMiddleware(cookies.Clone())}, credsOpt)
+	require.NoError(t, err)
+	defer client2.Close()
+
+	// validate clone worked
+	cookieMiddleware.expectedCookies = map[string]string{
+		"foo": "bar", "foo2": "bar2",
+	}
+	makeReq(client2, t)
+}