Refactor WrapRecordReader to main class

dremio-hub · Jan 2, 2025 · 70f2bd8 · 70f2bd8
1 parent 725b002
commit 70f2bd8
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 30 deletions.
diff --git 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
@@ -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
@@ -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