feat(go/adbc/driver/flightsql): propagate cookies to sub-clients (apa…

…che#1497) Closes apache#1194 This also bumps our dependency up to arrow v16
soumyadsanyal · Jan 31, 2024 · 168817e · 168817e
1 parent 3a50405
commit 168817e
Show file tree

Hide file tree

Showing 36 changed files with 172 additions and 143 deletions.
diff --git a/go/adbc/adbc.go b/go/adbc/adbc.go
@@ -40,8 +40,8 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 )

diff --git a/go/adbc/driver/driverbase/database.go b/go/adbc/driver/driverbase/database.go
@@ -21,7 +21,7 @@ import (
 	"context"
 
 	"github.com/apache/arrow-adbc/go/adbc"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"golang.org/x/exp/slog"
 )
 

diff --git a/go/adbc/driver/driverbase/driver.go b/go/adbc/driver/driverbase/driver.go
@@ -22,7 +22,7 @@ package driverbase
 
 import (
 	"github.com/apache/arrow-adbc/go/adbc"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 )
 
 // DriverImpl is an interface that drivers implement to provide

diff --git a/go/adbc/driver/flightsql/cmd/testserver/main.go b/go/adbc/driver/flightsql/cmd/testserver/main.go
@@ -31,11 +31,11 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/flight"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/flight"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"

diff --git a/go/adbc/driver/flightsql/flightsql_adbc_server_test.go b/go/adbc/driver/flightsql/flightsql_adbc_server_test.go
@@ -23,7 +23,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/google/uuid"
 	"net/textproto"
 	"os"
 	"strconv"
@@ -32,14 +31,16 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
+
 	"github.com/apache/arrow-adbc/go/adbc"
 	driver "github.com/apache/arrow-adbc/go/adbc/driver/flightsql"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/flight"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql/schema_ref"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/flight"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql/schema_ref"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"github.com/golang/protobuf/ptypes/wrappers"
 	"github.com/stretchr/testify/suite"
 	"golang.org/x/exp/maps"
