feat: Implement plugin Read (#1027)

#1026 should be merged first

Author: hermanschaaf

internal/clients/state/v3/state.go

@@ -72,30 +72,23 @@ func NewClient(ctx context.Context, pbClient pb.PluginClient, tableName string)
 		return nil, err
 	}
 
-	syncClient, err := c.client.Sync(ctx, &pb.Sync_Request{
-		Tables: []string{tableName},
+	readClient, err := c.client.Read(ctx, &pb.Read_Request{
+		Table: tableBytes,
 	})
 	if err != nil {
 		return nil, err
 	}
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 	for {
-		res, err := syncClient.Recv()
+		res, err := readClient.Recv()
 		if err != nil {
 			if err == io.EOF {
 				break
 			}
 			return nil, err
 		}
-		var insertMessage *pb.Sync_Response_Insert
-		switch m := res.Message.(type) {
-		case *pb.Sync_Response_MigrateTable:
-			continue
-		case *pb.Sync_Response_Insert:
-			insertMessage = m
-		}
-		rdr, err := ipc.NewReader(bytes.NewReader(insertMessage.Insert.Record))
+		rdr, err := ipc.NewReader(bytes.NewReader(res.Record))
 		if err != nil {
 			return nil, err
 		}

internal/servers/plugin/v3/plugin.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 
+	"github.com/apache/arrow/go/v13/arrow"
 	pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3"
 	"github.com/cloudquery/plugin-sdk/v4/message"
 	"github.com/cloudquery/plugin-sdk/v4/plugin"
@@ -63,6 +64,43 @@ func (s *Server) Init(ctx context.Context, req *pb.Init_Request) (*pb.Init_Respo
 	return &pb.Init_Response{}, nil
 }
 
+func (s *Server) Read(req *pb.Read_Request, stream pb.Plugin_ReadServer) error {
+	records := make(chan arrow.Record)
+	var syncErr error
+	ctx := stream.Context()
+
+	sc, err := pb.NewSchemaFromBytes(req.Table)
+	if err != nil {
+		return status.Errorf(codes.InvalidArgument, "failed to create schema from bytes: %v", err)
+	}
+	table, err := schema.NewTableFromArrowSchema(sc)
+	if err != nil {
+		return status.Errorf(codes.InvalidArgument, "failed to create table from schema: %v", err)
+	}
+	go func() {
+		defer close(records)
+		err := s.Plugin.Read(ctx, table, records)
+		if err != nil {
+			syncErr = fmt.Errorf("failed to sync records: %w", err)
+		}
+	}()
+
+	for rec := range records {
+		recBytes, err := pb.RecordToBytes(rec)
+		if err != nil {
+			return status.Errorf(codes.Internal, "failed to convert record to bytes: %v", err)
+		}
+		resp := &pb.Read_Response{
+			Record: recBytes,
+		}
+		if err := stream.Send(resp); err != nil {
+			return status.Errorf(codes.Internal, "failed to send read response: %v", err)
+		}
+	}
+
+	return syncErr
+}
+
 func (s *Server) Sync(req *pb.Sync_Request, stream pb.Plugin_SyncServer) error {
 	msgs := make(chan message.SyncMessage)
 	var syncErr error

plugin/plugin.go

@@ -27,7 +27,7 @@ func (UnimplementedDestination) Write(context.Context, <-chan message.WriteMessa
 }
 
 func (UnimplementedDestination) Read(context.Context, *schema.Table, chan<- arrow.Record) error {
-	return fmt.Errorf("not implemented")
+	return ErrNotImplemented
 }
 
 type UnimplementedSource struct{}

plugin/plugin_destination.go

@@ -37,3 +37,18 @@ func (p *Plugin) Write(ctx context.Context, res <-chan message.WriteMessage) err
 	}
 	return p.client.Write(ctx, res)
 }
+
+// Read is read data from the requested table to the given channel, returned in the same format as the table
+func (p *Plugin) Read(ctx context.Context, table *schema.Table, res chan<- arrow.Record) error {
+	if !p.mu.TryLock() {
+		return fmt.Errorf("plugin already in use")
+	}
+	defer p.mu.Unlock()
+	if p.client == nil {
+		return fmt.Errorf("plugin not initialized. call Init() first")
+	}
+	if err := p.client.Read(ctx, table, res); err != nil {
+		return fmt.Errorf("failed to read: %w", err)
+	}
+	return nil
+}