feat(go/adbc/driver/flightsql): support reuse-connection location (#1594

) Fixes #1588.
apache · Mar 7, 2024 · e2be3f1 · e2be3f1
1 parent bf42dc6
commit e2be3f1
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 3 deletions.
diff --git a/go/adbc/driver/flightsql/flightsql_adbc_server_test.go b/go/adbc/driver/flightsql/flightsql_adbc_server_test.go
@@ -1597,7 +1597,7 @@ func (server *MultiTableTestServer) GetFlightInfoTables(ctx context.Context, cmd
 	server.Alloc = memory.NewCheckedAllocator(memory.DefaultAllocator)
 	info := &flight.FlightInfo{
 		Endpoint: []*flight.FlightEndpoint{
-			{Ticket: &flight.Ticket{Ticket: desc.Cmd}},
+			{Ticket: &flight.Ticket{Ticket: desc.Cmd}, Location: []*flight.Location{{Uri: flight.LocationReuseConnection}}},
 		},
 		FlightDescriptor: desc,
 		Schema:           flight.SerializeSchema(schema, server.Alloc),

diff --git a/go/adbc/driver/flightsql/flightsql_connection.go b/go/adbc/driver/flightsql/flightsql_connection.go
@@ -64,10 +64,16 @@ func doGet(ctx context.Context, cl *flightsql.Client, endpoint *flight.FlightEnd
 	}
 
 	var (
-		cc interface{}
+		cc          interface{}
+		hasFallback bool
 	)
 
 	for _, loc := range endpoint.Location {
+		if loc.Uri == flight.LocationReuseConnection {
+			hasFallback = true
+			continue
+		}
+
 		cc, err = clientCache.Get(loc.Uri)
 		if err != nil {
 			continue
@@ -82,6 +88,10 @@ func doGet(ctx context.Context, cl *flightsql.Client, endpoint *flight.FlightEnd
 		return
 	}
 
+	if hasFallback {
+		return cl.DoGet(ctx, endpoint.Ticket, opts...)
+	}
+
 	return nil, err
 }
 

diff --git a/go/adbc/go.mod b/go/adbc/go.mod
@@ -20,7 +20,7 @@ module github.com/apache/arrow-adbc/go/adbc
 go 1.19
 
 require (
-	github.com/apache/arrow/go/v16 v16.0.0-20240129203910-c2ca9bcedeb0
+	github.com/apache/arrow/go/v16 v16.0.0-20240307132415-1c9a3122c603
 	github.com/bluele/gcache v0.0.2
 	github.com/golang/protobuf v1.5.3
 	github.com/google/uuid v1.6.0

diff --git a/go/adbc/go.sum b/go/adbc/go.sum
@@ -27,6 +27,8 @@ github.com/apache/arrow/go/v14 v14.0.2 h1:N8OkaJEOfI3mEZt07BIkvo4sC6XDbL+48MBPWO
 github.com/apache/arrow/go/v14 v14.0.2/go.mod h1:u3fgh3EdgN/YQ8cVQRguVW3R+seMybFg8QBQ5LU+eBY=
 github.com/apache/arrow/go/v16 v16.0.0-20240129203910-c2ca9bcedeb0 h1:ooLFCCZ/sq3KDyrcFBxWweB1wTr1oAIgjj1+Zl3WsRw=
 github.com/apache/arrow/go/v16 v16.0.0-20240129203910-c2ca9bcedeb0/go.mod h1:+HkSDKotr3KDBxj7gTVgj8Egy18Y1ECzQdnY5XsXwlQ=
+github.com/apache/arrow/go/v16 v16.0.0-20240307132415-1c9a3122c603 h1:UOXjIpzPxFAsxrtqUa+e8yuVdhMklFi+Uyo6oB+sDK4=
+github.com/apache/arrow/go/v16 v16.0.0-20240307132415-1c9a3122c603/go.mod h1:+HkSDKotr3KDBxj7gTVgj8Egy18Y1ECzQdnY5XsXwlQ=
 github.com/apache/thrift v0.17.0 h1:cMd2aj52n+8VoAtvSvLn4kDC3aZ6IAkBuqWQ2IDu7wo=
 github.com/apache/thrift v0.17.0/go.mod h1:OLxhMRJxomX+1I/KUw03qoV3mMz16BwaKI+d4fPBx7Q=
 github.com/aws/aws-sdk-go-v2 v1.24.1 h1:xAojnj+ktS95YZlDf0zxWBkbFtymPeDP+rvUQIH3uAU=