MINOR: [Go] Add gRPC status details to sample Flight SQL server (apac…

…he#37026) ### Rationale for this change Needed to test apache/arrow-adbc#963. ### What changes are included in this PR? Have the SQLite Flight SQL server sample emit a gRPC status detail. ### Are these changes tested? No. ### Are there any user-facing changes? No. Authored-by: David Li <[email protected]> Signed-off-by: David Li <[email protected]>
assignUser · Aug 8, 2023 · c6c21c0 · c6c21c0
1 parent 21e1e00
commit c6c21c0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/go/arrow/flight/flightsql/example/sql_batch_reader.go b/go/arrow/flight/flightsql/example/sql_batch_reader.go
@@ -31,6 +31,9 @@ import (
 	"github.com/apache/arrow/go/v13/arrow/flight/flightsql"
 	"github.com/apache/arrow/go/v13/arrow/internal/debug"
 	"github.com/apache/arrow/go/v13/arrow/memory"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
 func getArrowTypeFromString(dbtype string) arrow.DataType {
@@ -257,12 +260,19 @@ func (r *SqlBatchReader) Next() bool {
 	rows := 0
 	for rows < maxBatchSize && r.rows.Next() {
 		if err := r.rows.Scan(r.rowdest...); err != nil {
-			r.err = err
+			// Not really useful except for testing Flight SQL clients
+			detail := wrapperspb.StringValue{Value: r.schema.String()}
+			if st, sterr := status.New(codes.Unknown, err.Error()).WithDetails(&detail); sterr != nil {
+				r.err = err
+			} else {
+				r.err = st.Err()
+			}
 			return false
 		}
 
 		for i, v := range r.rowdest {
 			fb := r.bldr.Field(i)
+
 			switch v := v.(type) {
 			case *uint8:
 				fb.(*array.Uint8Builder).Append(*v)

diff --git a/go/arrow/flight/flightsql/example/sqlite_server.go b/go/arrow/flight/flightsql/example/sqlite_server.go
@@ -52,7 +52,9 @@ import (
 	"github.com/apache/arrow/go/v13/arrow/flight/flightsql/schema_ref"
 	"github.com/apache/arrow/go/v13/arrow/memory"
 	"github.com/apache/arrow/go/v13/arrow/scalar"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	_ "modernc.org/sqlite"
 )
@@ -462,6 +464,9 @@ type dbQueryCtx interface {
 func doGetQuery(ctx context.Context, mem memory.Allocator, db dbQueryCtx, query string, schema *arrow.Schema, args ...interface{}) (*arrow.Schema, <-chan flight.StreamChunk, error) {
 	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
+		// Not really useful except for testing Flight SQL clients
+		trailers := metadata.Pairs("afsql-sqlite-query", query)
+		grpc.SetTrailer(ctx, trailers)
 		return nil, nil, err
 	}