From e2be3f1bf1ce3f9ad607edfe77a296b651546846 Mon Sep 17 00:00:00 2001
From: David Li <li.davidm96@gmail.com>
Date: Thu, 7 Mar 2024 14:07:24 -0500
Subject: [PATCH] feat(go/adbc/driver/flightsql): support reuse-connection
 location (#1594)

Fixes #1588.
---
 .../driver/flightsql/flightsql_adbc_server_test.go   |  2 +-
 go/adbc/driver/flightsql/flightsql_connection.go     | 12 +++++++++++-
 go/adbc/go.mod                                       |  2 +-
 go/adbc/go.sum                                       |  2 ++
 4 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/go/adbc/driver/flightsql/flightsql_adbc_server_test.go b/go/adbc/driver/flightsql/flightsql_adbc_server_test.go
index 78ae01b441..3b5761495a 100644
--- 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
index 2b4ce93e9a..d0aa0b02bb 100644
--- 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
index 7e7b605ea6..93c96af4b8 100644
--- 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
index 75377ffe63..0911d39f76 100644
--- 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=