@@ -1078,7 +1079,8 @@ func (ts *TimeoutTests) TestDontTimeout() {
 type CookieTestServer struct {
 	flightsql.BaseServer
 
-	cur time.Time
+	cur  time.Time
+	addr string
 }
 
 func (server *CookieTestServer) GetFlightInfoStatement(ctx context.Context, cmd flightsql.StatementQuery, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) {
@@ -1095,7 +1097,13 @@ func (server *CookieTestServer) GetFlightInfoStatement(ctx context.Context, cmd
 	info := &flight.FlightInfo{
 		FlightDescriptor: desc,
 		Endpoint: []*flight.FlightEndpoint{
-			{Ticket: &flight.Ticket{Ticket: tkt}},
+			{
+				Ticket: &flight.Ticket{Ticket: tkt},
+				// passing a non-empty location uri so that the test client
+				// creates a sub-client and we test that the cookies are
+				// preserved and copied over.
+				Location: []*flight.Location{{Uri: server.addr}},
+			},
 		},
 		TotalRecords: -1,
 		TotalBytes:   -1,
@@ -1173,9 +1181,11 @@ type CookieTests struct {
 }
 
 func (suite *CookieTests) SetupSuite() {
-	suite.DoSetupSuite(&CookieTestServer{}, nil, map[string]string{
+	ts := &CookieTestServer{}
+	suite.DoSetupSuite(ts, nil, map[string]string{
 		driver.OptionCookieMiddleware: adbc.OptionValueEnabled,
 	})
+	ts.addr = "grpc://" + suite.s.Addr().String()
 }
 
 func (suite *CookieTests) TestCookieUsage() {

diff --git a/go/adbc/driver/flightsql/flightsql_adbc_test.go b/go/adbc/driver/flightsql/flightsql_adbc_test.go
@@ -42,12 +42,12 @@ import (
 	"github.com/apache/arrow-adbc/go/adbc"
 	driver "github.com/apache/arrow-adbc/go/adbc/driver/flightsql"
 	"github.com/apache/arrow-adbc/go/adbc/validation"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/flight"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql/example"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/flight"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql/example"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	"google.golang.org/grpc"
@@ -260,7 +260,7 @@ func (s *FlightSQLQuirks) GetMetadata(code adbc.InfoCode) interface{} {
 	case adbc.InfoVendorVersion:
 		return "sqlite 3"
 	case adbc.InfoVendorArrowVersion:
-		return "15.0.0-SNAPSHOT"
+		return "16.0.0-SNAPSHOT"
 	}
 
 	return nil

diff --git a/go/adbc/driver/flightsql/flightsql_connection.go b/go/adbc/driver/flightsql/flightsql_connection.go
@@ -27,12 +27,12 @@ import (
 
 	"github.com/apache/arrow-adbc/go/adbc"
 	"github.com/apache/arrow-adbc/go/adbc/driver/internal"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/flight"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql/schema_ref"
-	"github.com/apache/arrow/go/v15/arrow/ipc"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/flight"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql/schema_ref"
+	"github.com/apache/arrow/go/v16/arrow/ipc"
 	"github.com/bluele/gcache"
 	"google.golang.org/grpc"
 	grpccodes "google.golang.org/grpc/codes"

diff --git a/go/adbc/driver/flightsql/flightsql_database.go b/go/adbc/driver/flightsql/flightsql_database.go
@@ -30,9 +30,9 @@ import (
 
 	"github.com/apache/arrow-adbc/go/adbc"
 	"github.com/apache/arrow-adbc/go/adbc/driver/driverbase"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/flight"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/flight"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql"
 	"github.com/bluele/gcache"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
@@ -336,8 +336,7 @@ func (d *databaseImpl) Close() error {
 	return nil
 }
 
-func getFlightClient(ctx context.Context, loc string, d *databaseImpl) (*flightsql.Client, error) {
-	authMiddle := &bearerAuthMiddleware{hdrs: d.hdrs.Copy()}
+func getFlightClient(ctx context.Context, loc string, d *databaseImpl, authMiddle *bearerAuthMiddleware, cookies flight.CookieMiddleware) (*flightsql.Client, error) {
 	middleware := []flight.ClientMiddleware{
 		{
 			Unary:  makeUnaryLoggingInterceptor(d.Logger),
@@ -351,7 +350,7 @@ func getFlightClient(ctx context.Context, loc string, d *databaseImpl) (*flights
 	}
 
 	if d.enableCookies {
-		middleware = append(middleware, flight.NewClientCookieMiddleware())
+		middleware = append(middleware, flight.CreateClientMiddleware(cookies))
 	}
 
 	uri, err := url.Parse(loc)
@@ -378,17 +377,21 @@ func getFlightClient(ctx context.Context, loc string, d *databaseImpl) (*flights
 	}
 
 	cl.Alloc = d.Alloc
-	if d.user != "" || d.pass != "" {
-		var header, trailer metadata.MD
-		ctx, err = cl.Client.AuthenticateBasicToken(ctx, d.user, d.pass, grpc.Header(&header), grpc.Trailer(&trailer), d.timeout)
-		if err != nil {
-			return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "AuthenticateBasicToken")
-		}
+	if len(authMiddle.hdrs.Get("authorization")) > 0 {
+		d.Logger.DebugContext(ctx, "reusing auth token", "location", loc)
+	} else {
+		if d.user != "" || d.pass != "" {
+			var header, trailer metadata.MD
+			ctx, err = cl.Client.AuthenticateBasicToken(ctx, d.user, d.pass, grpc.Header(&header), grpc.Trailer(&trailer), d.timeout)
+			if err != nil {
+				return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "AuthenticateBasicToken")
+			}
 
-		if md, ok := metadata.FromOutgoingContext(ctx); ok {
-			authMiddle.mutex.Lock()
-			defer authMiddle.mutex.Unlock()
-			authMiddle.hdrs.Set("authorization", md.Get("Authorization")[0])
+			if md, ok := metadata.FromOutgoingContext(ctx); ok {
+				authMiddle.mutex.Lock()
+				defer authMiddle.mutex.Unlock()
+				authMiddle.hdrs.Set("authorization", md.Get("Authorization")[0])
+			}
 		}
 	}
 
@@ -401,7 +404,13 @@ type support struct {
 }
 
 func (d *databaseImpl) Open(ctx context.Context) (adbc.Connection, error) {
-	cl, err := getFlightClient(ctx, d.uri.String(), d)
+	authMiddle := &bearerAuthMiddleware{hdrs: d.hdrs.Copy()}
+	var cookies flight.CookieMiddleware
+	if d.enableCookies {
+		cookies = flight.NewCookieMiddleware()
+	}
+
+	cl, err := getFlightClient(ctx, d.uri.String(), d, authMiddle, cookies)
 	if err != nil {
 		return nil, err
 	}
@@ -411,10 +420,18 @@ func (d *databaseImpl) Open(ctx context.Context) (adbc.Connection, error) {
 		LoaderFunc(func(loc interface{}) (interface{}, error) {
 			uri, ok := loc.(string)
 			if !ok {
-				return nil, adbc.Error{Msg: fmt.Sprintf("Location must be a string, got %#v", uri), Code: adbc.StatusInternal}
+				return nil, adbc.Error{Msg: fmt.Sprintf("Location must be a string, got %#v",
+					uri), Code: adbc.StatusInternal}
 			}
 
-			cl, err := getFlightClient(context.Background(), uri, d)
+			var cookieMiddleware flight.CookieMiddleware
+			// if cookies are enabled, start by cloning the existing cookies
+			if d.enableCookies {
+				cookieMiddleware = cookies.Clone()
+			}
+			// use the existing auth token if there is one
+			cl, err := getFlightClient(context.Background(), uri, d,
+				&bearerAuthMiddleware{hdrs: authMiddle.hdrs.Copy()}, cookieMiddleware)
 			if err != nil {
 				return nil, err
 			}

diff --git a/go/adbc/driver/flightsql/flightsql_driver.go b/go/adbc/driver/flightsql/flightsql_driver.go
@@ -38,7 +38,7 @@ import (
 
 	"github.com/apache/arrow-adbc/go/adbc"
 	"github.com/apache/arrow-adbc/go/adbc/driver/driverbase"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"golang.org/x/exp/maps"
 	"google.golang.org/grpc/metadata"
 )

diff --git a/go/adbc/driver/flightsql/flightsql_statement.go b/go/adbc/driver/flightsql/flightsql_statement.go
@@ -25,11 +25,11 @@ import (
 	"time"
 
 	"github.com/apache/arrow-adbc/go/adbc"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/flight"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/flight"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"github.com/bluele/gcache"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"

diff --git a/go/adbc/driver/flightsql/record_reader.go b/go/adbc/driver/flightsql/record_reader.go
@@ -24,11 +24,11 @@ import (
 
 	"github.com/apache/arrow-adbc/go/adbc"
 	"github.com/apache/arrow-adbc/go/adbc/utils"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/flight"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/flight"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"github.com/bluele/gcache"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"

diff --git a/go/adbc/driver/flightsql/record_reader_test.go b/go/adbc/driver/flightsql/record_reader_test.go
@@ -24,12 +24,12 @@ import (
 	"testing"
 
 	"github.com/apache/arrow-adbc/go/adbc"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/flight"
-	"github.com/apache/arrow/go/v15/arrow/flight/flightsql"
-	"github.com/apache/arrow/go/v15/arrow/ipc"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/flight"
+	"github.com/apache/arrow/go/v16/arrow/flight/flightsql"
+	"github.com/apache/arrow/go/v16/arrow/ipc"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"github.com/bluele/gcache"
 	"github.com/stretchr/testify/suite"
 	"google.golang.org/grpc"

diff --git a/go/adbc/driver/internal/shared_utils.go b/go/adbc/driver/internal/shared_utils.go
@@ -26,9 +26,9 @@ import (
 	"time"
 
 	"github.com/apache/arrow-adbc/go/adbc"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 )
 
 type CatalogAndSchema struct {

diff --git a/go/adbc/driver/panicdummy/panicdummy_adbc.go b/go/adbc/driver/panicdummy/panicdummy_adbc.go
@@ -25,9 +25,9 @@ import (
 	"os"
 
 	"github.com/apache/arrow-adbc/go/adbc"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 )
 
 func maybePanic(fname string) {

diff --git a/go/adbc/driver/snowflake/bulk_ingestion.go b/go/adbc/driver/snowflake/bulk_ingestion.go
@@ -33,12 +33,12 @@ import (
 	"sync"
 
 	"github.com/apache/arrow-adbc/go/adbc"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
-	"github.com/apache/arrow/go/v15/arrow/memory"
-	"github.com/apache/arrow/go/v15/parquet"
-	"github.com/apache/arrow/go/v15/parquet/compress"
-	"github.com/apache/arrow/go/v15/parquet/pqarrow"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/memory"
+	"github.com/apache/arrow/go/v16/parquet"
+	"github.com/apache/arrow/go/v16/parquet/compress"
+	"github.com/apache/arrow/go/v16/parquet/pqarrow"
 	"github.com/snowflakedb/gosnowflake"
 	"golang.org/x/sync/errgroup"
 )

diff --git a/go/adbc/driver/snowflake/connection.go b/go/adbc/driver/snowflake/connection.go
@@ -30,8 +30,8 @@ import (
 
 	"github.com/apache/arrow-adbc/go/adbc"
 	"github.com/apache/arrow-adbc/go/adbc/driver/internal"
-	"github.com/apache/arrow/go/v15/arrow"
-	"github.com/apache/arrow/go/v15/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
 	"github.com/snowflakedb/gosnowflake"
 )
 

diff --git a/go/adbc/driver/snowflake/driver.go b/go/adbc/driver/snowflake/driver.go
@@ -24,7 +24,7 @@ import (
 
 	"github.com/apache/arrow-adbc/go/adbc"
 	"github.com/apache/arrow-adbc/go/adbc/driver/driverbase"
-	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/apache/arrow/go/v16/arrow/memory"
 	"github.com/snowflakedb/gosnowflake"
 	"golang.org/x/exp/maps"
 )