WIP - writing

Author: radeksimko

tfprotov6/internal/toproto/state_store.go

@@ -28,12 +28,12 @@ func ConfigureStateStore_Response(in *tfprotov6.ConfigureStateStoreResponse) *tf
 	}
 }
 
-func ReadStateBytes_ResponseChunk(in *tfprotov6.StateByteChunk) *tfplugin6.ReadStateBytes_ResponseChunk {
+func ReadStateBytes_Response(in *tfprotov6.ReadStateByteChunk) *tfplugin6.ReadStateBytes_Response {
 	if in == nil {
 		return nil
 	}
 
-	return &tfplugin6.ReadStateBytes_ResponseChunk{
+	return &tfplugin6.ReadStateBytes_Response{
 		Diagnostics: Diagnostics(in.Diagnostics),
 		Bytes:       in.Bytes,
 		TotalLength: in.TotalLength,

tfprotov6/state_store.go

@@ -20,6 +20,8 @@ type StateStoreServer interface {
 	// ReadStateBytes streams byte chunks of a given state file from a state store
 	ReadStateBytes(context.Context, *ReadStateBytesRequest) (*ReadStateBytesStream, error)
 
+	WriteStateBytes(context.Context, *WriteStateBytesStream) (*WriteStateBytesResponse, error)
+
 	// GetStates returns a list of all states (i.e. CE workspaces) managed by a given state store
 	GetStates(context.Context, *GetStatesRequest) (*GetStatesResponse, error)
 
@@ -51,12 +53,28 @@ type ReadStateBytesRequest struct {
 }
 
 type ReadStateBytesStream struct {
-	Chunks iter.Seq[StateByteChunk]
+	Chunks iter.Seq[ReadStateByteChunk]
+}
+
+// type ChunkIterator func(StateByteChunkRequest) StateByteChunk
+
+type WriteStateBytesStream struct {
+	Chunks iter.Seq[WriteStateByteChunk]
+}
+
+type WriteStateBytesResponse struct {
+	Diagnostics []*Diagnostic
+}
+
+type WriteStateByteChunk = StateByteChunk
+
+type ReadStateByteChunk struct {
+	StateByteChunk
+	Diagnostics []*Diagnostic
 }
 
 type StateByteChunk struct {
 	Bytes       []byte
-	Diagnostics []*Diagnostic
 	TotalLength int64
 	Range       StateByteRange
 }

tfprotov6/tf6server/server.go

@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"os/signal"
 	"regexp"
@@ -1598,7 +1599,7 @@ func (s *server) ConfigureStateStore(ctx context.Context, protoReq *tfplugin6.Co
 	return protoResp, nil
 }
 
-func (s *server) ReadStateBytes(protoReq *tfplugin6.ReadStateBytes_Request, protoStream grpc.ServerStreamingServer[tfplugin6.ReadStateBytes_ResponseChunk]) error {
+func (s *server) ReadStateBytes(protoReq *tfplugin6.ReadStateBytes_Request, protoStream grpc.ServerStreamingServer[tfplugin6.ReadStateBytes_Response]) error {
 	rpc := "ReadStateBytes"
 	ctx := protoStream.Context()
 	ctx = s.loggingContext(ctx)
@@ -1628,12 +1629,13 @@ func (s *server) ReadStateBytes(protoReq *tfplugin6.ReadStateBytes_Request, prot
 
 	for chunk := range stream.Chunks {
 		select {
+		// TODO: check how interruptions are handled
 		case <-ctx.Done():
 			logging.ProtocolTrace(ctx, "all chunks sent")
 			return nil
 
 		default:
-			protoChunk := toproto.ReadStateBytes_ResponseChunk(&chunk)
+			protoChunk := toproto.ReadStateBytes_Response(&chunk)
 			if err := protoStream.Send(protoChunk); err != nil {
 				logging.ProtocolError(ctx, "Error sending chunk", map[string]any{logging.KeyError: err})
 				return err
@@ -1644,6 +1646,69 @@ func (s *server) ReadStateBytes(protoReq *tfplugin6.ReadStateBytes_Request, prot
 	return nil
 }
 
+func (s *server) WriteStateBytes(srv grpc.ClientStreamingServer[tfplugin6.WriteStateBytes_RequestChunk, tfplugin6.WriteStateBytes_Response]) error {
+	rpc := "WriteStateBytes"
+	ctx := srv.Context()
+	ctx = s.loggingContext(ctx)
+	ctx = logging.RpcContext(ctx, rpc)
+	// ctx = logging.StateStoreContext(ctx, protoReq.TypeName)
+	ctx = s.stoppableContext(ctx)
+	// logging.ProtocolTrace(ctx, "Received request")
+	// defer logging.ProtocolTrace(ctx, "Served request")
+
+	ctx = tf6serverlogging.DownstreamRequest(ctx)
+
+	server, ok := s.downstream.(tfprotov6.StateStoreServer)
+	if !ok {
+		err := status.Error(codes.Unimplemented, "ProviderServer does not implement WriteStateBytes")
+		logging.ProtocolError(ctx, err.Error())
+		return err
+	}
+
+	var iteratorErr error
+
+	// TODO: what about error handling per chunk and providers having the ability to do cleanup on interruption?
+
+	iterator := func(yield func(tfprotov6.WriteStateByteChunk) bool) {
+		for {
+			chunk, err := srv.Recv()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				iteratorErr = err
+				srv.SendMsg(&tfplugin6.WriteStateBytes_Response{
+					// Diagnostics: ,
+				})
+				return
+			}
+
+			yield(tfprotov6.WriteStateByteChunk{
+				Bytes:       chunk.Bytes,
+				TotalLength: chunk.TotalLength,
+				Range: tfprotov6.StateByteRange{
+					Start: chunk.Range.Start,
+					End:   chunk.Range.End,
+				},
+			})
+
+		}
+	}
+
+	resp, err := server.WriteStateBytes(ctx, &tfprotov6.WriteStateBytesStream{
+		Chunks: iterator,
+	})
+	if err != nil {
+		return err
+	}
+
+	err = srv.SendAndClose(&tfplugin6.WriteStateBytes_Response{
+		// Diagnostics: resp.Diagnostics,
+	})
+
+	return nil
+}
+
 func (s *server) GetStates(ctx context.Context, protoReq *tfplugin6.GetStates_Request) (*tfplugin6.GetStates_Response, error) {
 	rpc := "GetStates"
 	ctx = s.loggingContext(ctx)