feat: Implement transformations logic. (#1800)

This is a first draft of sdk changes to add transformations support, so that transformer plugins can be created.

I don't want to merge it yet, because the channel logic is not well thought out. I'd rather want to do a test sync locally with some transformation.

Author: marianogappa

examples/simple_plugin/plugin/client.go

@@ -64,6 +64,11 @@ func (*Client) Read(context.Context, *schema.Table, chan<- arrow.Record) error {
 	return nil
 }
 
+func (*Client) Transform(_ context.Context, _ <-chan arrow.Record, _ chan<- arrow.Record) error {
+	// Not implemented, just used for testing destination packaging
+	return nil
+}
+
 func Configure(_ context.Context, logger zerolog.Logger, spec []byte, opts plugin.NewClientOptions) (plugin.Client, error) {
 	if opts.NoConnection {
 		return &Client{

internal/memdb/memdb.go

@@ -306,6 +306,10 @@ func (c *client) deleteRecord(_ context.Context, msg *message.WriteDeleteRecord)
 	c.memoryDB[tableName] = filteredTable
 }
 
+func (*client) Transform(_ context.Context, _ <-chan arrow.Record, _ chan<- arrow.Record) error {
+	return nil
+}
+
 func evaluatePredicate(pred message.Predicate, record arrow.Record) bool {
 	sc := record.Schema()
 	indices := sc.FieldIndices(pred.Column)

internal/reversertransformer/reversertransformer.go

@@ -0,0 +1,91 @@
+package reversertransformer
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/memory"
+	"github.com/cloudquery/plugin-sdk/v4/plugin"
+	"github.com/rs/zerolog"
+)
+
+// client is mostly used for testing the destination plugin.
+type client struct {
+	plugin.UnimplementedDestination
+	plugin.UnimplementedSource
+}
+
+type Option func(*client)
+
+type Spec struct {
+}
+
+func GetNewClient(options ...Option) plugin.NewClientFunc {
+	c := &client{}
+	for _, opt := range options {
+		opt(c)
+	}
+	return func(context.Context, zerolog.Logger, []byte, plugin.NewClientOptions) (plugin.Client, error) {
+		return c, nil
+	}
+}
+
+func (*client) GetSpec() any {
+	return &Spec{}
+}
+
+func (*client) Close(context.Context) error {
+	return nil
+}
+
+func (c *client) Transform(ctx context.Context, recvRecords <-chan arrow.Record, sendRecords chan<- arrow.Record) error {
+	for {
+		select {
+		case record, ok := <-recvRecords:
+			if !ok {
+				return nil
+			}
+			reversedRecord, err := c.reverseStrings(record)
+			if err != nil {
+				return err
+			}
+			sendRecords <- reversedRecord
+		case <-ctx.Done():
+			return nil
+		}
+	}
+}
+
+func (*client) reverseStrings(record arrow.Record) (arrow.Record, error) {
+	for i, column := range record.Columns() {
+		if column.DataType().ID() != arrow.STRING {
+			continue
+		}
+		newColumnData := []string{}
+		for i := 0; i < column.Len(); i++ {
+			if !column.IsValid(i) {
+				continue
+			}
+			s := column.ValueStr(i)
+			runes := []rune(s)
+			for i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 {
+				runes[i], runes[j] = runes[j], runes[i]
+			}
+			newColumnData = append(newColumnData, string(runes))
+		}
+		fmt.Println("new column data is ", newColumnData)
+		mem := memory.NewGoAllocator()
+		bld := array.NewStringBuilder(mem)
+
+		// create an array with 4 values, no null
+		bld.AppendValues(newColumnData, nil)
+		var err error
+		record, err = record.SetColumn(i, bld.NewStringArray())
+		if err != nil {
+			return nil, err
+		}
+	}
+	return record, nil
+}

internal/reversertransformer/reversertransformer_test.go

@@ -0,0 +1,106 @@
+package reversertransformer
+
+import (
+	"context"
+	"io"
+	"testing"
+
+	"github.com/apache/arrow/go/v16/arrow"
+	"github.com/apache/arrow/go/v16/arrow/array"
+	"github.com/apache/arrow/go/v16/arrow/memory"
+	pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3"
+	internalPlugin "github.com/cloudquery/plugin-sdk/v4/internal/servers/plugin/v3"
+	"github.com/cloudquery/plugin-sdk/v4/plugin"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+var mem = memory.NewGoAllocator()
+
+func TestReverserTransformer(t *testing.T) {
+	p := plugin.NewPlugin("test", "development", GetNewClient())
+	s := internalPlugin.Server{
+		Plugin: p,
+	}
+	_, err := s.Init(context.Background(), &pb.Init_Request{
+		Spec:         []byte("{}"),
+		NoConnection: true,
+		InvocationId: "26b550f9-c6f8-4b4b-9ec4-773bab288ee6",
+	})
+	require.NoError(t, err)
+	requests := makeRequestsFromStrings("hello", "world")
+	stream := mockTransformServer{incomingMessages: requests}
+	require.NoError(t, s.Transform(&stream))
+	require.Equal(t, 2, len(stream.outgoingMessages))
+
+	record1, err := pb.NewRecordFromBytes(stream.outgoingMessages[0].Record)
+	require.NoError(t, err)
+	record2, err := pb.NewRecordFromBytes(stream.outgoingMessages[1].Record)
+	require.NoError(t, err)
+
+	require.Equal(t, "olleh", record1.Column(0).ValueStr(0))
+	require.Equal(t, "dlrow", record2.Column(0).ValueStr(0))
+}
+
+func makeRequestsFromStrings(s ...string) []*pb.Transform_Request {
+	requests := make([]*pb.Transform_Request, len(s))
+	for i, str := range s {
+		requests[i] = makeRequestFromString(str)
+	}
+	return requests
+}
+
+func makeRequestFromString(s string) *pb.Transform_Request {
+	record := makeRecordFromString(s)
+	bs, _ := pb.RecordToBytes(record)
+	return &pb.Transform_Request{Record: bs}
+}
+
+func makeRecordFromString(s string) arrow.Record {
+	str := array.NewStringBuilder(mem)
+	str.AppendString(s)
+	arr := str.NewStringArray()
+	schema := arrow.NewSchema([]arrow.Field{{Name: "col1", Type: arrow.BinaryTypes.String}}, nil)
+
+	return array.NewRecord(schema, []arrow.Array{arr}, 1)
+}
+
+type mockTransformServer struct {
+	grpc.ServerStream
+	incomingMessages []*pb.Transform_Request
+	outgoingMessages []*pb.Transform_Response
+}
+
+func (*mockTransformServer) SendAndClose(*pb.Transform_Response) error {
+	return nil
+}
+func (s *mockTransformServer) Recv() (*pb.Transform_Request, error) {
+	if len(s.incomingMessages) > 0 {
+		msg := s.incomingMessages[0]
+		s.incomingMessages = s.incomingMessages[1:]
+		return msg, nil
+	}
+	return nil, io.EOF
+}
+func (s *mockTransformServer) Send(resp *pb.Transform_Response) error {
+	s.outgoingMessages = append(s.outgoingMessages, resp)
+	return nil
+}
+func (*mockTransformServer) SetHeader(metadata.MD) error {
+	return nil
+}
+func (*mockTransformServer) SendHeader(metadata.MD) error {
+	return nil
+}
+func (*mockTransformServer) SetTrailer(metadata.MD) {
+}
+func (mockTransformServer) Context() context.Context {
+	return context.Background()
+}
+func (mockTransformServer) SendMsg(any) error {
+	return nil
+}
+func (mockTransformServer) RecvMsg(any) error {
+	return nil
+}

internal/servers/plugin/v3/plugin.go

@@ -394,6 +394,87 @@ func (s *Server) Write(stream pb.Plugin_WriteServer) error {
 	}
 }
 
+func (s *Server) Transform(stream pb.Plugin_TransformServer) error {
+	var (
+		recvRecords       = make(chan arrow.Record)
+		sendRecords       = make(chan arrow.Record)
+		pluginStopsWriter = make(chan struct{})
+		doneReading       = false
+		ctx               = stream.Context()
+		eg, gctx          = errgroup.WithContext(ctx)
+	)
+
+	// Run the plugin's transform with both channels.
+	//
+	// When the plugin is done, it must return with either an error or nil.
+	// The plugin must not close either channel.
+	eg.Go(func() error {
+		err := s.Plugin.Transform(gctx, recvRecords, sendRecords)
+		close(pluginStopsWriter)
+		doneReading = true
+		return err
+	})
+
+	// Write transformed records from transformer to destination.
+	//
+	// Currently the `sendRecords` channel is never closed. Instead, the plugin finishes this goroutine
+	// when it returns, either with an error or null.
+	//
+	// The reading never closes the writer, because it's up to the Plugin to decide when to finish
+	// writing, regardless of if the reading finished.
+	eg.Go(func() error {
+		for {
+			select {
+			case record := <-sendRecords:
+				recordBytes, err := pb.RecordToBytes(record)
+				if err != nil {
+					return status.Errorf(codes.Internal, "failed to convert record to bytes: %v", err)
+				}
+
+				if err := stream.Send(&pb.Transform_Response{Record: recordBytes}); err != nil {
+					return fmt.Errorf("error sending response: %w", err)
+				}
+			case <-pluginStopsWriter:
+				return nil
+			}
+		}
+	})
+
+	// Read records from source to transformer
+	//
+	// If there's an error receiving or deserialising records, or if there are no more records,
+	// the `recvRecords` channel will be closed. This will tell the plugin's transformer that
+	// no more transforming can be done.
+	//
+	// The writer cannot stop the reader even on error, but the plugin will when it returns,
+	// by setting `doneReading` to true.
+	eg.Go(func() error {
+		for {
+			req, err := stream.Recv()
+			if err == io.EOF {
+				close(recvRecords)
+				return nil
+			}
+			if err != nil {
+				close(recvRecords)
+				return fmt.Errorf("Error receiving request: %v", err)
+			}
+			if doneReading {
+				return nil
+			}
+			record, err := pb.NewRecordFromBytes(req.Record)
+			if err != nil {
+				close(recvRecords)
+				return status.Errorf(codes.InvalidArgument, "failed to create record: %v", err)
+			}
+
+			recvRecords <- record
+		}
+	})
+
+	return eg.Wait()
+}
+
 func (s *Server) Close(ctx context.Context, _ *pb.Close_Request) (*pb.Close_Response, error) {
 	return &pb.Close_Response{}, s.Plugin.Close(ctx)
 }

plugin/plugin.go

@@ -26,9 +26,12 @@ type NewClientFunc func(context.Context, zerolog.Logger, []byte, NewClientOption
 type Client interface {
 	SourceClient
 	DestinationClient
+	TransformerClient
 }
 
-type UnimplementedDestination struct{}
+type UnimplementedDestination struct {
+	UnimplementedTransformer
+}
 
 func (UnimplementedDestination) Write(context.Context, <-chan message.WriteMessage) error {
 	return ErrNotImplemented
@@ -38,7 +41,9 @@ func (UnimplementedDestination) Read(context.Context, *schema.Table, chan<- arro
 	return ErrNotImplemented
 }
 
-type UnimplementedSource struct{}
+type UnimplementedSource struct {
+	UnimplementedTransformer
+}
 
 func (UnimplementedSource) Sync(context.Context, SyncOptions, chan<- message.SyncMessage) error {
 	return ErrNotImplemented
@@ -48,6 +53,12 @@ func (UnimplementedSource) Tables(context.Context, TableOptions) (schema.Tables,
 	return nil, ErrNotImplemented
 }
 
+type UnimplementedTransformer struct{}
+
+func (UnimplementedTransformer) Transform(context.Context, <-chan arrow.Record, chan<- arrow.Record) error {
+	return ErrNotImplemented
+}
+
 // Plugin is the base structure required to pass to sdk.serve
 // We take a declarative approach to API here similar to Cobra
 type Plugin struct {

plugin/plugin_test.go

@@ -56,6 +56,9 @@ func (c *testPluginClient) Write(_ context.Context, res <-chan message.WriteMess
 func (*testPluginClient) Close(context.Context) error {
 	return nil
 }
+func (*testPluginClient) Transform(context.Context, <-chan arrow.Record, chan<- arrow.Record) error {
+	return nil
+}
 
 func TestPluginSuccess(t *testing.T) {
 	ctx := context.Background()

plugin/plugin_transformer.go

@@ -0,0 +1,15 @@
+package plugin
+
+import (
+	"context"
+
+	"github.com/apache/arrow/go/v16/arrow"
+)
+
+type TransformerClient interface {
+	Transform(context.Context, <-chan arrow.Record, chan<- arrow.Record) error
+}
+
+func (p *Plugin) Transform(ctx context.Context, recvRecords <-chan arrow.Record, sendRecords chan<- arrow.Record) error {
+	return p.client.Transform(ctx, recvRecords, sendRecords)
+}