From 70f2bd8670bc297a1cf74338b84e52412d5adc50 Mon Sep 17 00:00:00 2001
From: Valter Francisco <valter.francisco@dremio.com>
Date: Thu, 2 Jan 2025 11:44:50 +0000
Subject: [PATCH] Refactor WrapRecordReader to main class

---
 go/example.go                            | 16 ++++++++++++++--
 go/example_test.go                       | 12 ++++++------
 go/implementations/mock_record_reader.go |  4 ----
 go/interfaces/record_reader.go           | 18 ------------------
 4 files changed, 20 insertions(+), 30 deletions(-)
 delete mode 100644 go/interfaces/record_reader.go

diff --git a/go/example.go b/go/example.go
index fe2a54f7..19e6d727 100644
--- a/go/example.go
+++ b/go/example.go
@@ -16,6 +16,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"github.com/apache/arrow-go/v18/arrow"
 	"log"
 	"net"
 
@@ -49,6 +50,17 @@ Options:
   --certs=<path>      Path to trusted certificates for encrypted connection.
   --project_id=<project_id>   Dremio project ID`
 
+type RecordReader interface {
+	Next() bool
+	Record() arrow.Record
+	Err() error
+	Release()
+}
+
+func WrapRecordReader(stream flight.FlightService_DoGetClient) (RecordReader, error) {
+	return flight.NewRecordReader(stream)
+}
+
 func main() {
 	args, err := docopt.ParseDoc(usage)
 	if err != nil {
@@ -97,13 +109,13 @@ func main() {
 	}
 	defer client.Close()
 
-	if err := run(config, client, interfaces.WrapRecordReader); err != nil {
+	if err := run(config, client, WrapRecordReader); err != nil {
 		log.Fatal(err)
 	}
 }
 
 func run(config interfaces.FlightConfig, client flight.Client,
-	readerCreator func(flight.FlightService_DoGetClient) (interfaces.RecordReader, error),
+	readerCreator func(flight.FlightService_DoGetClient) (RecordReader, error),
 ) error {
 
 	// Two WLM settings can be provided upon initial authentication with the dremio
diff --git a/go/example_test.go b/go/example_test.go
index 95ce0da9..6e5a10be 100644
--- a/go/example_test.go
+++ b/go/example_test.go
@@ -144,7 +144,7 @@ func TestRun(t *testing.T) {
 		Query: "SELECT * FROM test",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
 		return implementations.NewMockRecordReader([]arrow.Record{record}), nil
 	}
 
@@ -231,7 +231,7 @@ func TestRunWithPAT(t *testing.T) {
 		ProjectID: "test_project_id",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
 		return implementations.NewMockRecordReader([]arrow.Record{record}), nil
 	}
 
@@ -307,7 +307,7 @@ func TestRunWithPATNoProjectID(t *testing.T) {
 		Query: "SELECT * FROM test",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
 		return implementations.NewMockRecordReader([]arrow.Record{record}), nil
 	}
 
@@ -335,7 +335,7 @@ func TestInvalidCredentials(t *testing.T) {
 		Query: "SELECT 1",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
 		t.Fatal("Reader creator should not be called due to authentication failure")
 		return nil, nil
 	}
@@ -373,7 +373,7 @@ func TestInvalidHost(t *testing.T) {
 		Query: "SELECT 1",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
 		t.Fatal("Reader creator should not be called due to authentication failure")
 		return nil, nil
 	}
@@ -411,7 +411,7 @@ func TestInvalidPort(t *testing.T) {
 		Query: "SELECT 1",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
 		t.Fatal("Reader creator should not be called due to authentication failure")
 		return nil, nil
 	}
diff --git a/go/implementations/mock_record_reader.go b/go/implementations/mock_record_reader.go
index 95428739..cf797281 100644
--- a/go/implementations/mock_record_reader.go
+++ b/go/implementations/mock_record_reader.go
@@ -2,12 +2,8 @@ package implementations
 
 import (
 	"github.com/apache/arrow-go/v18/arrow"
-
-	"arrow-flight-client-example/interfaces"
 )
 
-var _ interfaces.RecordReader = &MockRecordReader{}
-
 // MockRecordReader for testing purposes
 type MockRecordReader struct {
 	records      []arrow.Record
diff --git a/go/interfaces/record_reader.go b/go/interfaces/record_reader.go
deleted file mode 100644
index 336ccd7e..00000000
--- a/go/interfaces/record_reader.go
+++ /dev/null
@@ -1,18 +0,0 @@
-package interfaces
-
-import (
-	"github.com/apache/arrow-go/v18/arrow"
-	"github.com/apache/arrow-go/v18/arrow/flight"
-)
-
-// RecordReader interface to abstract record reading functionality
-type RecordReader interface {
-	Next() bool
-	Record() arrow.Record
-	Err() error
-	Release()
-}
-
-func WrapRecordReader(stream flight.FlightService_DoGetClient) (RecordReader, error) {
-	return flight.NewRecordReader(stream)
-}