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 <[email protected]>
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 <[email protected]>
Signed-off-by: Matt Topol <[email protected]>
---
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)
+}
